Регистры важны для диффузионных трансформеров в пиксельном пространстве

Аннотация

Известно, что Vision Transformers (ViTs) демонстрируют высоконормовые выбросы патч-токенов, ухудшающие качество карт признаков, — проблема, эффективно решаемая с помощью токенов-регистров. Поскольку диффузионные модели всё чаще используют архитектуры трансформеров и переходят к обучению в пиксельном пространстве, они становятся ближе по форме к ViT, что ставит вопрос о том, полезны ли токены-регистры также для Diffusion Transformers (DiTs). В данной работе мы показываем, что DiTs принципиально отличаются от ViTs: они не проявляют выбросов патч-токенов, но всё же выигрывают от использования регистров. Интересно, что регистры более эффективны для DiTs в пиксельном пространстве, чем для DiTs в латентном пространстве. Анализируя промежуточные представления, мы обнаруживаем, что токены-регистры создают более чистые карты признаков при высоких уровнях шума, что может способствовать их эффективности при генерации в пиксельном пространстве. Мы также отмечаем, что современные архитектуры DiTs в пиксельном пространстве неявно включают механизмы, подобные регистрам, что может частично объяснять их высокую эмпирическую производительность. Основываясь на этих наблюдениях, мы предлагаем Register Guidance — метод, усиливающий вклад токенов-регистров, ответственных за улучшение визуальной структуры и согласованности.

English

Vision Transformers (ViTs) are known to exhibit high-norm patch-token outliers that degrade feature map quality, a problem effectively mitigated by register tokens. As diffusion models increasingly adopt transformer architectures and move toward pixel-space training, they become closer in form to ViTs, raising the question of whether register tokens are also useful for Diffusion Transformers (DiTs). In this work, we show that DiTs differ from ViTs in a key respect: they do not exhibit patch-token outliers but still benefit from registers. Interestingly, registers are more effective in pixel-space DiTs than in latent-space DiTs. By analyzing intermediate representations, we find that register tokens produce cleaner feature maps at high noise levels, which may contribute to their effectiveness in pixel-space generation. We further observe that recent pixel-space DiT architectures implicitly incorporate register-like mechanisms, which may partially account for their strong empirical performance. Motivated by these observations, we propose Register Guidance, a technique that amplifies the contribution of register tokens responsible for improving visual structure and coherence.