跳转至

Test-Time Registers as Global Priors for Tokenized Image Generation

会议: ECCV 2026
论文: ECCV 原文
项目: https://y-research-sbu.github.io/RegToken
领域: 图像生成
关键词: 测试时寄存器 / Tokenized图像生成 / 注意力槽 / 全局先验 / 视觉Transformer

一句话总结

本文提出免训练的 RegToken 框架,通过归一化特征范数(NFN)定位注意力槽层级、借助 TokenRank 提取寄存器子空间并执行投影守恒更新,将原本被视作注意力伪影的高范数寄存器特征转化为全局低频先验,在完全冻结 ViT 主干与解码器的情况下显著提升了紧凑一维离散 Token 图像生成的质量与对齐精度。

研究背景与动机

基于自注意力机制的视觉 Transformer(ViT)与多模态大模型在深层特征处理中普遍存在“注意力槽”(attention sinks)现象:极少数 token 会无序地吸收绝大部分注意力权重,并积累异常巨大的激活范数。早期研究多将其视为网络缺乏垃圾回收机制导致的数值伪影,并通过在训练阶段显式引入寄存器 token(Register Tokens)来吸收这些冗余激活以稳定稠密特征表示;后续工作进一步发现,无需重新训练,仅在测试时利用少量离散神经元(register neurons)即可在冻结模型中复现类似效应。然而,学界以往对寄存器的探索大多停留在可解释性分析、注意力图平滑或线性探测等判别式评测层面,尚未回答一个核心问题:这些在注意力槽中汇聚的高范数表征,是否包含可直接复用于下游生成任务的结构化信息?

深入分析该表征的频域特性后可以发现,注意力槽中的离群 token 并非不可控的数值噪声,而是系统性地编码了图像的平滑全局语义与低频场景统计信息(如主色调、全局光照与粗粒度空间布局)。在紧凑一维离散视觉 token 生成(如 TiTok、HCT)架构中,模型仅依赖数十个离散 token 重构高分辨率图像,缺乏轻量级、无需微调的全局上下文调控信号。现有方案要么依赖判别导向强烈的 [CLS] token 作为条件,要么直接依赖高成本的全量测试时优化,容易在微小生成步数内陷入局部模式崩塌或失去整体场景和谐度。

面对将冻结主干中的寄存器结构转化为即插即用生成先验的挑战,本文的关键切入点在于避开繁琐的手工启发式搜索,建立一整套无须训练的层级定位、子空间投影与离散码本对齐机制。核心 idea:利用归一化特征范数(NFN)与 TokenRank 自动定位 ViT 中注意力槽涌现的层级与特征子空间,通过投影守恒插值将离群特征凝练为紧凑的寄存器全局先验,在完全冻结主干和解码器权重的条件下直接引导离散 Token 图像生成。

方法详解

整体框架

RegToken 的核心流程分为“寄存器子空间定位与提取”和“生成管线全局先验注入”两个核心阶段。在定位阶段,模型仅需前向遍历冻结的视觉 Transformer 主干,通过归一化特征范数(NFN)自动识别注意力槽显著放大的候选层级,并在局部窗口内锁定离群 patch token,随后结合基于马尔可夫链的 TokenRank 排序筛选出主导激活的寄存器神经元集合;在特征整合阶段,利用正交投影与均值守恒机制,将离群 patch 中的寄存器子空间能量平滑转移至新插入的测试时寄存器 token 中;最后,在紧凑一维 token 生成阶段,将凝聚了全局低频信息的寄存器特征对齐并映射至离散码本空间,作为初始先验或软偏置直接注入冻结的解码器中指导图像生成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像与冻结ViT主干"] --> B["NFN层级与神经元定位<br/>前向计算NFN范数比筛选候选层"]
    B --> C["TokenRank子空间重要度量化<br/>马尔可夫转移矩阵计算全局中心度"]
    C --> D["投影与守恒特征插值<br/>正交重分配离群patch至寄存器token"]
    D --> E["离散码本对齐与先验注入<br/>映射至一维Token生成管线z0"]
    E --> F["冻结解码器生成目标图像"]

关键设计

1. NFN 引导的层级与神经元子空间定位:摆脱人工启发式搜索并精准捕捉注意力槽

