Dense Supervision Is Not Enough: The Readout Blind Spot in Looped Language Models

背景

循环语言模型（Looped Language Models）是一种将隐藏状态转化为运行时状态的架构。在这种模型中，每个时间步的隐藏状态不仅用于当前的预测解码，还会被反馈到后续的计算中，形成递归或循环结构。这种设计旨在通过多次迭代深化模型对序列信息的处理，从而在有限的计算深度下获得更丰富的表征。

然而，这种架构带来了一个基本的监督问题：交叉熵损失（Cross-Entropy Loss）究竟控制了哪些状态变量？在标准的监督学习中，我们假设损失函数能够直接优化模型内部的所有关键变量。但在循环结构中，情况变得复杂。如果损失函数无法完全控制循环内部的关键动态变量（如隐藏状态的尺度），模型可能会出现训练不稳定或性能瓶颈。

核心内容

本文通过理论分析和实验验证，揭示了循环语言模型中存在的“读取盲区”（Readout Blind Spot）。研究指出，密集的每循环交叉熵监督（Dense per-loop cross-entropy）仅能控制由读取头（Readout）暴露出来的变量，而无法直接控制循环转换过程中活跃的所有变量。

隐藏状态尺度（Hidden-State Scale）的失效模式

研究的核心发现是，隐藏状态的尺度（Scale）是一个关键的失效模式。在许多现代循环架构中，读取头通常使用尺度不变的正规范化方法，如 RMSNorm 或 LayerNorm。这些规范化操作会隐藏径向尺度信息，使得直接的交叉熵损失无法感知隐藏状态的绝对大小。

与此同时，预归一化残差循环（Pre-norm residual recurrence）结构会继续携带并更新这些被隐藏的尺度信息。这意味着，尽管每循环的损失函数在优化预测精度，但它实际上是在“盲目”地训练出口（Exits），而无法控制循环内部的尺度增长。

实验证据

在包含 44M 和 129M 参数的循环 Transformer 模型中（且循环间没有归一化层），即使通过 RMSNorm 读取头进行每循环交叉熵监督，最终的隐藏状态范数（Norms）仍然会激增到数千甚至数万。这表明，密集的每循环监督虽然能让早期的退出机制（Early exits）变得可用，但它并没有控制住循环内部的尺度爆炸问题。

解决方案与设计准则

为了解决这一问题，研究提出了两种互补的架构修复方案：

1. 尺度可见的读取头（Scale-visible readouts）：使用不隐藏尺度的读取方式，使损失函数能够直接感知并优化隐藏状态的尺度。
2. 显式的范数惩罚（Explicit norm penalties）：在损失函数中加入对隐藏状态范数的惩罚项，直接约束尺度增长。

此外，研究还指出，“移除尺度的循环”（Scale-removing recurrence）是另一种互补的架构修复手段。

由此得出的设计准则非常简单：密集监督训练出口；而循环尺度的控制，要么让尺度对损失函数可见，要么将其从循环中移除。

遵循这一准则的尺度控制变体，在变量深度基准测试中，在匹配推理深度操作点时，实现了更低的困惑度（Perplexity）。

关键要点

• 监督局限性：在循环语言模型中，密集的每循环交叉熵损失只能控制由读取头（Readout）暴露的变量，无法直接控制循环转换中所有活跃的隐藏状态变量。
• 尺度盲区：使用 RMSNorm 或 LayerNorm 等尺度不变规范化技术的读取头，会隐藏隐藏状态的径向尺度，导致交叉熵损失无法感知和控制该尺度。
• 尺度爆炸：在预归一化残差循环结构中，即使损失函数在优化预测，隐藏状态的范数仍可能失控增长至数千或数万，尤其是在缺乏循环间归一化的情况下。
• 双重解决策略：
  1. 使尺度对损失函数可见（通过尺度可见的读取头或显式范数惩罚）。
  2. 从循环结构中移除尺度（通过特定的循环架构设计）。
• 性能提升：遵循“尺度控制”准则设计的模型，在相同的推理深度下，能够显著降低困惑度，证明尺度控制对于循环语言模型的性能至关重要。

意义与影响

这项研究对循环语言模型的设计和优化具有重要的指导意义。它挑战了“密集监督足以优化所有内部状态”的传统假设，指出了在循环架构中，隐藏状态尺度的控制是一个独立且关键的问题。

对于模型架构师而言，这意味着在设计循环 Transformer 或其他递归神经网络时，不能仅依赖标准的交叉熵损失来保证训练稳定性。必须显式地考虑隐藏状态的尺度管理，要么通过修改读取头使其对尺度敏感，要么在架构层面引入尺度抑制机制。

此外，这一发现也为理解早期退出（Early Exit）机制与循环内部动态之间的关系提供了新的视角。它表明，优化早期退出并不等同于优化了整个循环的动态过程，特别是在尺度管理方面。未来的研究可以在此基础上，探索更高效的尺度控制方法，以进一步提升循环语言模型在长序列处理和复杂推理任务中的性能。