跳转至

One Model, Two Roles: Emergent Specialization in a Shared Recurrent Transformer

会议: NeurIPS 2026
arXiv: 2605.17811
OpenReview: https://openreview.net/forum?id=5g3g2p3Zwl
代码: https://github.com/juchengshen/air
领域: 其他
关键词: 循环Transformer, 涌现特化, 隐式工作空间, 抽象推理, 状态注意力机制

一句话总结

提出非对称输入循环(AIR)架构,证明参数完全共享的单模型无需物理模块解耦,仅靠输入注入差异或结构化状态标识即可自发涌现出“提议生成”与“局部修正”的双重计算分工。

研究背景与动机

循环推理模型与权重共享架构近年来在符号推理与算法推演等任务中展现出了强大的紧凑表达能力。传统的分层推理架构(如 HRM)往往依靠物理上分离的两个独立网络分别维护高层与低层潜状态,以实现结构化解耦;而全共享模型(如 TRM)虽将参数量压缩至单网络,却普遍缺乏对其内部计算是否真正形成明确分工的深入机制探究。这就带来了一个核心基础问题:参数完全共享的循环 Transformer 究竟是只能以同质化的模式盲目更新潜状态,还是有能力在无参数隔离的前提下自主涌现出异质且稳定的功能角色?

这一问题的核心挑战在于,现有大部分模块化系统均假设计算异质性必须源自显式结构解耦,一旦所有更新步骤绑死在同一个参数矩阵上,模型往往会趋于表征退化或同质化坍缩。若要使共享网络涌现出功能分工,系统必须在循环更新轨迹中提供稳定且可辨识的潜状态身份信号。然而以往研究未能厘清,究竟需要多强的结构诱导才能打破这种对称性瓶颈。

本文的切入角度是构建极简的受控测试基准——非对称输入循环(Asymmetric Input Recurrence, AIR)。在严格统一参数、循环拓扑与单步截断梯度训练流的前提下,仅保留更新规则中“是否注入输入特征”这单一微小差异。核心 idea:通过在循环更新序列中引入非对称输入注入(或独立的结构化层级 token),为共享参数的单一 Transformer 提供清晰的状态身份信号,从而自然诱导出全局承诺解生成(\(\mathbf{z}_H\))与局部不确定性暂存/修正(\(\mathbf{z}_L\))的涌现功能特化。

方法详解

整体框架

AIR 维持两个维度相同的潜状态序列:高层状态 \(\mathbf{z}_H \in \mathbb{R}^{B \times S \times D}\) 与低层状态 \(\mathbf{z}_L \in \mathbb{R}^{B \times S \times D}\)。整个系统仅由一个 4 层的标准 Transformer 块 \(f(\cdot; \mathbf{\theta})\) 构成,所有更新步骤共享同一套权重 \(\mathbf{\theta}\)。输入序列经过编码器得到特征嵌入 \(\tilde{\mathbf{x}}\)。在每个高层更新周期内,模型依次交替执行 \(C_L\) 次 L-update 与 1 次 H-update,并在一个大周期内循环重复 \(C_H\) 次(默认 \(C_L=2, C_H=2\),即单个周期展开为 L, L, H, L, L, H 共 6 个子步)。

在整个前向推演循环中,L-update 会显式注入编码输入 \(\tilde{\mathbf{x}}\),而 H-update 则不注入任何外部输入。虽然两个阶段调用完全相同的网络权重,但这一非对称输入差异赋予了模型区分当前执行角色的明确状态身份。模型最终仅通过接在 \(\mathbf{z}_H\) 上的单个线性分类头进行结果预测与监督,\(\mathbf{z}_L\) 则作为隐式工作暂存区被纯粹的循环动态间接塑形。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入序列 x<br/>(数独/迷宫状态)"] --> Enc["输入编码器<br/>映射为特征 x~"]
    Enc --> Init["潜状态初始化<br/>固定截断正态向量 (zL, zH)"]
    Init --> LStep["非对称注入更新 (L-update)<br/>输入 zL + zH + x~"]
    LStep --> HStep["无输入承诺更新 (H-update)<br/>输入 zH + zL"]
    HStep -->|重复交替子步调度| Loop{"达到截断周期步数?"}
    Loop -->|否: 续传隐状态| LStep
    Loop -->|是: 终止循环| OutHead["共享输出解码头<br/>对 zH 进行线性投影解码"]
    OutHead --> Sol["最终解输出 / 监督损失"]

