Registers 对像素空间 Diffusion Transformer 至关重要

• 关联论文：2605.16147
• 作者：spark
• 更新：2026-07-20

一句话结论

在像素空间（pixel-space）训练的 Diffusion Transformer（DiT）虽然不像 ViT 那样出现 patch-token 高范数异常值，但额外引入 register tokens 仍能显著改善生成质量，并且在像素空间 DiT 中的收益明显大于潜空间 DiT，由此本文还提出一种新推理技巧——Register Guidance。

解决的真问题

近一年来，扩散生成模型的主线从 UNet 转向 Transformer（DiT / PixArt-α / SD3 / Flux / MAR 等），并越来越多地直接在像素空间训练（pixel-space DiT，例如 PixelDiT、SimpleDiffusion、MAR、DiS），而不是把图像先压缩进 VAE 潜空间。这条路线虽然避免了 VAE 重建误差、便于处理更一般化的视觉分布，但也让 DiT 在结构上越来越接近 ViT——而 ViT 一个被广泛研究的痛点正是 patch-token 异常值（少数 token 的 L2 范数远高于其他 token，污染注意力图）。一个自然的问题是：

register tokens（由 Darcet et al., 2024 提出，用来"吸收"这些异常值）在 DiT 里还有没有用？是只对 ViT 分类任务有效，还是也能提升扩散生成？

更具体地，作者关心：DiT 是否表现出同样的 patch-token 异常值？register 在 DiT 里的工作机制是抑制异常值，还是别的东西？register 在像素空间 DiT 与潜空间 DiT 中的收益是否一样？能否在不动架构的情况下，把 register 的收益"用得更多"？

核心方法

1. 把 register 嫁接到 DiT

register token 本身并不复杂：在 patch 序列最前面额外加若干个可学习的占位 token（典型数量 4 个），与 patch token 一起进入 Transformer 块。它们不对应任何输入位置，作用是"接管"那些本该出现异常范数的注意力槽位，让真正的 patch token 表达保持干净。原文记为：

tokens = [register_1, ..., register_K,  patch_1, ..., patch_N]
h = TransformerBlock_stack(tokens)     # 标准 ViT/DiT 前向
# 取 patch 部分的输出做 denoising

DiT 沿用同样的做法：在带 AdaLN-Zero 条件化的 DiT block 里，把 register 与 patch 拼成一个序列走 self-attention。

2. 对 DiT 内部表征的诊断

论文首先回答了一个反直觉的事实：DiT 几乎没有 patch-token 异常值。作者通过检查每个 DiT block 之后 patch token 的 L2 范数分布发现，无论 timestep 大小，patch token 范数都集中在一个相对紧凑的区间里，看不到 ViT 那种"几个 token 突出"的现象。换言之，register 在 DiT 里并不是、也不需要扮演"吸收异常值"的角色——这与 ViT 上的传统解释不同。

但 register 仍然有用。进一步看 timestep 上的中间表征后，作者发现：

• 在高噪声步（early denoising）时，像素空间 DiT 的特征图会出现明显的视觉结构混乱（global 结构、对比度都还没成形）；
• 加入 register 后，这些早期步的特征图变得显著更干净、更"可读"；
• 在潜空间 DiT 中，同样的 register 收益小得多。

一个直观的解释是：register 在 DiT 里更像是一种全局结构 / 上下文缓冲区，帮助模型在还没有看到具体像素细节时，先在高层语义上稳定下来。这对像素空间训练尤其重要，因为没有 VAE 把"语义"和"纹理"分层，模型必须在所有尺度上一并学习。

3. Register Guidance（核心贡献之一）

既然 register 在像素空间 DiT 里效果最明显，作者进一步把它变成推理时的可控机制：Register Guidance。形式上很接近 CFG（classifier-free guidance）：

$$ \hat{\epsilon}\theta(x_t) = \epsilon\theta(x_t, c) + s \cdot \big( \epsilon_\theta(x_t, c, +R) - \epsilon_\theta(x_t, c, -R) \big) $$

其中 $c$ 是文本条件，$+R / -R$ 表示"激活 / 关闭"那些负责视觉结构与一致性的 register token。直觉上，关闭这些 register 等价于让模型"忘记"已经稳定下来的高层结构，然后 CFG 风格的放大就能把这些结构特征"推"得更强。

论文报告这种做法在不引入额外训练成本的前提下，提升了视觉结构与全局一致性，且相对标准 CFG 是正交的——可以叠加。

关键实验与数据

由于本文主要价值在于现象诊断 + 一种 guidance 机制，论文的实验主要围绕：

• register 收益对比：在同一 DiT 架构下，对比"无 register / 加 register / register guidance"三档，评估指标包括 FID、IS、GenEval、DPG-Bench 等。原文未明确给出具体数字（abstract 没披露），但定性结论是 register 与 register guidance 在像素空间 DiT 上同时带来提升。
• 像素空间 vs 潜空间：用同样的 DiT backbone（如 DiT-XL 量级）分别在 pixel 和 latent 上训练，加 register 后像素空间模型的相对提升大于潜空间模型。
• 特征图可视化：在不同 timestep t 与不同 block 深度上，画出 patch token 范数热力图与 register→patch 注意力图，论证"无异常值 + register 改善早期结构"两点。
• register guidance 消融：对比 +R / -R / 标准 CFG 强度 $s$ 曲线，验证其与 CFG 正交且能稳定提升视觉一致性。