以往的测试时寄存器方法高度依赖人工设定的层级窗口与激活阈值,跨模型泛化能力薄弱。为实现通用的自动化层级定位,本文引入归一化特征范数(Normalized Feature Norms, NFN)。对于主干网络中任意投影权重 \(W\) 及其输入特征 \(z_{\text{in}}(x)\),NFN 通过对比真实特征响应与同尺度高斯随机输入 \(\tilde{z}_{\text{in}}(x)\) 响应的二范数之比来评估特征放大程度:

\[NFN(W; x) = \frac{\|W z_{\text{in}}(x)\|_2}{\|W \tilde{z}_{\text{in}}(x)\|_2}\]

通过在验证子集上取各头中位数并最大化子模块得分 \(S_\ell = \max_m NFN_{\ell, m}\),系统自动选取激活响应最集中的前 \(\kappa\) 个层级作为候选层 \(\mathcal{L}_{\text{cand}}\)。在候选层邻近窗口内,利用局部范数统计锁定持续处于高范数的离群空间位置集合 \(\Omega_\ell\),并沿前置层级累积绝对激活幅度,精准提取出前 \(k\) 个主导通道构成的寄存器神经元集合 \(\mathcal{R}_\ell = \text{top-}k(\{r_d\})\)。

2. Value-TokenRank 驱动的注意力重要度量化:区分真实信息流通与伪影累积

注意力矩阵中的高权重容易造成“虚假繁荣”——某些 token 虽然分配到大量注意力权重,但其对应的 value 向量模长极小,并不承载有效语义。为解决这一脱节,方法将注意力矩阵 \(A\) 视作离散时间马尔可夫链的转移概率矩阵,求解其平稳分布 \(\pi^\top = \pi^\top A\) 得到每个 token 的 TokenRank 全局中心度;同时,将各位置的写入质量定义为 \(WriteMass(t) = \sum_i A_{i,t} \|V_t\|_2^2\)。通过联合监控马尔可夫链的第二大特征值 \(\lambda_2\)(谱间隙 \(1-\lambda_2\) 越小表示信息混合越慢、注意力汇聚越强烈),模型能量化各通道的实际全局流通贡献,为跨多层寄存器特征加权融合提供了理论支撑。

3. 投影与守恒的特征重分配插值:无损吸收全局低频能量并抑制空间伪影

在确定寄存器神经元子空间 \(\mathcal{R}_\ell\) 与离群 patch 集合 \(\mathcal{P}\) 后,直接抹除离群值会导致特征剧烈抖动,而简单复制则无法净化局部空间表征。本文构造子空间正交投影矩阵 \(P_{\mathcal{R}}\),计算离群 token 在寄存器子空间上的平均响应 \(\bar{v}_{\mathcal{R}} = \frac{1}{|\mathcal{P}|} \sum_{p \in \mathcal{P}} P_{\mathcal{R}} V_p\)。随后,通过参数化尺度 \(s\) 与守恒因子 \(\alpha\) 进行两极特征迁移:

\[\begin{aligned} V_{t_{\mathrm{reg}}} &\leftarrow V_{t_{\mathrm{reg}}} + s\,\bar{v}_{\mathcal{R}} \\ V_p &\leftarrow V_p - \alpha P_{\mathcal{R}} V_p, \quad \forall p \in \mathcal{P} \end{aligned}\]

该更新规则将原本分散且阻塞局部注意力的全局离群分量正交剥离,统一汇聚至新插入的测试时寄存器 token \(t_{\text{reg}}\) 中,在数值上近似保持一阶投影能量守恒,既平滑了常规 patch 的特征分布,又提炼出纯净的全局场景基底。

4. 离散码本对齐与轻量级全局先验注入:即插即用赋能一维 Token 生成解码器

