像素空间扩散Transformer中的寄存器令牌机制 Registers Matter for Pixel-Space Diffusion Transformers
寄存器令牌在无异常值的DiT中仍提升像素空间生成质量，并提出Register Guidance
前置知识
Vision Transformer (ViT) 与 patch token
Vision Transformer 将图像切分为固定大小的图块（patch），每个图块线性投影成一个 token，再像序列一样通过自注意力处理。在自监督训练的 ViT（如 DINOv2）中，研究发现少数 patch token 会演变成范数异常高的离群点，常聚集在背景等低信息量区域，进而污染注意力图、削弱密集预测性能。
本文的研究起点正是 ViT 的高范数离群点问题，理解 ViT 中 token 行为是理解作者为何把寄存器令牌迁移到扩散模型的前提。
寄存器令牌（register token）
寄存器令牌是附加在 patch-token 序列之后的若干额外可学习 token，它们不参与最终的损失监督，从而为网络提供“空闲槽位”。在 ViT 中，这些槽位会吸收原本要由 patch token 承担的全局信息和高范数异常，使 patch token 的特征图更干净、注意力图更可解释。本文沿用这一机制：在 DiT 的 patch 序列后附加若干寄存器令牌，训练目标只施加在 patch token 上。
寄存器令牌是贯穿全文的核心概念，论文的全部发现与方法（包括 Register Guidance）都围绕它在 DiT 中是否有效、为何有效展开。
扩散Transformer (DiT) 与流匹配 (flow matching)
DiT 用 Transformer 替代 U-Net 作为扩散模型主干。流匹配是一种连续流形式的生成目标：给定数据样本 $x$ 和噪声 $\epsilon$，构造线性插值 $z = t\,x + (1-t)\,\epsilon$，模型学习速度场 $v=(x-z)/(1-t)$，推理时沿该速度场把噪声积分到数据。时间步 $t$ 接近 0 对应高噪声、$t$ 接近 1 对应低噪声；高噪声步对高维像素空间模型尤为重要，因为它塑造图像的主要内容。
本文的 pDiT、PixelDiT、JiT、SiT、RAE 等模型都用流匹配训练，且关键分析（如特征平滑度随 $t$ 变化）必须结合扩散时间步来理解。
像素空间 vs 潜在空间扩散
潜在扩散（latent diffusion）先用预训练自编码器（VAE 或基于 DINOv2 的 RAE）把图像压到低维潜在空间再扩散；像素空间扩散则直接在原始像素上建模。像素空间目标维度更高、噪声更不可感知，中间特征更不平滑、训练更困难，但省去了自编码器。代表工作有 PixelDiT、JiT、VDM++ 等。
论文最核心的发现之一是寄存器令牌在像素空间最有效、潜在空间次之、RAE 空间反而有害，理解两种空间的差异是把握这一结论的关键。
Classifier-Free Guidance (CFG) 与 AutoGuidance
CFG 在采样时用条件预测与无条件预测之差来放大条件信号：$v = v_u + w_{cfg}(v_{cond}-v_u)$。AutoGuidance 进一步用一个人工缩放的“弱化版”自身模型替代无条件模型，认为弱模型与强模型的差异是更高质量的引导方向。本文的 Register Guidance 把“去掉寄存器令牌的同一模型”视为天然的弱模型，是 AutoGuidance 思想的具体实例。
Register Guidance 的设计直接借鉴 CFG 与 AutoGuidance，并与 CFG 互补叠加，不理解这两个概念就无法读懂论文的方法与实验。
研究动机
在视觉表征学习中，Vision Transformer (ViT) 普遍存在高范数 patch-token 离群点问题：少数 token 的特征范数远高于其余，常聚集在背景等低信息区域，导致注意力图出现伪影、密集预测性能下降。Darcot 等人提出用寄存器令牌（register token）来吸收这些离群点以缓解问题。与此同时，扩散模型正大规模采用 Transformer 架构（DiT）并日益向像素空间训练发展（如 PixelDiT、JiT、VDM++），使其在形式上越来越接近 ViT。这就自然引出两个悬而未决的问题：（1）DiT 是否也继承了 ViT 那样的高范数 patch-token 离群点？（2）即便没有离群点，寄存器令牌在 DiT 中是否仍能带来收益？此外，attention sink 与特殊 token 行为在生成式 Transformer 中已被广泛研究，但在图像扩散 Transformer 中，寄存器类 token 的存在与作用几乎未被系统探讨，其作用机制完全不清楚。
本文的目标是本文的目标是系统厘清寄存器令牌在图像扩散 Transformer（尤其像素空间 DiT）中的作用，并据此设计一个实用的采样技术以进一步提升生成质量。具体目标包括：诊断 DiT 是否存在 patch-token 异常值；无论存在与否，定量评估寄存器令牌对生成质量的影响；从特征范数、空间平滑度（Total Variation）、线性探测准确率等角度揭示寄存器令牌改善生成的具体机制；比较寄存器令牌在像素空间、VAE 空间、RAE 空间三种训练空间的效果差异；剖析现有高性能像素空间 DiT（如 JiT 的 in-context conditioning）是否隐含了寄存器机制；最终提出 Register Guidance (RG) 方法，放大寄存器令牌对结构与连贯性的有益贡献，并与 CFG 组合取得更好结果。
与已有工作不同的是，本文的独特切入角度是“反直觉的实证发现 + 深度机制分析 + 实用方法”三位一体。与社区预期相反，作者发现 DiT 根本不存在 patch-token 离群点，按理不需要寄存器令牌，但实验却证明寄存器令牌反而更有效，尤其对像素空间模型。这一反差驱动了深入机制分析，并揭示寄存器令牌在 DiT 中扮演双重角色（范数吸收槽 + 全局语义载体），与 ViT 中仅作为异常缓冲截然不同。更具洞察力的是，作者证明近期高性能像素空间 DiT（JiT 的 in-context conditioning）的大部分收益来自寄存器类行为而非额外类别信息，从而重新解释了已有工作的成功机理。最后 RG 把“同一模型去寄存器”当作天然弱模型，填补了 AutoGuidance 中“如何获得弱模型”这一难点。
核心方法
方法整体分两条线：先做大量诊断性分析以理解寄存器令牌的作用机制，再据此设计实用的 Register Guidance。分析层面，作者训练像素空间 pDiT（标准 DiT 架构叠加常用 Transformer 改进），用 flow matching 在 ImageNet 上以 $256\times256$ 分辨率、patch size 16 训练，模型规模从 B 到 H，并系统性对比有无寄存器令牌的版本；同时也分析 PixelDiT、SiT、RAE 等潜在空间架构作为对照。寄存器令牌实现为附加在 patch-token 序列后的可学习 token，不参与训练损失。作者通过测量 token-wise 特征范数分布、用 Total Variation (TV) 量化中间特征的空间平滑度、对寄存器令牌做线性探测分类等多种手段定位其作用。实用方法层面，受 AutoGuidance 启发，把无寄存器令牌的模型视为有寄存器令牌模型的“弱化版”，用二者速度预测之差作为引导方向，得到 Register Guidance，并可自然地与 CFG 叠加。
核心创新点是把“去掉寄存器令牌的同一模型”当作“保留寄存器令牌模型”的天然弱化版本，并用二者预测之差作为引导方向。这与 CFG（用无条件模型）和 AutoGuidance（用人工缩放的弱模型）有本质区别：RG 的弱模型是自然得到的，无需人为设计缩放，且二者预测保持相似图像内容但结构连贯性不同，因此其差值构成有意义的“质量方向”。为消除两个独立模型预测不对齐所带来的伪影，作者训练单一统一模型：训练时以概率 $p$ 随机丢弃所有寄存器令牌，使同一模型既能以有寄存器令牌方式、又能以无寄存器令牌方式推理，显著提升两路预测的对齐度。关键洞察是寄存器令牌改善的是物体结构和视觉连贯性而保留主体内容，这正是可被引导放大的、与 CFG 正交的质量维度。
方法步骤详情
方法分分析与 RG 两部分。分析流程：训练 pDiT/PixelDiT/SiT/RAE 各自有无寄存器令牌的版本；在每个 block 后提取特征统计 token 范数分布；在不同时间步测 patch-token 范数；用 TV 计算（有/无寄存器令牌）比值评估空间平滑度；对寄存器令牌做线性探测判断其语义角色并可视化 attention map；在 JiT 上对 in-context token 做同样分析。RG 推理（Snippet 1）：对样本 $z$ 分别前向得 $x_r$（regs=True）与 $x_{nr}$（regs=False），速度 $v=(x-z)/(1-t)$，合成 $v_{rg}=v_{nr}+w_{rg}(v_r-v_{nr})$，沿 $z_{next}=z+(t_{next}-t)v_{rg}$ 更新。RG+CFG（Snippet 2）引入无条件 $v_u$。RG 训练（Snippet 3）：$z=tx+(1-t)\epsilon$、目标 $v=(x-z)/(1-t)$，以 $\mathrm{rand()}\geq p$ 决定是否丢弃寄存器令牌，两 regime 用独立 batch 训同一模型。
技术新颖性
技术新颖性体现在四个层面。第一是反直觉的实证发现：DiT 本不存在 patch-token 离群点，按理不需要寄存器令牌，但寄存器令牌仍会自发形成自身的高范数 token 并显著提升质量，颠覆了“寄存器令牌只为吸收既有异常值”的传统认知。第二是机制层面的新解释：作者证明寄存器令牌在 DiT 中通过降低所有 patch-token 范数、提升高噪声步（$t\in[0,0.2]$）的中间特征平滑度（TV 比值小于 1）来起作用，并发现寄存器令牌只在深层（block 4–11）有效、所需数量远多于 ViT（32 个 vs ViT 的 4 个）。第三是重新解释了已有工作：JiT 的 in-context conditioning 收益主要来自寄存器类行为而非额外类别信息，大尺度文本到图像模型中部分 text token 也表现为 sink。第四是 RG 方法本身作为 AutoGuidance 的具体实例，给出了获得弱模型的自然途径，并能与 CFG 互补叠加，在 ImageNet 256/512 多个规模上一致改进。
实验结果
核心发现带具体数字。其一，DiT 无 patch-token 异常值：Figure 2 显示 DINOv2 有显著范数离群 token，而 pDiT 范数分布近乎均匀、注意力图无背景伪影。其二，寄存器令牌仍大幅提升像素空间质量：Table 1 中 pDiT-B/16 epoch 200 的 FID 从 7.39 降到 5.30、IC 达 4.71；pDiT-H/16 epoch 600 从 2.35 到 2.02、IC 达 1.90。其三，效果随空间变化：Table 2 显示像素空间提升最大、VAE 中等、RAE 反受损（large 4.44 有 vs 3.91 无）。其四，机制：寄存器令牌降低 patch-token 范数（Figure 3），在高噪声步使特征更平滑（Figure 4 中 TV 比值低至 0.76–0.85）。其五，双重角色：线性探测显示部分寄存器令牌高准确率（约 0.9）编码全局语义、部分低准确率（约 0.02）高范数作 sink，且只在深层有效、需 32 个 token（Table 3）。其六，RG 在 $p=0.03$ 最优，RG+CFG 达 3.32 优于 CFG 3.71；Table 6 中 JiT-B/16 CFG+RG 达 2.96（基线 3.66）、JiT-H/16 达 1.80。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| pDiT-B/16 生成质量（epoch 200，ImageNet 256） | FID↓ | 5.30（有寄存器令牌）/ 4.71（in-context） | 7.39（无寄存器令牌） | 约 28%（有寄存器令牌）/ 36%（in-context） |
| pDiT-H/16 生成质量（epoch 600） | FID↓ | 2.02（有寄存器令牌）/ 1.90（in-context） | 3.52（无寄存器令牌） | 约 43% / 46% |
| JiT-B/16 ImageNet 256×256（CFG+RG） | FID↓ | 2.96 | 3.66（JiT 基线） | 约 19% |
| JiT-H/16 ImageNet 256×256（CFG+RG） | FID↓ | 1.80 | 1.86（JiT 基线） | 约 3% |
| RG 与 CFG 互补（JiT-B/16，p=0.03） | FID↓ | 3.32（RG+CFG） | 3.71（仅 CFG） | 约 11% |
| JiT-H/32 ImageNet 512×512（CFG+RG） | FID↓ | 1.87 | 1.94（JiT 基线） | 约 4% |
局限与改进
作者承认的局限包括：Register Guidance 主要在类条件像素空间设置（JiT/ImageNet）上验证，向文本到图像扩展是明确的未来方向但尚未实现；作者初步探索的若干高效双流 DiT 架构并未观察到明显改进。基于论文内容，我观察到几点额外局限：其一，机制分析主要依赖线性探测、TV 等代理指标，对为何在 RAE 空间寄存器令牌反而有害的解释（潜在空间更结构化、更低维）仍偏定性，缺乏更严格的因果论证；其二，RG 需要在训练时以概率 $p$ 丢弃寄存器令牌、引入 $p$ 与 $w_{rg}$ 两个新超参，增加了训练复杂度和调参成本；其三，所有实验局限于 ImageNet 类条件，未在更大规模真实世界数据集或更高分辨率多样性上验证寄存器机制与 RG 的普适性；其四，RG 相对 CFG 的增益随模型变大而缩小（B/16 约 0.70，H/16 仅 0.06），暗示大模型边际收益递减。
独立分析的弱点
第一个弱点是 RG 收益随模型规模递减。Table 6 中 JiT-B/16 从 3.66 降到 2.96（降 0.70），但 JiT-H/16 仅从 1.86 降到 1.80（降 0.06），说明大模型的寄存器机制可能已被其它设计部分吸收。改进方向是研究大模型中寄存器令牌与已有 token（如 text token、in-context token）的交互，设计更精细的引导或显式分离的 token 角色。第二个弱点是 RAE 空间寄存器令牌有害的机制未完全厘清，仅给出“潜在空间更结构化”的定性解释。改进方向是用受控实验分离潜在维度、噪声结构、范数分布等因子，量化各自贡献。第三个弱点是实验仅限 ImageNet 类条件。改进方向是扩展到 text-to-image——作者已指出部分 text token 会变为 sink（Figure 25），RG 有潜力但要解决 text CFG 与 RG 的协调。第四个弱点是训练需双 regime（$p$-drop）与独立 batch，增加成本。改进方向是用注意力掩码替代独立 batch，或研究更高效的对齐训练策略。
未来方向
作者明确提出的未来方向包括：开发双流 DiT 架构以更有效地分别处理寄存器令牌与 patch token（初步尝试未成功但作者认为值得继续探索）；将 Register Guidance 扩展到文本到图像模型；更深入理解寄存器令牌在 DiT 中的全部潜力。基于本文成果可延伸的方向有：其一，把寄存器令牌机制与 REPA 等表示对齐方法结合，作者在附录 D.3 已初步探索二者是否互补，值得系统化；其二，研究寄存器令牌的最优数量、起始层、初始化与模型规模、分辨率的 scaling law；其三，将寄存器令牌“sink vs 语义载体”的特化思想用于设计可解释的条件化 token；其四，把 RG 推广到其它“强/弱模型对”自然存在的场景，如不同去噪步、不同分辨率子网络的组合；其五，构建 attention sink 理论在图像生成中的统一框架，把寄存器令牌、in-context token、text-token sink 统一起来理解。
复现评估
复现性中等偏上但并非完全开放。论文给出了核心实验设置：pDiT 基于标准 DiT 架构叠加常用 Transformer 改进，用 flow matching 在 ImageNet 256×256、patch size 16 训练，寄存器令牌为附加可学习 token 且不参与损失；RG 的推理与训练以伪代码 Snippet 1–3 清晰给出（含 $z=tx+(1-t)\epsilon$、$v=(x-z)/(1-t)$ 等关键公式），超参 $p$、$w_{rg}$、$w_{cfg}$ 的选取在附录 B.2 与 Tables 4–5 中给出。但论文未明确提及是否开源代码或检查点；为达到 Tables 6–7 的最优结果，作者使用了原始 JiT 检查点并做 fine-tune，复现者需先获取这些检查点；计算资源未明确披露（训练 H/G 规模像素空间 DiT 成本高昂）；部分诊断实验（如线性探测的 probe 训练细节、TV 的具体计算）散落在附录，个别参数仍需推断。总体而言，核心机制分析与小规模实验可复现，但完整复现 SOTA 数字需要相当算力和未公开的 JiT 检查点。
论文图表
图对比了输入图像在多种模型下的注意力图。与 ViT（DINOv2、RAE）在背景等低信息区域出现明显异常不同，SiT、pDiT 这类扩散 Transformer 的注意力主要集中在主体物体上，没有 ViT 式的背景伪影。这说明 DiT 不像 ViT 那样在低信息区产生注意力异常。
这张图是全文的出发点：它直观地证明 DiT 没有 ViT 的注意力异常问题，从而引出“既然没有异常，寄存器令牌是否还有用”这一核心反直觉问题。