Depth-Recurrent Attention Mixtures: Giving Latent Reasoning the Attention it Deserves¶

会议: NeurIPS 2026
arXiv: 2601.21582
领域: Reasoning
关键词: 深度循环; 隐式推理; 深度注意力; 专家混合; 架构扩展

一句话总结¶

针对深度循环架构中层规模受限与隐状态表征瓶颈两大缺陷，本文提出模块化深度循环注意力混合体（Dreamer），通过将序列、深度与稀疏专家统一建模为三维注意力混合机制，在严格资源对齐下实现了 2–8 倍的数据效率提升。

研究背景与动机¶

大语言模型的复杂推理能力高度依赖于思维链（Chain-of-Thought, CoT）对测试时计算的扩展，但传统的离散自然语言 CoT 会引入冗长且计算昂贵的显式序列生成，限制了表达效率。在连续隐空间内借助参数跨层共享的深度循环（Depth-Recurrence, DR）展开多步隐式推理（Latent Reasoning），被视为消除离散词元冗余、促进知识复用与组合泛化的潜在路径。然而，现有的深度循环方法大多依赖密集网络或外挂非循环的编解码层，始终未能建立起公平严谨的扩展基准。

制约深度循环模型扩展存在两大核心瓶颈：其一是层规模瓶颈（Layer-Size Bottleneck），传统的稠密 DR 模型在每一深度步都需要完整执行全部参数，直接增大单层参数量会导致计算量急剧爆炸；其二是隐状态容量瓶颈（Hidden-Size Bottleneck），当原本沿序列展开的推理过程被大量压缩至深度循环维度时，固定维度的隐状态向量（\(h\)）成为信息累积与交换的信息窄颈，极大地阻碍了更深层次的复杂隐式推导。

针对这双重瓶颈，本文将注意力机制推广至统一的信息访问范式：借鉴序列注意力解决循环神经网络记忆瓶颈的逻辑，引入纵向的深度注意力来打破隐状态容量约束，同时以低秩且带位置编码的稀疏专家注意力解耦参数量与计算量。核心 idea：将序列注意力、纵向深度注意力与稀疏专家注意力统一正交融合为单层单核心的深度循环注意力混合体（Dreamer），彻底打破隐状态与层规模双重瓶颈，实现隐式推理的高效端到端扩展。

方法详解¶

整体框架¶

Dreamer 的核心设计是将所有维度的信息流动统一抽象为注意力运算。模型仅由单一的可复用层堆叠循环构成，输入词元经过深度迭代循环计算，在每个深度步同时沿着序列、深度与专家维度调度知识。序列注意力（SA）维持传统的跨时间步因果交互；深度注意力（DA）沿深度轴检索当前词元在历史深度的激活状态；稀疏专家注意力（EA）则通过路由动态激活特定的 FFN 子专家。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    IN["输入词元嵌入 $x_0$"] --> NORM["残差层归一化 RMSNorm"]
    NORM --> DASA["双路协同注意力<br/>深度注意力 DA + 序列注意力 SA"]
    DASA --> RES1["中间残差融合"]
    RES1 --> EA["稀疏专家注意力 EA<br/>Top-1/Top-$k$ 路由与动态负载均衡"]
    EA --> OUTCHECK{"是否达到最大循环深度 $L$"}
    OUTCHECK -->|否：迭代下一步 $l+1$| NORM
    OUTCHECK -->|是：循环收敛| OUT["输出隐状态及预测分布"]

在推理计算时，序列注意力的键值缓存跨词元持久保留，而深度注意力由于仅在当前词元的深度迭代步之间寻址，其 KV Cache 可以在生成下一个词元时直接覆写丢弃，因此深度注意力带来的显存占用与序列长度完全解耦，额外内存开销恒定且极小。

关键设计¶

1. 深度注意力机制：基于深度反向旋转位置编码纵向打破表征瓶颈 针对深度循环过程中恒定隐状态尺寸 \(h\) 限制深层信息流通的问题，本文构建了沿深度维度检索的深度注意力（Depth Attention, DA）。对于当前深度 \(l\) 的输入特征 \(x_l \in \mathbb{R}^{b \times s \times h}\)，将其视角转换为以批次与序列展平为外部维度的张量 \(\mathbb{R}^{(b \cdot s) \times 1 \times h}\)，从而直接复用底层高度优化的点积自注意力内核。为了给深度循环提供全局时序感知，DA 采用旋转位置编码（RoPE）对深度步编号进行编码，并创新性地将一半的 RoPE 维度进行反向排列，使其能够自底向上编码距离最大深度的前瞻相对距离。在具体实现中，DA 的残差连接采用与 SA 并行计算的结构以优化硬件吞吐量：