关键设计

1. 非对称输入注入机制:极简差异打破计算对称性

为了以最小代价打破参数共享带来的同质化坍缩,AIR 在更新规则中显式剥离了显式网络解耦,仅通过指示权重控制外部编码输入 \(\tilde{\mathbf{x}}\) 的注入: $$ \mathbf{z}_L \leftarrow f(\mathbf{z}_L + \mathbf{z}_H + \alpha_L \tilde{\mathbf{x}}; \mathbf{\theta}), \quad \mathbf{z}_H \leftarrow f(\mathbf{z}_H + \mathbf{z}_L + \alpha_H \tilde{\mathbf{x}}; \mathbf{\theta}) $$ 默认设定采用 \((\alpha_L, \alpha_H) = (1, 0)\),即编码输入仅在 L-update 注入,H-update 仅接收双潜状态相加之和。通过这一非对称输入信号,共享网络在执行计算时能够获知当前所处的更新语境:解码中间隐状态显示,\(\mathbf{z}_H\) 表现为一个完全承诺的提议生成状态(在数独中每一步都解码为填满数字的候选解,在迷宫中解码为连贯路径);而 \(\mathbf{z}_L\) 则表现为未完全承诺的草稿暂存区(保留大量代表不确定性的 BLANK 或 PAD 占位符,且位置随推理步骤不断动态迁移)。

2. 结构化层级标识符对比设计:解构状态身份信号的物理载体

针对“功能特化是否必须依赖输入注入”的深层假设,论文进一步探究了对称输入基准(即两个阶段均注入输入,\(\Delta = 0\))下引入额外状态标识信号的有效范式。研究对比了三种不同机制: - 逐元素相加(Addition):为每个 token 向量加上可学习的状态向量 \(\tau_m\); - 独立拼接并剥离(Prepend-strip):在输入序列头部插入专用的可学习 level token,允许其在自注意力机制中参与全交互,并在该子步前向结束时即刻剥离,不带入下一时刻; - 持续拼接(Prepend-no strip):将该 token 持久化保留在潜状态序列中持续演化。

理论与实验揭示,逐元素相加会导致交叉项与 token 本身内容产生位置依赖混叠,使注意力头无法单独读取状态标记;而持久拼接则会导致状态标记累积跨步历史信息而失去瞬时性。唯有 Prepend-strip 提供了结构上正交可分离的 per-sub-step 状态标记,能近乎完全恢复非对称注入带来的性能优势。

3. 跨状态注意力局部性与违规敏感性分工

