When Worse is Better: Navigating the compression-generation tradeoff in visual tokenization

GenAI Diffusion TDIG
当前的图像生成方法,例如潜在扩散和基于离散标记的生成,依赖于两阶段的训练方法。在第一阶段,自动编码器被训练以将图像压缩到潜在空间;在第二阶段,生成模型被训练以学习该潜在空间上的分布。大多数工作集中在独立于第二阶段来最大化第一阶段的性能,假设更好的重建总是能带来更好的生成效果。然而,我们证明这并不完全正确。即使重建性能下降,较小的第二阶段模型可以从更压缩的第一阶段潜在表示中受益,显示出压缩和生成建模能力之间存在根本性的权衡。 为了更好地优化这一权衡,我们引入了因果正则化分词(Causally Regularized Tokenization, CRT),它利用第二阶段生成建模过程的知识,在第一阶段的潜在表示中嵌入有用的归纳偏见。这种正则化使第一阶段的重建性能变差,但通过使标记更容易建模,从而提高了第二阶段的生成性能:我们能够将计算效率提高2-3倍,并且使用不到一半的每张图像标记数(256对576)以及总模型参数量仅为四分之一(7.75亿对31亿)的情况下,匹配最先进的离散自回归ImageNet生成(FID 2.18),超过了之前的最先进水平(LlamaGen)。
Current image generation methods, such as latent diffusion and discrete token-based generation, depend on a two-stage training approach. In stage 1, an auto-encoder is trained to compress an image into a latent space; in stage 2, a generative model is trained to learn a distribution over that latent space. Most work focuses on maximizing stage 1 performance independent of stage 2, assuming better reconstruction always leads to better generation. However, we show this is not strictly true. Smaller stage 2 models can benefit from more compressed stage 1 latents even if reconstruction performance worsens, showing a fundamental trade-off between compression and generation modeling capacity. To better optimize this trade-off, we introduce Causally Regularized Tokenization (CRT), which uses knowledge of the stage 2 generation modeling procedure to embed useful inductive biases in stage 1 latents. This regularization makes stage 1 reconstruction performance worse, but makes stage 2 generation performance better by making the tokens easier to model: we are able to improve compute efficiency 2-3$\times$ over baseline and match state-of-the-art discrete autoregressive ImageNet generation (2.18 FID) with less than half the tokens per image (256 vs. 576) and a fourth the total model parameters (775M vs. 3.1B) as the previous SOTA (LlamaGen).
许愿