\[y_{\operatorname{DA+SA},l} = x_l + \operatorname{DA}(x_{:l}) + \operatorname{SA}(x_l)\]

2. 稀疏专家注意力：低秩解耦计算与容量的动态路由 为了解决稠密层参数扩大带来的层规模瓶颈，本文将前馈网络（FFN）解构为专家注意力（Expert Attention, EA）。每个专家的前向网络对应为值向量变换，而专家键向量由全局可学习参数表达。与传统 MoE 采用高维全连接投影不同，EA 沿用了注意力的点积低秩投影与深度 RoPE 编码，使得专家选择可以与所处深度阶段紧密相关。针对稀疏门控带来的负载不平衡与不可微问题，EA 采用 Sigmoid 激活与无丢弃的偏置调节路由机制：

\[\sigma(x) = \operatorname{norm}\left(\operatorname{TopK}\left(\operatorname{sigmoid}(x), x + b, k\right)\right)\]

其中偏置向量 \(b\) 在训练中根据专家使用计数与使用计数中位数的差值符号进行动态微调：\(b \leftarrow b + \lambda \operatorname{sign}(\operatorname{median}(N) - N)\)。这确保了路由偏置在全局统计下的均值为零，平稳保障了上千个极小专家的利用率平衡。

3. 注意力投影线性 MoE 与可吸收共享专家设计 在深度循环架构下，若在所有深度步强制复用单一投影权重矩阵，会导致注意力头的表达能力出现明显退化。为了在维持单一循环层的前提下增强深度适应性，Dreamer 将每个融合的 Query/Key/Value 投影以及多头注意力聚合变换也改造为轻量级线性 MoE。两个线性 MoE 共享路由注意力得分以消除额外路由损耗。为了解决 Top-1 路由导致的备选专家梯度截断问题，模型引入了一个恒常激活且利用截断梯度缩放的共享基底专家 \(W_{\text{shared}}\)：

\[\sigma x W + \max(\widehat{\sigma}) x W_{\text{shared}} = \sigma x (W + W_{\text{shared}})\]

得益于线性专家的线性叠加性质，共享专家的权重矩阵可以在训练完成后离线重参数化直接加和到各个路由专家权重中，在推理部署阶段实现了零额外计算与显存开销。

损失函数 / 训练策略¶

模型采用标准自回归因果语言建模交叉熵损失进行训练。为消除显式思维链带来的低信噪比文本噪音并倒逼隐式推理能力，训练数据集过滤了冗长的显式 CoT 轨迹，仅保留包含精简解答指导的自然数学样本。优化器采用 AdamW，批大小为 \(2^{20}\) tokens，最大学习率设为 \(4 \times 10^{-4}\)，在经过 500 步线性预热后保持常数衰减；网络残差流引入 RMSNorm 归一化以抑制多步循环积累的残差发散。在模型对比前，采用双变量坐标下降算法对 EA 中间层维度与专家数量进行精准对齐，确保所有同深度基线之间的 FLOPs、参数量与显存占用完全一致。

实验关键数据¶

主实验¶

实验基于 100B tokens 的数学与常识数据集展开，对比经典分层 Transformer（LA）、标准深度循环（DR）以及结合深度注意力的循环模型（DR+DA）。测试覆盖 GSM8K、MATH、MMLU（数学子集）与 MathQA 等推理评测集。

模型配置	循环深度	参数量 (B)	FLOPs (B/tok)	GSM8K Acc (0-shot)	MATH Acc (0-shot)	MMLU-Math Acc	平均准确率	数据效率提升 (DE)
LA (分层基准)	16	1.1708	0.9389	43.8%	40.9%	26.6%	32.3%	1.0×
DR (深度循环)	16	1.1704	0.9389	47.6%	47.8%	44.5%	45.7%	4.0×
DR+DA (本文)	16	1.1708	0.9413	51.7%	50.7%	38.5%	43.4%	2.9×
LA (分层基准)	32	2.0790	1.6150	49.7%	48.2%	37.2%	39.5%	1.0×
DR (深度循环)	32	2.0794	1.6199	51.5%	51.3%	45.3%	47.6%	2.6×
DR+DA (本文)	32	2.0788	1.6130	56.3%	54.5%	50.0%	51.4%	3.5×