为了让提取的连续寄存器特征兼容冻结的离散生成模型(如 TiTok/HCT),RegToken 采用白化与余弦相似度最近邻搜索(或轻量 Procrustes 正交对齐),将融合后的寄存器表示 \(\tilde{r}\) 投射到离散码本向量空间中。在解码阶段,提供两种灵活的注入范式:其一为首位标记插值(Init-prior),即直接将生成序列的首个 token 替换为插值结果 \(z_0 \leftarrow (1-\gamma)z_0 + \gamma \tilde{r}\);其二为全流程软偏置(Soft-bias),通过余弦相似度指数加权 \(p'(c) \propto p(c) \cdot \exp(\beta \cos(e_c, \hat{u}))\) 引导离散采样。在可选的测试时微调阶段,仅对注入的单个全局 token \(z_0\) 执行少量梯度更新,其余 token 与解码器权重完全冻结,实现了极高效率的全局语义对齐。

损失函数 / 训练策略

本方法主体为完全免训练(training-free)架构,仅在前置离线阶段使用 1000 张 ImageNet 验证集图像前向推断计算 NFN 统计量以固定关键层与通道索引。在包含测试时优化(token opt.)的实验配置下,保持预训练 HCT 生成解码器权重与后置 token 序列 \(z_{1:T}\) 完全固定,仅将 \(z_0\) 设为可学习变量,利用冻结的 CLIP 文本-图像对齐损失函数反向传播极少迭代步数(如 50 步),学习率设为小步长更新,避免破坏已学习到的离散先验分布。

实验关键数据

主实验

主实验在 ImageNet-1k 评估基准上进行,采用冻结的 HCT 解码器配合 VQ-LL-32 紧凑码本,在相同的测试时解码超参数下对比不同全局先验对图像生成真实感(FID-5k、IS)及文本-图像图文对齐质量(CLIPScore、SigLIP)的影响。

骨干网络 注入先验方法 FID-5k (↓) IS (↑) CLIP (↑) SigLIP (↑)
DINOv2-L/14 HCT 无先验 (Beyer et al.) 21.2 281 0.40 3.5
DINOv2-L/14 HCT + 随机噪声先验 22.5 278 0.39 3.4
DINOv2-L/14 HCT + [CLS] 先验 20.5 281 0.39 3.6
DINOv2-L/14 HCT + 传统测试时寄存器 (TTR) 21.5 283 0.40 3.6
DINOv2-L/14 HCT + 监督训练寄存器 (Darcet et al.) 20.4 285 0.41 3.8
DINOv2-L/14 HCT + RegToken (本文) 20.1 289 0.45 3.9
OpenCLIP-L/14 HCT + [CLS] 先验 20.8 283 0.40 3.5
OpenCLIP-L/14 HCT + 传统测试时寄存器 (TTR) 21.1 284 0.41 3.7
OpenCLIP-L/14 HCT + RegToken (本文) 20.3 287 0.42 3.8

消融实验

针对 NFN 定位所提取的寄存器神经元数量以及不同先验数据来源,论文在特征范数分离度与生成迁移性上进行了充分的消融对比。

表 1:特征范数与注意力分离度(验证寄存器子空间吸收离群激活的有效性):

骨干网络 寄存器定位配置 范数差距 (Norm Gap) 注意力差距 (Attention Gap)
OpenCLIP 传统基准 (50 神经元) 62.03 0.58
OpenCLIP NFN 引导 (8 神经元) 61.36 0.52
OpenCLIP NFN 引导 (16 神经元) 64.42 0.55
OpenCLIP NFN 引导 (50 神经元) 70.09 0.58
DINOv2 传统基准 (50 神经元) 468.83 0.60
DINOv2 NFN 引导 (8 神经元) 485.23 0.59
DINOv2 NFN 引导 (16 神经元) 494.63 0.63
DINOv2 NFN 引导 (50 神经元) 633.71 0.66

表 2:先验来源对生成指标的影响(验证寄存器先验承载类别级与实例级全局信息):

先验输入源 FID-5k (↓) IS (↑) CLIP (↑) SigLIP (↑)
同图像寄存器先验 (Same-image) 20.1 289 0.45 3.9
同类别跨图像先验 (Cross-image, same class) 20.6 285 0.41 3.8
随机类别图像先验 (Random-image, any class) 21.1 283 0.40 3.7

关键发现

  • 低频信息高度富集:频域 FFT 分析表明,在 DINOv2 与 OpenCLIP 中,测试时寄存器特征的低频能量占比(\(\rho_{1:15}\) 分别达 0.507 与 0.524)与低频平坦度(\(\bar{\text{dB}}_{\text{low}}\))均显著超越 [CLS] 和 Patch 平均特征(\(p < 10^{-16}\)),证明其本质上承载平滑的宏观场景基调而非局部细粒度纹理。
  • 抑制伪影并加速注意力混合:NFN 选取的神经元能更彻底地降低离群 patch 比例(DINOv2 上中位数离群占比减少 -0.0312,显著优于传统测试时寄存器的 -0.0215),并将马尔可夫链第二特征值大幅降低(\(\Delta\lambda_2 = -0.0946\)),使得注意力网络混合更充分、避免注意力陷入局部奇点。
  • 显著提升测试时优化效率:在 HCT 解码测试时优化中,RegToken 将达到预设 CLIP 相似度阈值所需的步数从 74 步缩减至 52 步(耗时骤降近 30%),且在 1000 组随机种子下的表现方差明显低于其他基准,表明其提供了极佳的全局解空间初始化。

亮点与洞察

  • 将“负面伪影”转变为“正面生成资产”:过去的研究往往仅将注意力槽作为一种需要被消除或平滑的计算异常,本文首次洞察到其低频能量汇聚的本质,并将其作为紧凑全局先验赋能冻结解码器,视角转换极具启发性。
  • 免训练的几何正交守恒机制:提出的投影与守恒更新公式巧妙地在局部离群 patch 与全局寄存器 token 之间构建了一阶统计守恒,既消除了原图的注意力塌陷,又完整保留了全局信息。
  • 跨模型架构与生成任务的即插即用性:无需修改主干网络权重,也无需对生成解码器做针对性微调,该流程可无缝迁移至不同视觉自监督主干(DINOv2、OpenCLIP)以及不同生成范式(HCT、MaskGIT-VQGAN)。

局限与展望

  • 依赖预计算的特征统计集:NFN 自动化定位层级需要小规模验证集(约 1000 张图)进行前向推断以统计中位数,尚无法在单张未知域图像冷启动下完成完全自适应的即时层级定位。
  • 单 token 全局容量瓶颈:当前生成端仅将寄存器特征融合成单个全局 token(\(z_0\))进行注入,对于包含多个复杂前景与多光源的宏观复杂场景,单 token 难以完整表达非均质的空间先验。
  • 未来方向:未来可探索将多个层次的寄存器解耦为空间自适应的多全局先验序列,并尝试在连续扩散模型(Diffusion Transformers, DiT)的潜空间注意力槽中验证该先验控制机制。

相关工作与启发

  • vs Darcet et al. (ICLR 2024, Vision Transformers Need Registers):Darcet 等人通过在预训练期显式追加可学习 register token 来吸收注意力离群值;本文则完全无需训练,直接从预训练冻结主干中提取注意力槽特征,并进一步拓展到了生成任务的先验引导。
  • vs Jiang et al. (NeurIPS 2025, Vision Transformers Don't Need Trained Registers):Jiang 等人提出了基于启发式规则定位离群通道的测试时寄存器(TTR),但主要停留在特征平滑与分类探针评估;RegToken 引入 NFN 与 TokenRank 实现了更理论自洽的自动化定位,并在图像生成这一高阶下游任务上展现出实用价值。
  • vs Beyer et al. (ICML 2025, HCT):HCT 依赖紧凑 1D 离散 token 执行图像生成,但由于缺少强先验,测试时优化速度慢且易不稳定;RegToken 为其提供了一个高匹配度的低频全局基底,显著优化了生成收敛轨迹与对齐指标。

评分

  • 新颖性: ⭐⭐⭐⭐ [首次将注意力槽与测试时寄存器转化为图像生成的低频全局先验]
  • 实验充分度: ⭐⭐⭐⭐⭐ [涵盖频域分析、范数分离度、ImageNet主生成指标及跨图像泛化消融]
  • 写作质量: ⭐⭐⭐⭐⭐ [逻辑链条严密,从现象诊断、理论建模到实验印证层层递进]
  • 价值: ⭐⭐⭐⭐ [为免训练的冻结表征重利用与轻量化图像生成控制提供了实用范式]