Реконструкция против Генерации: Укрощение Дилеммы Оптимизации в Моделях Латентной Диффузии

Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models

January 2, 2025

Авторы: Jingfeng Yao, Xinggang Wang

cs.AI

Аннотация

Модели латентной диффузии с архитектурами Трансформера отличаются в создании изображений высокой точности. Однако недавние исследования выявляют дилемму оптимизации в этом двухэтапном дизайне: увеличение размерности признаков на токен визуальных токенизаторов улучшает качество восстановления, но требует значительно более крупных моделей диффузии и большего числа итераций обучения для достижения сопоставимой производительности генерации. В результате существующие системы часто останавливаются на субоптимальных решениях, либо порождают визуальные артефакты из-за потери информации в токенизаторах, либо не могут полностью сойтись из-за высоких вычислительных затрат. Мы считаем, что эта дилемма проистекает из врожденной сложности обучения неограниченных высокоразмерных латентных пространств. Для решения этой проблемы мы предлагаем выравнивание латентного пространства с предварительно обученными моделями визуальных основ при обучении визуальных токенизаторов. Наш предложенный VA-VAE (Выравненный вариационный автокодировщик с моделью визуальных основ) значительно расширяет границу восстановления-генерации латентных моделей диффузии, обеспечивая более быструю сходимость Диффузионных Трансформеров (DiT) в высокоразмерных латентных пространствах. Для полного использования потенциала VA-VAE мы создаем улучшенную базовую модель DiT с улучшенными стратегиями обучения и архитектурными конструкциями, названную LightningDiT. Интегрированная система достигает передовой производительности на генерации ImageNet 256x256 со значением FID 1.35, продемонстрировав замечательную эффективность обучения, достигнув значения FID 2.11 всего за 64 эпохи - представляя более чем в 21 раз ускоренную сходимость по сравнению с оригинальным DiT. Модели и код доступны по ссылке: https://github.com/hustvl/LightningDiT.

English

Latent diffusion models with Transformer architectures excel at generating high-fidelity images. However, recent studies reveal an optimization dilemma in this two-stage design: while increasing the per-token feature dimension in visual tokenizers improves reconstruction quality, it requires substantially larger diffusion models and more training iterations to achieve comparable generation performance. Consequently, existing systems often settle for sub-optimal solutions, either producing visual artifacts due to information loss within tokenizers or failing to converge fully due to expensive computation costs. We argue that this dilemma stems from the inherent difficulty in learning unconstrained high-dimensional latent spaces. To address this, we propose aligning the latent space with pre-trained vision foundation models when training the visual tokenizers. Our proposed VA-VAE (Vision foundation model Aligned Variational AutoEncoder) significantly expands the reconstruction-generation frontier of latent diffusion models, enabling faster convergence of Diffusion Transformers (DiT) in high-dimensional latent spaces. To exploit the full potential of VA-VAE, we build an enhanced DiT baseline with improved training strategies and architecture designs, termed LightningDiT. The integrated system achieves state-of-the-art (SOTA) performance on ImageNet 256x256 generation with an FID score of 1.35 while demonstrating remarkable training efficiency by reaching an FID score of 2.11 in just 64 epochs--representing an over 21 times convergence speedup compared to the original DiT. Models and codes are available at: https://github.com/hustvl/LightningDiT.