在微观计算机制层面,AIR 涌现出的双重角色对应着高度正规化的注意力拓扑分离。注意力层探测指标定义为: $$ \Delta_{\mathrm{nbr}}(q) = \mathrm{mass}_L^{\mathrm{nbr}}(q) - \mathrm{mass}H^{\mathrm{nbr}}(q), \quad \Delta) $$ 其中 }}(q) = \mathrm{entropy}(p_q^{(H)}) - \mathrm{entropy}(p_q^{(L)\(\mathcal{N}(q)\) 表示查询单元 \(q\) 的约束邻域。在数独任务中,L-update 在所有网络层中均呈现出比 H-update 更高的局部邻域注意力聚集度(\(\Delta_{\mathrm{nbr}} > 0\))与更低的注意力熵(\(\Delta_{\mathrm{ent}} > 0\))。在深层网络(Layer 2-3),L-update 进一步自发涌现出针对当前规则冲突单元(\(\mathcal{V}\))的定向高权分配(\(\Delta_{\mathrm{viol}} > 0\));而在 H-update 中,注意力分布始终保持全局发散,用于统筹整体棋盘候选方案的生成。

损失函数 / 训练策略

模型采用基于 HRM 的单步截断反向传播机制(One-step Truncated Backpropagation)。在每个训练周期内,前 4 个子步(L, L, H, L)以无梯度模式前向推进,仅在周期的最后一次 L-update 与最后一次 H-update 展开梯度计算并反向传播。

优化器选用 AdamATan2(\(\beta_1=0.9, \beta_2=0.95\)),基础学习率为 \(1 \times 10^{-4}\),配以 2,000 步线性 warmup 及权重衰减 1.0。分类损失使用数值稳定性更优的 stablemax cross-entropy。训练过程中监督信号仅直接施加于 \(\mathbf{z}_H\) 的解码输出,\(\mathbf{z}_L\) 完全不接受显式标签损失。

实验关键数据

主实验

在两个极具代表性的高难度抽象符号推理基准 Sudoku-Extreme(仅 17 个初始提示数的 9×9 数独极难集)与 Maze(30×30 网格硬迷宫寻路)上进行全量测试。表 1 对比了不同输入注入非对称性配置(\((n_L, n_H)\) 表示各阶段注入 \(\tilde{\mathbf{x}}\) 的份数,\(\Delta = |n_L - n_H|\))与双独立网络基线 HRM 的准确率表现:

配置 / 变体 输入注入 \((n_L, n_H)\) 非对称度 \(\Delta\) Sudoku-Extreme (%) Maze-30×30 (%) 架构参数规模
HRM (Wang et al., 2025) 独立双网络基线 — 55.0 74.5 2× 参数量
AIR: Lx_H (默认) (1, 0) 1 60.0 ± 2.0 71.0 ± 6.3 1× 参数量
AIR: L_Hx (0, 1) 1 58.7 ± 3.3 75.3 ± 3.2 1× 参数量
AIR: L_H2x (0, 2) 2 58.6 ± 1.9 75.6 ± 1.9 1× 参数量
AIR: L2x_H (2, 0) 2 59.1 ± 2.4 71.1 ± 6.4 1× 参数量
AIR: Lx_H2x (1, 2) 1 59.6 ± 0.9 70.9 ± 2.4 1× 参数量
AIR: L2x_Hx (2, 1) 1 58.6 ± 2.9 74.5 ± 1.6 1× 参数量
非对称组均值 (Δ>0) 多个非对称组合 ≥1 59.1 73.1 1× 参数量
对称消融: Lx_Hx (1, 1) 0 52.1 ± 1.6 69.4 ± 2.5 1× 参数量
对称消融: L2x_H2x (2, 2) 0 50.9 ± 2.9 70.3 ± 4.2 1× 参数量
对称组均值 (Δ=0) 对称无差异注入 0 51.5 69.9 1× 参数量

消融实验

针对状态标识机制、单状态退化以及算子形式控制的消融实验结果汇总如下:

配置分类 变体名称 核心机制设计 Sudoku-Extreme 准确率 (%) 相对默认增减 (pt)
状态标识控制 Base (L2x_H2x) 对称基线,无标识 50.9 ± 2.9 -9.1
状态标识控制 + Addition 逐元素相加可学习向量 \(\tau_m\) 50.0 ± 1.9 -10.0
状态标识控制 + Prepend (strip) 拼接独立 level token,前向后剥离 57.5 ± 1.3 -2.5
状态标识控制 + Prepend (no strip) 拼接 level token 并持久跨步保留 47.8 ± 1.6 -12.2
状态槽位数量 单状态 \(\mathbf{z}\) (周期注入) 仅 1 个潜状态,周期交替注入 51.1 ± 1.8 -8.9
注入算子形式 线性变换 \(g(\tilde{\mathbf{x}})\) 注入项改为可学习线性映射 59.6 ± 1.3 -0.4
注入算子形式 非线性变换 \(\tanh(g(\tilde{\mathbf{x}}))\) 注入项加激活函数 59.9 ± 1.5 -0.1
注入算子形式 符号反转 \(-\tilde{\mathbf{x}}\) 注入项改为减去编码输入 59.2 ± 1.2 -0.8
注入算子形式 哈达玛积 \((\mathbf{z}_H+\mathbf{z}_L) \odot \tilde{\mathbf{x}}\) 元素级相乘调制 58.3 ± 1.0 -1.7

关键发现

  • 非对称性是性能跃升的关键杠杆:在数独任务中,只要非对称度 \(\Delta > 0\),模型准确率稳定在约 59%~60%,而一旦沦为对称注入(\(\Delta=0\)),准确率剧降至约 51%(下降近 8~9 个百分点)。此外,默认 AIR(仅 50% 参数量)以 60.0% 的成绩显著超越了双网络独立参数的 HRM 基线(55.0%)。
  • 状态相互依赖的动力学反常:冻结实验表明两状态间存在紧密纠缠的动力学闭环。在数独中,冻结 \(\mathbf{z}_H\) 会大幅抑制 \(\mathbf{z}_L\) 的内容变动(步间 token 改变总数从 1235 骤减至 323),而冻结 \(\mathbf{z}_L\) 则反向加剧 \(\mathbf{z}_H\) 的变动(从 275 上升至 551)。在两任务中任意冻结一个状态均导致最终准确率直接归零(0%),证明特化是协同共生而非孤立冗余。
  • 纯局部性 vs 违规导向机制的解耦:注意力分析证明,L-update 的高局部性在所有浅层与深层普遍成立(最深层局部注意力增益高达 47%);但针对约束冲突(violation)的特异性聚焦是在深层(Layer 2-3)才晚期涌现的高阶推理能力。

亮点与洞察

  • 单一模型内生“双重角色”:打破了“功能特化必然要求物理模块分离”的思维定势。仅需引入轻微的非对称信号,共享参数的 Transformer 即可自发形成高阶全局提议者与低阶局部审计者的有机分工。
  • 结构化隔离是标识有效性的本质:Prepend-strip 的成功与 Addition 的失败从代数层面给出了极佳启示——通过自注意力机制独立读取并在步末剥离,避免了标识向量与语义内容的非线性污染,为未来循环架构与统一多任务网络的状态路由设计提供了极简范式。
  • 超强参数利用率的推演框架:AIR 仅用单一 4 层小型 Transformer 便在极度困难的数独与迷宫任务上击败了双网络模型,展示出循环权重共享结合动态隐式草稿在端到端高效推理系统中的巨大潜力。

局限与展望

  • 基准环境的合成局限:所有实证结论均建立在规则完全确定、全状态可观测的合成网格组合任务(Sudoku-Extreme 与 30×30 Maze)上,尚未拓展至部分可观测或大尺度连续自然语言复杂推理场景。
  • 架构深度与周期调度探索不足:论文仅针对固定 4 层、默认 \(C_L=2, C_H=2\) 的固定循环调度进行深入剖析,更深层网络以及自适应动态启停(ACT)对角色分工稳定性的影响仍待后续探索。
  • 注意力权重的机制解释局限:注意力图谱的聚焦虽然直观反映了局部/全局特征,但注意力分布仅代表计算流的一部分,MLP 内部的表征特化机理尚未完全明晰。

相关工作与启发

  • vs HRM (Wang et al., 2025): HRM 采用两套独立的 Transformer 物理网络维护分层状态;AIR 证明完全共享同一套参数权重不仅能涌现出同等甚至更清晰的职能特化,且能在减少一半参数的同时在复杂数独任务上取得更高准确率。
  • vs TRM (Jolicoeur-Martineau, 2025): TRM 虽然使用了共享权重循环网络,但采用端到端全时步 BPTT 训练且未对两状态的职能分工进行因果切分;AIR 引入单步截断梯度并证明了非对称输入注入是驱动异质角色分工的核心归纳偏置。
  • vs COCONUT (Hao et al., 2025): COCONUT 通过连续潜空间隐式思维链进行推理,但依赖专门设计的隐状态传递与逐步替换策略;AIR 展现了在无显式阶段引导下,循环潜状态自发分化为局部草稿(\(\mathbf{z}_L\))与承诺解(\(\mathbf{z}_H\))的涌现特性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙证明并剖析了共享参数循环模型自发涌现功能特化的极简诱导条件,对网络模块化认知有重要理论贡献。
  • 实验充分度: ⭐⭐⭐⭐⭐ 包含跨时步隐状态解码、注入不对称消融、状态标识结构控制、成对冻结因果动力学分析及逐层注意力机制探测,论证链条严丝合缝。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨清晰,问题定义锐利,从宏观现象到微观代数展开层层递进。
  • 价值: ⭐⭐⭐⭐☆ 为轻量级紧凑型循环推理模型的设计以及未来大模型隐式思维工作空间提供了极具启发性的架构参考。