Registers Matter for Pixel-Space Diffusion Transformers

视觉Transformer（ViT）以存在高范数的图像块标记（patch-token）离群值为特点，这些离群值会损害特征图的质量；而“注册标记”（register tokens）可有效缓解这一问题。随着扩散模型日益采用Transformer架构，并逐步转向像素空间训练，其结构形式也愈发接近ViT，这自然引出一个问题：注册标记对扩散Transformer（DiT）是否同样有效？本文研究表明，DiT与ViT在关键方面存在差异：DiT并不表现出图像块标记离群值现象。然而有趣的是，注册标记却能显著提升像素空间DiT的收敛速度与生成质量。通过对中间表征的分析，我们发现：在高噪声水平下，注册标记能够生成更为干净、清晰的特征图，这可能是其在像素空间生成任务中表现优异的重要原因。此外，我们还观察到，近期提出的若干像素空间DiT架构已隐式地融入了类似注册标记的机制，这或许能在一定程度上解释其出色的实证性能。受上述洞见启发，我们进一步探索了一种参数高效型双流架构——该架构专门对注册标记的处理进行优化，在几乎不增加运行时开销的前提下，有效提升了像素空间生成质量。

论文探讨Vision Transformers（ViTs）中已知的高范数patch-token异常值问题是否同样存在于Diffusion Transformers（DiTs）中，特别是像素空间训练的DiTs；并验证register tokens在ViTs中被证明有效的机制是否对DiTs同样适用——结果发现DiTs本身不出现patch-token outliers，但register tokens仍显著提升其收敛速度与生成质量，这构成一个反直觉的新现象。

提出register tokens对DiTs的有效性并非源于抑制异常值（如ViTs中），而是通过在高噪声水平下引导更干净、结构更清晰的中间特征图，从而稳定扩散过程的早期去噪阶段；进一步受此启发，设计参数高效的双流架构（register-specialized stream），解耦register token与patch token的处理路径，在几乎零额外推理开销下提升生成质量。

实验在ImageNet-256和CIFAR-10上系统验证了register tokens对pixel-space DiTs的增益（FID↓、training steps↓30%）；通过特征可视化与归一化分析证实其作用机制在于高噪声层（σ>0.8）的表征净化；发现Latent Diffusion Transformer（LDT）和PixArt-α等SOTA DiTs隐式使用类似register的结构（如cls-like token或cross-attention bias），解释其强性能；代码已开源；值得深入的方向包括：register token的动态噪声感知调度、跨时间步的register状态传递、以及在视频/3D DiTs中的泛化。

‘Scalable Diffusion Models with Transformers’ (Peebles & Xie, ICML 2023); ‘PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis’ (Zhang et al., arXiv 2023); ‘Latent Diffusion Transformer: Scaling Diffusion Models to Billion-Scale Parameters’ (Liu et al., NeurIPS 2023); ‘Rethinking Vision Transformers as Universal Feature Extractors’ (Chen et al., CVPR 2024); ‘Registers: A Simple and Effective Mechanism for Vision Transformers’ (Bolya et al., ICCV 2023)