消融实验¶

针对注意力混合机制并行度以及各模块开销与性能的消融分析如下：

配置架构	深度步数	相对单步吞吐提升	语言建模困惑度 (PPL)	GSM8K (0-shot)	特性说明
串行架构 (DA \(\to\) SA \(\to\) EA)	32	0.0% (基准)	5.88	56.5%	精度最高但单步时延最长
部分并行架构 (DA+SA 并行，EA 串行)	32	+15.2%	5.90	56.3%	本文默认方案，兼顾计算吞吐与推理精度
全并行架构 (DA + SA + EA 并行)	32	+30.8%	6.24	51.2%	计算最快，但在小尺寸参数下精度有明显损失
DR+DA 去除残差 RMSNorm	32	-	发散 (NaN)	-	深层循环时残差幅度剧烈漂移致使训练失败
DR 移除注意力投影线性 MoE	32	+2.1%	6.45	45.1%	跨深度强制复用权重限制了注意力的阶段性表征

关键发现¶

打破深层表征瓶颈：在 16 层浅层场景下，DR 与 DR+DA 的平均表现互有胜负；但当扩展至 32 层深层时，DR+DA 相对纯 DR 取得显著优势（平均分从 47.6% 跃升至 51.4%），验证了深度注意力在更深隐式推理链条中突破容量瓶颈的核心价值。
超越翻倍规模基线：深度为 16 的 DR+DA 模型在 GSM8K（51.7% vs 49.7%）和 MATH（50.7% vs 48.2%）等全部核心推理任务上全面超越了参数量与 FLOPs 接近其两倍的 32 层分层基准（LA-32）。
专家调度多样性跃迁：洛伦兹曲线分析表明，DR 架构在各个深度步调用的独特专家数量达到了传统分层 MoE 的 2 到 11 倍，同时全局基尼系数低至 0.075，兼具超高知识调度灵活性与极佳的负载平衡性。

亮点与洞察¶

注意力机制的全维度归一：将序列上下文、深层循环轨迹与参数空间中的专家动态调度全部抽象为 Query-Key-Value 交互，构建了高度优雅且对称的通用注意力混合统一框架。
零推理代价的共享专家融合：针对 Top-1 稀疏注意力线性专家设计了与最大得分同比例缩放的共享基底专家，在训练期平稳提供备选专家梯度，在部署期通过算子重参数化直接吸收合并，实现零附加时延。
深度注意力随用随抛的高效 KV 缓存：DA 仅需维护当前生成词元内部的历史深度键值对，词元完成计算后即可原地覆写，彻底摆脱了序列级 KV 缓存随长度二次暴涨的显存痛点。

局限与展望¶

作者承认的局限：目前验证的最大深度仅为 32 步，模型在超长隐式推理步数下的动态循环泛化能力及自适应早停停止机制尚未完全建立严密的扩展定律（Scaling Laws）。
实验设计与场景局限：当前的实验重点聚焦于自然语言数学推理任务，尚未在通用代码生成、大规模多轮对话及长文本多模态理解等非结构化复杂任务上进行泛化验证。
未来改进思路：探索结合局部滑动窗口与扩张稀疏深度注意力，进一步压缩深层步数的显存访问带宽；将深度循环步数与自适应测试时计算分配相联合，实现根据问题动态调节思考深度的弹性推理。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [将序列、深度与专家统一抽象为正交注意力混合体，极具理论前瞻性与架构美感]
实验充分度: ⭐⭐⭐⭐⭐ [执行了极为严格的双变量坐标下降 FLOPs/参数/显存完全对齐消融实验]
写作质量: ⭐⭐⭐⭐⭐ [概念阐述清晰透彻，可视化分析深层专家与注意力行为丰富严谨]
价值: ⭐⭐⭐⭐⭐ [为替代显式 CoT、探索端到端常数显存开销的纯隐空间测试时扩展提供了极具潜力的架构范式]