RATS! Patches Talk Through Registers: Emergent Parts in Register Attention Transformers¶

会议: NeurIPS 2026
arXiv: 2606.14701
OpenReview: https://openreview.net/forum?id=Ylj622Nr12
代码: https://github.com/yangtiming/RATS
领域: 自监督/表示学习
关键词: 自监督学习, 视觉Transformer, 寄存器Token, 部件发现, 无监督分割

一句话总结¶

提出寄存器注意力Transformer（RATS），通过头内硬划分的寄存器Token构建“压缩-交互-广播”三阶低秩注意力瓶颈，在纯自监督DINO目标下诱导寄存器自发特化为语义部件，在5项无监督分割基准上超越既有基线超12 mIoU。

研究背景与动机¶

人类在观察视觉世界时具有天然的组合认知本能：看到一只鸟时，看到的不仅是全局的“鸟”这一分类标签，而是头部、翅膀和爪子等跨物种复用的结构化部件装配体；这种部件级语义在猫耳与兔耳、自行车轮与汽车轮之间高度通用与可迁移。以DINO为代表的自监督视觉模型虽然通过特征PCA投影展现出潜在的跨图像部件对应关系，但标准Vision Transformer（ViT）在架构上缺乏显式的部件级分组与表征归纳偏置。其全局全连接自注意力机制最终将整图信息压缩至单一的全局分类Token（[CLS]），而密集的Patch Token作为扁平特征袋缺乏显式聚类承诺，导致部件级结构只能隐式弥散而无法作为可解释的显式表征单元被导出。

为了将Patch分组为语义单元，以往方法尝试了多种路径但各有瓶颈：基于超像素的方法依赖低级图像边界，缺乏跨全局上下文辨识部件身份的能力；基于槽位注意力的目标中心方法（如Slot Attention、AdaSlot）主要依靠像素重建驱动，强竞争机制容易割裂连贯的语义区域，诱导出的多为整物掩码而非局部语义部件；GroupViT依赖图文对监督，特化的是文本类别而非纯视觉部件；Perceiver IO将输入压缩为潜在数组，但所有注意力头联合建模导致潜在变量高度纠缠，无法形成独立的语义基元。而常规ViT引入Register Token（寄存器Token）最初仅用于吸收和丢弃特征伪影与异常背景，其潜在的表征与路由潜力未被挖掘。

面对“如何在纯无监督条件下诱导视觉模型自发展开部件级语义解耦”这一矛盾，本文重新审视了注意力通信机制：既然[CLS]是极端的全局信息汇聚点，完全可以将其解构为一组具备局部特化能力的中间路由通道。核心 idea：将Transformer自注意力重塑为头内排他划分的寄存器注意力瓶颈，通过“压缩-交互-广播”（$L \to N \to N \to L$）三阶段通信强制Patch间经由独立子空间的寄存器交换信息，在不加任何部件标注与辅助损失的前提下让寄存器自发涌现为跨类别通用的结构化语义部件。

方法详解¶

整体框架¶

RATS保留了ViT的主干骨架与预归一化残差连接结构，其核心改造在于将标准多头自注意力（MHSA）替换为寄存器注意力模块（R-Attn）。在输入层，图像被切分为 $L$ 个Patch Token $P \in \mathbb{R}^{L \times D}$ 并叠加位置编码，同时拼接 $N$ 个不带位置编码的可学习寄存器Token $R \in \mathbb{R}^{N \times D}$ 构成联合序列 $X=[R; P]$。在每个Transformer块内，信息通过投影、头内划分、三步注意力（压缩、交互、广播）并在头间拼接输出投影后，更新寄存器与Patch表示。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列<br/>X = [R; P] (N个寄存器 + L个Patch)"] --> B["共享QKV线性投影与多头排他切分<br/>每个头分配独立子集 R^h 与共享 P"]
    B --> C["阶段1：压缩 (Compress)<br/>R^h 跨注意力聚合局部 Patch 信息"]
    C --> D["阶段2：交互 (Communicate)<br/>R^h 内部自注意力交互，部件上下文绑定"]
    D --> E["阶段3：广播 (Broadcast)<br/>更新后的寄存器向所有 Patch 广播写回"]
    E --> F["多头拼接与输出投影<br/>恢复联合序列维度更新 [R; P]"]