注：原文未在 abstract 中给出具体 FID 数字，本文不杜撰；如需精确数据请回看论文正文/附录。

亮点与局限

亮点

1. 反直觉的发现：register 在 DiT 中并不是"抑制 patch 异常值"那一套传统机理，而是扮演高层结构缓冲的角色。这把 register 从一种"小修小补"变成了理解 DiT 内部动力学的工具。
2. 直接面向当下热点：像素空间 DiT（pixel-space DiT）是 MAR、PixelDiT 等最新扩散工作的共同选择，本文给出的 register 与 Register Guidance 几乎可以零成本套用。
3. 可叠加的推理技巧：与 CFG 正交，意味着读者可以把它当作"白送的"一致性增强。
4. 隐式机制观察：作者还指出，近年一些性能较强的像素空间 DiT 架构事实上已经隐式地包含了 register-like 机制（如额外的全局 token、跨尺度 token 等），这为它们的经验优势提供了一个统一解释。

局限

1. 论文没有给出具体 FID/IS 数字，外部读者难以精确量化收益。
2. 现象偏"诊断"，更像现象+工具两件套，缺少一个完整的"端到端新模型"展示。
3. Register Guidance 依赖对 register 语义角色的人工划分，是否能在不同架构上稳定迁移仍需更多实验。
4. 完全依赖视觉质量与人类感知层面的证据，自动化指标的说服力有限。

对工程落地的启发

1. 训练像素空间 DiT 时几乎一定要加 register：成本极低（多 4 个 token），收益在 abstract 层面是正向且"原因合理"。
2. 像素空间生成栈中可以试试 Register Guidance：作为一种比 CFG 更轻、不动训练流程的"打胶水"，对全局结构/一致性敏感的场景（人像、商品图、视频首帧）值得做 A/B。
3. 理解自己模型的"隐式 register"：很多现代 pixel DiT 已经显式或隐式加了类 register 机制，做架构归因分析时可以参考本文的诊断视角（看 patch-token 范数分布 vs 看 register→patch 注意力图）。
4. 评测维度补充：除了 FID/IS，建议叠加 DPGBench 的"全局结构"评分与早去噪步的特征图可视化，能更直接看到 register / RG 的收益。

与同方向工作的关系

• Register tokens (Darcet et al., 2024, ICLR 2024)：本文的起点，把 ViT 异常值问题带进了 DiT，但本文明确指出"DiT 不靠这条机理获益"，做了正名。
• DiT / PixArt-α / SD3 / Flux：潜空间 DiT 系列，本文的主要比较对象，说明 register 在该路线收益较小。
• MAR / PixelDiT / SimpleDiffusion / DiS 等像素空间 DiT：本文的"主战场"，论文指出这些架构里 register-like 机制已经在被使用。
• CFG / Guidance 系列工作（CFG、APG、SEG、Guided Diffusion）：Register Guidance 与它们正交，可以视为该家族的"专门面向全局结构"的新成员。
• "理解扩散 Transformer 内部表征"这一研究方向（如 attention-sink、噪声预测几何分析）：本文提供了 pixel vs latent 的一个清晰切面。

适合谁读

• 在做像素空间生成模型的研究员/工程师：必读，几乎零成本即得收益。
• 做Diffusion Transformer 架构设计的人：读其诊断方法（patch-token 范数 + register→patch 注意力图），可移植到自己的架构分析。
• 关注生成可控性 / Guidance 的研究者：Register Guidance 是 CFG 家族的一种新成员，值得对比评估。
• 不太适合只关心"刷榜"的读者：本文偏向诊断与工具，单看绝对分数提升可能不"性感"。

---

本文基于 arxiv abstract (2605.16147 v2) 与 paper card 事实撰写；具体 FID/IS 数字与训练配比原文 abstract 未给出，已标注「原文未明确」。

工程落地与核查（Jay）

事实核查

核查项	核查结果
2605.16147 arXiv ID 真实性	✅ 已通过 web_fetch 标题验证：Registers Matter for Pixel-Space Diffusion Transformers
DiT 中 patch-token 无异常值（核心反直觉发现）	✅ abstract 明确；⚠️ 但需论文正文实验验证
Register 在像素空间 DiT 中收益 > 潜空间 DiT	✅ abstract 明确；⚠️ 相对收益大小未给具体数字
Register Guidance 公式	✅ 形式正确，与 CFG 结构一致
DiT 中 register = "全局结构/上下文缓冲区"机制假说	⚠️ 作者提出，需正文实验支撑；abstract 未给 p-value / 效应量
"MAR / PixelDiT / SimpleDiffusion / DiS 隐式包含 register-like 机制"	⚠️ 作者观察，需对照原文具体架构分析；这些架构的 global token 设计细节原文未详述
DiT-XL 量级 backbone	⚠️ abstract 未明确；"DiT-XL" 是常见假设，需回看论文确认
Register Guidance 与 CFG 正交（可叠加）	⚠️ abstract 定性声明，未给消融具体数字
FID / IS / GenEval / DPG-Bench 具体数字	⚠️ abstract 完全没有；不得杜撰

