Управление диффузионным трансформером с помощью его внутренней динамики

Guiding a Diffusion Transformer with the Internal Dynamics of Itself

December 30, 2025

Авторы: Xingyu Zhou, Qifan Li, Xiaobin Hu, Hai Chen, Shuhang Gu

cs.AI

Аннотация

Модели диффузии демонстрируют высокую способность к захвату полного (условного) распределения данных. Однако из-за недостатка обучения и данных для освоения маловероятных областей модель штрафуется за невозможность генерации высококачественных изображений, соответствующих этим областям. Для повышения качества генерации такие стратегии управления, как классификаторное управление без классификатора (CFG), могут направлять сэмплы в области высокой вероятности на этапе сэмплирования. Тем не менее, стандартный CFG часто приводит к излишне упрощённым или искажённым результатам. С другой стороны, альтернативный подход управления диффузионной моделью через её "плохую" версию ограничен необходимостью тщательного проектирования стратегий деградации, дополнительного обучения и увеличения шагов сэмплирования. В данной работе мы предлагаем простую, но эффективную стратегию — внутреннее управление (IG), которая вводит вспомогательный контроль на промежуточном слое во время обучения и экстраполирует выходы промежуточных и глубоких слоёв для получения генеративных результатов на этапе сэмплирования. Эта простая стратегия обеспечивает значительное улучшение как эффективности обучения, так и качества генерации на различных базовых моделях. На ImageNet 256×256 модель SiT-XL/2+IG достигает FID=5,31 и FID=1,75 на 80 и 800 эпохах соответственно. Более впечатляюще, LightningDiT-XL/1+IG достигает FID=1,34, что значительно превосходит все остальные методы. В сочетании с CFG модель LightningDiT-XL/1+IG устанавливает новый рекорд, достигая state-of-the-art значение FID=1,19.

English

The diffusion model presents a powerful ability to capture the entire (conditional) data distribution. However, due to the lack of sufficient training and data to learn to cover low-probability areas, the model will be penalized for failing to generate high-quality images corresponding to these areas. To achieve better generation quality, guidance strategies such as classifier free guidance (CFG) can guide the samples to the high-probability areas during the sampling stage. However, the standard CFG often leads to over-simplified or distorted samples. On the other hand, the alternative line of guiding diffusion model with its bad version is limited by carefully designed degradation strategies, extra training and additional sampling steps. In this paper, we proposed a simple yet effective strategy Internal Guidance (IG), which introduces an auxiliary supervision on the intermediate layer during training process and extrapolates the intermediate and deep layer's outputs to obtain generative results during sampling process. This simple strategy yields significant improvements in both training efficiency and generation quality on various baselines. On ImageNet 256x256, SiT-XL/2+IG achieves FID=5.31 and FID=1.75 at 80 and 800 epochs. More impressively, LightningDiT-XL/1+IG achieves FID=1.34 which achieves a large margin between all of these methods. Combined with CFG, LightningDiT-XL/1+IG achieves the current state-of-the-art FID of 1.19.