Apple Machine Learning Research(RSS)
56AI 编辑部评分,满分 100

驯服扩散 Transformer 中的离群 token:Dual-Stage Registers 干预

2026-08-05 08:00· 22小时前
AI 导读

研究发现扩散 Transformer(DiT)图像生成流程中,预训练 ViT 编码器和 DiT 去噪器均会产生离群 token,尤其在中间层,且简单掩蔽高范数 token 无法改善性能,问题与局部 patch 语义损坏相关。为此提出 Dual-Stage Registers(DSR)干预方法,在 ImageNet 和文生图任务上持续减少离群伪影并提升生成质量。

We study outlier tokens in Diffusion Transformers (DiTs) for image generation. Prior work has shown that Vision Transformers (ViTs) can produce a small number of high-norm tokens that attract disproportionate attention while carrying limited local information, but their role in generative models remains underexplored. We show that this phenomenon appears in both the encoder and denoiser of modern Representation Autoencoder (RAE)-DiT pipelines: pretrained ViT encoders can produce outlier representations, and DiTs themselves can develop internal outlier tokens, especially in intermediate layers. Moreover, simply masking high-norm tokens does not improve performance, indicating that the problem is not only caused by a few extreme values, but is more closely related to corrupted local patch semantics. To address this issue, we introduce Dual-Stage Registers (DSR), a register-based intervention for both components: trained registers when available, recursive test-time registers otherwise, and diffusion registers for the denoiser. Across ImageNet and large-scale text-to-image generation, these interventions consistently reduce outlier artifacts and improve generation quality. Our results highlight outlier-token control as an important ingredient in building stronger DiTs.

  • † Rice University
  • * Equal contribution

Related readings and updates.

DiT-Air: Revisiting the Efficiency of Diffusion Model Architecture Design in Text to Image Generation

December 11, 2025research area Computer Vision

In this work, we empirically study Diffusion Transformers (DiTs) for text-to-image generation, focusing on architectural choices, text-conditioning strategies, and training protocols. We evaluate a range of DiT-based architectures—including PixArt-style and MMDiT variants—and compare them with a standard DiT variant which directly processes concatenated text and noise inputs. Surprisingly, our findings reveal that the performance of standard…

On Inductive Biases That Enable Generalization of Diffusion Transformers

September 22, 2025research area Computer Visionconference NeurIPS

Recent work studying the generalization of diffusion models with UNet-based denoisers reveals inductive biases that can be expressed via geometry-adaptive harmonic bases. However, in practice, more recent denoising networks are often based on transformers, e.g., the diffusion transformer (DiT). This raises the question: do transformer-based denoising networks exhibit inductive biases that can also be expressed via geometry-adaptive harmonic…

Bottom banner

Discover opportunities in Machine Learning.

来源:Apple Machine Learning Research(RSS) · machinelearning.apple.com

驯服扩散 Transformer 中的离群 token:Dual-Stage Registers 干预

Apple Machine Learning Research(RSS)·2026-08-05 08:00·22小时前
AI 导读

研究发现扩散 Transformer(DiT)图像生成流程中,预训练 ViT 编码器和 DiT 去噪器均会产生离群 token,尤其在中间层,且简单掩蔽高范数 token 无法改善性能,问题与局部 patch 语义损坏相关。为此提出 Dual-Stage Registers(DSR)干预方法,在 ImageNet 和文生图任务上持续减少离群伪影并提升生成质量。

原文 · 保持原样,未翻译

We study outlier tokens in Diffusion Transformers (DiTs) for image generation. Prior work has shown that Vision Transformers (ViTs) can produce a small number of high-norm tokens that attract disproportionate attention while carrying limited local information, but their role in generative models remains underexplored. We show that this phenomenon appears in both the encoder and denoiser of modern Representation Autoencoder (RAE)-DiT pipelines: pretrained ViT encoders can produce outlier representations, and DiTs themselves can develop internal outlier tokens, especially in intermediate layers. Moreover, simply masking high-norm tokens does not improve performance, indicating that the problem is not only caused by a few extreme values, but is more closely related to corrupted local patch semantics. To address this issue, we introduce Dual-Stage Registers (DSR), a register-based intervention for both components: trained registers when available, recursive test-time registers otherwise, and diffusion registers for the denoiser. Across ImageNet and large-scale text-to-image generation, these interventions consistently reduce outlier artifacts and improve generation quality. Our results highlight outlier-token control as an important ingredient in building stronger DiTs.

  • † Rice University
  • * Equal contribution

Related readings and updates.

DiT-Air: Revisiting the Efficiency of Diffusion Model Architecture Design in Text to Image Generation

December 11, 2025research area Computer Vision

In this work, we empirically study Diffusion Transformers (DiTs) for text-to-image generation, focusing on architectural choices, text-conditioning strategies, and training protocols. We evaluate a range of DiT-based architectures—including PixArt-style and MMDiT variants—and compare them with a standard DiT variant which directly processes concatenated text and noise inputs. Surprisingly, our findings reveal that the performance of standard…

On Inductive Biases That Enable Generalization of Diffusion Transformers

September 22, 2025research area Computer Visionconference NeurIPS

Recent work studying the generalization of diffusion models with UNet-based denoisers reveals inductive biases that can be expressed via geometry-adaptive harmonic bases. However, in practice, more recent denoising networks are often based on transformers, e.g., the diffusion transformer (DiT). This raises the question: do transformer-based denoising networks exhibit inductive biases that can also be expressed via geometry-adaptive harmonic…

Bottom banner

Discover opportunities in Machine Learning.

来源:Apple Machine Learning Research(RSS)· machinelearning.apple.com