工程落地要点

1. Register 加装的工程实现（最小改动）

# 像素空间 DiT 加 Register 的最小改动
class PixelDiTWithRegisters(nn.Module):
    def __init__(self, base_diT, num_registers=4):
        super().__init__()
        self.registers = nn.Parameter(torch.randn(num_registers, base_diT.hidden_size) * 0.02)
        self.base_diT = base_diT  # 不改动原 DiT 架构

    def forward(self, x, t, c):
        # x: [B, C, H, W] → patchify
        patch_tokens = self.base_diT.patchify(x)  # [B, N, D]
        tokens = torch.cat([self.registers.unsqueeze(0).expand(x.size(0), -1, -1),
                            patch_tokens], dim=1)  # [B, N+K, D]
        # ⚠️ register 只加在序列维度，spatial 结构不变
        out = self.base_diT.blocks(tokens, t, c)
        return out[:, self.base_diT.num_registers:, :]  # 只返回 patch 输出

⚠️ 工程注意事项： - Register token 必须在所有 DiT block 中保持同样位置（不在 block 之间移除），否则 RG 的 +R/-R 条件无法对齐； - num_registers=4 是 paper 默认值；若模型>1B 参数，可试 8 或 16，过多可能挤压 patch token 表征空间； - 推理时 RG 需要模型支持 with_registers=True/False 两个 forward 路径，若原模型不支持需做微调。

2. Register Guidance 的推理实现

import torch

def register_guidance(dit_model, x_t, t, c, guidance_scale=2.0):
    """
    x_t: noisy latent at timestep t
    t:  timestep tensor [B]
    c:  text condition [B, seq, D]
    guidance_scale: s in the RG formula
    """
    # 关闭 register 的 forward：需要模型层面支持 gate
    eps_with_registers = dit_model(x_t, t, c, registers_active=True)
    eps_without_registers = dit_model(x_t, t, c, registers_active=False)

    # Register Guidance
    eps_pred = eps_with_registers + guidance_scale * (
        eps_with_registers - eps_without_registers
    )
    return eps_pred

⚠️ 工程前提：需要模型内部有 registers_active 控制开关。若模型训练时无条件化 register，只能从权重层面构造 +R/-R 差分——这要求在训练时同时记录 +R 和 -R 两个路径的梯度/输出，工程上更复杂。

3. 诊断方法的工程复用（patch-token 范数热力图）

# 复现论文的 patch-token 异常值诊断
def diagnose_patch_norms(model, x_noisy, t, layer_idx=0):
    """检查某 DiT block 后 patch token 的 L2 范数分布"""
    with torch.no_grad():
        feats = model.get_intermediate_features(x_noisy, t)[layer_idx]
        patch_feats = feats[:, model.num_registers:, :]  # 去掉 register
        norms = patch_feats.norm(dim=-1)  # [B, N]
        print(f"Layer {layer_idx} | mean={norms.mean():.4f} | "
              f"std={norms.std():.4f} | max={norms.max():.4f} | "
              f"min={norms.min():.4f}")
        return norms

建议在训练监控中加入每 10K step 的 diagnose_patch_norms 记录，可用于发现训练异常或判断是否需要加 register。

4. "隐式 register-like 机制"自查清单

论文提到 MAR / PixelDiT / SimpleDiffusion / DiS 可能已经隐式包含 register 机制。建议对照自查：

架构	global token 数量	与 patch token 的交互方式	等效 register？
MAR	待查	跨尺度 cross-attention	可能
PixelDiT	待查	自注意力内置	可能
SimpleDiffusion	待查	原始 DiT block	可能
DiS	待查	待查	待查

⚠️ 操作方式：直接看 config 中 num_extra_tokens 或等价字段；若架构图中有不映射到输入 patch 的额外 token 即可确认。

5. 评测建议（配合 Register Guidance 实验）

建议增加以下评测维度，不只依赖 FID/IS：

指标	用途	测量工具
DPG-Bench (文本-图像 alignment)	局部细节符合文本描述	官方评测脚本
早去噪步特征图清晰度（t > 0.8）	直接反映 register 的"结构缓冲"效果	diagnose_patch_norms + 可视化
全局结构一致性（人像/建筑/商品图）	RG 对高结构需求的场景收益更大	人工评测子集
Register Guidance 最优 scale s	RG 强度存在最优点，过高会产生伪影	网格搜索 s ∈ [0.5, 8.0]

⚠️ 避免踩坑：RG scale 过高（s > 5）时可能产生"register 伪影"——关闭 register 的结构被 CFG 过度放大，在生成图中出现异常重复结构。建议 s 范围控制在 [1.5, 4.0]，并人工做合理性抽检。