关键设计¶

1. 头内排他寄存器划分：构建解耦的低秩注意力子空间

标准ViT的注意力头往往学到高度冗余或完全混合的特征，若直接让所有注意力头共享并更新全局寄存器池，极易陷入类似Perceiver IO的表示纠缠困境。针对这一问题，RATS将全部 $N$ 个寄存器Token在 $H$ 个注意力头之间进行互斥的硬划分（Hard-Partitioning），每个头仅拥有独立且互不相交的 $n = N/H$ 个寄存器子集 $R^h = R[hn : (hn+n)]$。在单头计算中，Patch投影在所有头之间共享，而寄存器仅在所属头的低维子空间（维度 $d = D/H$）中运算。由于每个头的通信瓶颈被严格限制在至多秩 $n$ 的低秩子空间内，不同头被迫捕获图像中不同区域与不同粒度的空间语义，杜绝了跨头干扰，为部件特化的自发涌现奠定了结构先验。

2. “压缩-交互-广播”三步通信瓶颈：实现部件级语义聚合与回写

针对扁平全连接注意力无法显式聚类局部语义的痛点，RATS在每个注意力头内部将Patch之间的直接全连接交互阻断，改为严格的 $L \to n \to n \to L$ 三阶段路由： - 压缩（Compress）：以寄存器投影 $Q_R \in \mathbb{R}^{n \times d}$ 为Query、Patch投影 $K_P, V_P \in \mathbb{R}^{L \times d}$ 为Key和Value执行跨注意力，将全图Patch信息根据与寄存器的语义亲和度汇聚至各个寄存器中： $$\hat{R} = \mathrm{softmax}\left(\frac{Q_R K_P^\top}{\sqrt{d}}\right) V_P \in \mathbb{R}^{n \times d}$$ - 交互（Communicate）：如果仅有压缩与广播，Patch之间只能单向透传，无法捕获共现部件之间的上下文关联。因此在寄存器之间引入参数共享的自注意力通信，使各个已初步聚合区域信息的寄存器相互感知、协同绑定更大尺度的部件组合： $$\bar{R} = \mathrm{softmax}\left(\frac{\hat{R}\hat{R}^\top}{\sqrt{d}}\right) \hat{R} \in \mathbb{R}^{n \times d}$$ - 广播（Broadcast）：以Patch投影 $Q_P \in \mathbb{R}^{L \times d}$ 为Query、更新后的寄存器 $\bar{R}$ 为Key和Value，将带有全局协同上下文的部件表征写回所有Patch： $$O_P = \mathrm{softmax}\left(\frac{Q_P \bar{R}^\top}{\sqrt{d}}\right) \bar{R} \in \mathbb{R}^{L \times d}$$ 单头输出汇聚并拼接后，经由输出投影矩阵 $W_O$ 统一重组为 $[O_R \mid O_P]W_O$，完成本层的Token表征更新。

3. 无监督部件字典构建与推理分割：从寄存器激活到显式语义词表

为验证并利用训练好的寄存器所涌现出的局部语义，模型在最后一层直接评估寄存器与Patch之间的空间选择性。每个寄存器 $R_i$ 与网格上各位置Patch $P_{u,v}$ 计算余弦相似度图 $S_i(u,v) = \frac{\langle R_i, P_{u,v} \rangle}{\|R_i\| \|P_{u,v}\|}$，选取峰值响应前 $\tau$ 显著的寄存器经过高斯平滑滤波后进行逐像素 $\mathrm{argmax}$ 分配，即可在无需任何监督的情况下输出高精度的无监督部件分割图。进一步地，由于物体部件具有跨样本复用性，将这些寄存器特征在数据集级别聚合聚类，可提炼出紧凑的可复用部件字典（Part Dictionary）。字典条目不仅在超类内部表现出严格的部件一致性（如稳定激活不同鸟类的头部），还在跨超类场景下涌现出功能类比特性（如在汽车与自行车间自动迁移“车轮”表征）。

损失函数 / 训练策略¶

RATS完全采用标准DINO自蒸馏目标进行自监督预训练，不引入任何辅助损失、区域对比正则或像素重建约束。输入图像经过全局与局部数据增强裁剪后分别输入学生与教师网络；教师网络参数通过学生参数的指数移动平均（EMA，动量 $\lambda = 0.996 \to 1.0$）平滑更新。最终层全部 $N$ 个寄存器Token的均值 $z = \frac{1}{N}\sum_{i=1}^N R_i$ 被送入DINO投影头，映射为关于 $K=65,536$ 个可学习原型的概率分布，并与教师分布计算交叉熵损失： $$\mathcal{L}_{\text{DINO}} = \sum_{x \in \mathcal{V}^g} \sum_{x' \in \mathcal{V}, x' \neq x} - p_t(x)^\top \log p_s(x')$$ 消融实验证实，跨全部注意头的全局共享DINO损失明显优于各头独立的DINO损失，保证了全局语义与局部部件约束的良好平衡。

在下游密集预测任务微调时（如Mask2Former分割），由于预训练的寄存器瓶颈注意力秩最高为 $n$，为了恢复密集预测对高分辨率全秩表征的表达能力，RATS并联恢复了一条残差全秩自注意力通路：$X_{\text{out}} = (L \to N \to N \to L)(X) + (L \to L)(X)$，同时将预训练的寄存器特征作为Mask2Former的优质初始化Query。

实验关键数据¶

主实验¶

在5项无监督分割基准测试中，以多对一（Many-to-One, M2O）mIoU（%）和调整兰德指数（ARI）为主要评估指标，各基准均在 $512 \times 512$ 分辨率下测试，对比标准ViT、带寄存器的Mamba、超像素模型与槽位注意力模型：

方法	主干网络	寄存器数 $N$	COCO (mIoU)	ADE20K (mIoU)	ImageNet (mIoU)	IN-S919 (mIoU)	PartIN (mIoU)	平均 mIoU	平均 ARI
DINOv1	ViT-S/16	6	8.53	15.43	11.21	48.28	25.58	21.81	8.7
DINOv3	ViT-S/16	6	8.17	13.69	11.15	50.25	27.67	22.19	11.1
DINOv1	ViT-B/16	12	8.55	15.87	11.05	46.94	25.82	21.65	7.9
DINOv3	ViT-B/16	12	10.49	18.46	14.13	56.43	32.57	26.41	12.3
Mamba-R	Mam.R-B/16	12	10.07	18.95	13.44	55.54	30.11	25.62	12.8
SPFormer	ViT-S/16	6	8.37	14.65	12.10	51.92	27.21	22.85	11.5
SPiT	ViT-B/16	12	8.41	12.89	12.37	53.47	27.83	23.00	8.5
AdaSlot	ViT-B/16	33	12.31	23.87	15.49	54.83	30.74	27.45	19.3
RATS (Ours)	RATS-S/16	6	9.91	17.02	13.97	58.84	33.93	26.73	20.4
RATS (Ours)	RATS-B/16	12	12.99	21.45	17.31	67.89	41.06	32.14	22.0
RATS (Ours)	RATS-S/16	96	17.31	26.98	21.87	73.18	46.43	37.15	21.7
RATS (Ours)	RATS-B/16	192	19.32	30.77	23.47	75.39	48.77	39.54	19.6

消融实验¶

在相同参数量与输入设定（NVIDIA A5000, $512 \times 512$, Batch 16, top-$\tau=6$）下评估注意力架构设计的有效性，以及每头寄存器数量 $n$ 对过分割的影响：

架构配置	参数量 (M)	显存 (GB)	延迟 (ms)	相对开销	M2O mIoU (%)	O2O mIoU (%)	ARI
ViT-S/16 (DINO 基线)	21.67	1.08	172.8	$2.12\times$	21.8	19.3	8.7
RATS 基线 ($n=16$)	21.70	0.49	81.7	$1.00\times$	26.6	23.1	19.2
RATS + 可学习交互投影	22.59	0.49	82.6	$1.01\times$	25.5	21.9	17.1
RATS (所有头共享寄存器)	21.70	0.52	91.7	$1.12\times$	25.0	21.2	15.1
Perceiver-IO 风格	21.73	0.31	12.4	$0.15\times$	20.2	18.6	10.0

在单头寄存器数量 $n \in \{8, 16, 24, 32\}$ 消融中，RATS-S（$H=6$）在 $n=16$ 时达到最佳平均 M2O mIoU（36.62%），$n=8$ 时为 35.36% 但具有更高的 O2O mIoU（25.83% 对比 25.60%），表明较小的 $n$ 能有效抑制过分割；而 $n \ge 24$ 时由于表征容量过剩，过分割程度加剧，指标逐渐饱和回落。

关键发现¶

头内硬划分是部件涌现的决定性前提：在消融实验中，一旦去除头内排他约束改为跨头共享寄存器，M2O 从 26.6 骤降至 25.0，ARI 从 19.2 跌至 15.1，证实独立子空间是防止语义过度混合、迫使多头特化到不同部件的关键机制。
参数免训练交互兼具效率与纯度：在“交互”阶段使用无参数的纯自注意力机制表现优于引入额外学习投影矩阵的变体（M2O 26.6 vs 25.5），避免了在低秩通信路径上过拟合。
下游迁移双重增益：微调 Mask2Former 时，RATS-S 在 ADE20K 达到 46.52 mIoU（优于 DINO 的 45.41），COCO 实例分割达到 38.1 $\text{AP}^{\text{m}}$（优于 DINO 的 37.9）；仅替换主干时仍带来净提升，证明增益同时来源于更强的高质量特征表示与寄存器作为语义 Query 的初始引导。
分辨率与训练动态存在相变区间：推理分辨率从 224 增至 640 时分割性能单调提升，超过 640 后因 Patch 过度稀释而饱和；训练在 160-180 epoch 达到峰值性能，超过该区间后过度蒸馏导致寄存器多样性减退。

亮点与洞察¶

寄存器功能范式转换：将过去仅用于吸纳背景孤立伪影的 Register Token 转变为显式引导“压缩-交互-广播”的核心通信瓶颈，彻底打破了以往无监督部件发现依赖繁杂聚类后处理或像素级重建先验的局限。
零额外监督自然涌现：不改变 DINO 自蒸馏损失函数的任何数学形式，仅凭 Transformer 内部连接通路的结构归纳偏置（拓扑瓶颈与头内正交划分），就让纯数据驱动的模型自发学出符合人类直觉的可解释视觉部件。
部件词表的组合泛化：从训练完的寄存器中导出的 Part Dictionary 不仅展现出同类内部一致性与系统分类学邻近性，更能实现对未见组合物体（如“飞马 Pegasus”由马身与鸟翼组成）的零样本组合解构。

局限与展望¶

固有低秩瓶颈对密集预测精度的限制：由于每个注意力头的通信带宽被压缩至 $n \le 16$，表征秩受限。虽然作者在微调阶段引入平行的 $L \to L$ 全连接路径予以弥补，但在冻结特征表征用于超密集像素级任务时仍显容量不足。
无监督边界与人类概念分类的对齐偏差：部件切分纯由统计视觉共现驱动，导致发现的语义边界与人类带语义标签的细粒度定义并非严格等价（例如难以区分紧密粘连但功能独立的附属物）。
未来方向：探索动态分配寄存器容量的自适应瓶颈机制，以及将 RATS 架构先验推广至视频时序部件跟踪与多模态视觉大模型（MLLM）的可解释底层编码。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 架构设计极为精炼，开创性地将寄存器改造为结构化部件通信瓶颈，无监督自发涌现部件令人惊艳。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 5 项权威分割基准、2 项下游密集预测迁移、详尽消融实验与词表语义结构分析。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严密，动机生动，图表清晰直观，分析深入透彻。
价值: ⭐⭐⭐⭐⭐ 为无监督视觉表示学习提供了极具前景的结构先验，对可解释性与组合视觉推理具有重大推动作用。