CosmicFish-HRM：通过紧凑语言模型中的层级循环机制实现自适应推理

# CosmicFish-HRM：通过分层循环机制实现紧凑语言模型的自适应推理
来源：https://arxiv.org/html/2605.28919

###### 摘要

大型语言模型已展现出强大的推理能力，但这通常以巨大的参数数量和昂贵的推理为代价。在本工作中，我们探索了一个不同的方向：紧凑语言模型中的自适应推理深度。

我们提出了 CosmicFish-HRM，这是一个围绕分层推理模块（HRM）构建的紧凑语言模型，该模块在推理过程中动态分配计算资源。模型并非对每个输入应用固定的计算量，而是通过高层和低层推理循环进行迭代，并学习根据输入复杂度决定何时停止。

CosmicFish-HRM 将这一自适应推理核心与现代 Transformer 组件（包括分组查询注意力、RoPE 和 SwiGLU 激活函数）相结合。虽然额外的推理基础设施在较小规模下会引入开销，但我们假设随着模型规模增长，HRM 核心的相对成本降低，这种权衡将变得越来越有利。

我们的结果表明，模型学会了非均匀的推理行为，根据任务和输入分配不同数量的推理步骤。这些发现表明，自适应推理深度可能为仅依赖参数规模来获得推理能力提供了一种有前景的替代方案。

## 1 引言

现代语言建模的主导趋势一直是规模。从 GPT-4 [15] 到 LLaMA [25]，进展基本上遵循一种熟悉的模式：更大的参数数量、更大的数据集以及日益昂贵的训练过程。结果令人印象深刻，尤其是在推理和指令遵循方面，尽管这一进展也将最先进的系统推离了在边缘硬件、移动设备和其他资源受限环境中的实际部署。运行现代语言模型通常需要大量内存、专门的加速器和推理基础设施，而这些在数据中心之外几乎不存在。

大多数现有的效率方法都集中在压缩上。量化、剪枝和知识蒸馏 [10] 减小了已有大型模型的尺寸，但底层的计算结构通常保持不变。标准 Transformer 仍然对每个输入应用相同的计算深度，无论其难度如何。一个简短的事实性补全问题和多步推理问题本质上经过相同的流程，具有相同数量的层和注意力操作。然而在实践中，并非每个提示都需要相同的推理量。

我们持不同观点。有些任务几乎是反射性的。其他任务则需要迭代、修正或更长的内部计算链。一个能够在必要时才分配额外推理步骤的紧凑模型，可能与简单地无限增加参数数量提供不同的权衡。本工作的核心思想很直接：推理能力可能不仅取决于模型规模，还取决于推理过程中计算使用的灵活性。

为了探索这一想法，我们提出了 CosmicFish-HRM，这是一个紧凑的 82.77M 参数语言模型，围绕分层推理模块（HRM）构建。模型并非作为一个纯粹固定深度的 Transformer 运行，而是通过迭代循环在高（HH）和低（LL）推理状态之间交替。一个学习到的停止头决定在继续生成之前是否需要额外的推理步骤。困难的输入可能触发更深的推理轨迹，而简单的提示可以提前停止并避免不必要的计算。

该架构将这一自适应推理核心与现代 Transformer 组件相结合，包括旋转位置嵌入（RoPE）[24]、分组查询注意力（GQA）[1] 和 SwiGLU 激活函数 [23]。CosmicFish-HRM 在 CosmicSet 上进行了训练，这是一个精心策划的 10B token 数据集，涵盖网络文本、维基百科、代码、数学和研究论文。

在紧凑规模下，引入循环推理基础设施会产生一种权衡。模型的部分容量被分配给自适应推理行为，而不是纯粹的原始语言建模能力。这可能会限制与传统相同规模的 Transformer 相比的基准性能。然而，我们假设随着模型规模的增大，推理基础设施的相对开销会减少，而自适应计算行为保持不变，这种权衡将变得越来越有利。

本文做出了以下贡献：

- • 我们提出了 CosmicFish-HRM，一种紧凑的语言模型架构，它将分层推理模块集成到 Transformer 主干中，以实现推理过程中的自适应推理深度。
- • 我们引入了一种学习到的自适应停止机制，允许模型跨输入动态变化推理深度，而不是统一应用固定计算。
- • 我们分析了推理步骤分配与输入复杂度之间的关系，表明模型在不同任务上展现出非均匀的推理行为。
- • 我们讨论了在紧凑规模下引入自适应推理基础设施的权衡，并论证了随着模型规模增大，这些成本可能会更有利地被摊销。

本文的其余部分组织如下。第 2 节回顾了自适应计算和高效语言模型的相关工作。第 3 节详细描述了 CosmicFish-HRM 架构。第 4 节展示了实验结果和评估。第 5 节分析了模型的自适应推理行为，第 6 节讨论了局限性和未来方向。

## 2 相关工作

### 2.1 神经网络中的自适应计算

神经网络应根据输入难度调整其计算工作的想法并不新鲜。Graves [8] 引入了自适应计算时间 (ACT)，使循环网络能够决定在产生输出之前使用多少计算步骤。核心直觉简单但重要：不同的输入可能需要不同量的处理。然而，ACT 是在循环网络的背景下开发的，不能自然地扩展到基于 Transformer 的现代语言模型。

后来的工作探索了相关形式的自适应计算。PonderNet [2] 引入了一种概率性的停止机制，通过额外的正则化目标进行训练，鼓励模型仅在必要时“思考”。Universal Transformers [7] 通过重复应用同一层进行多次迭代，将循环深度引入 Transformer 架构，但没有明确的学到的停止策略。Confident Adaptive Language Modeling [21] 探索了 Transformer 解码器中的早期退出，允许一旦预测置信度足够高时，推理在中间层停止。

Mixture of Depths [19] 从不同角度处理了该问题，通过动态地将 tokens 路由到 Transformer 层的子集中。这种方法并非迭代推理，而是专注于在 token 和层之间不均匀地分配计算，以减少总体推理成本。

CosmicFish-HRM 的动机略有不同。我们研究的不是早期退出或 token 路由本身，而是紧凑语言模型是否可以从生成进行之前在多个抽象级别上运行的迭代推理循环中受益。

### 2.2 分层推理模型

与我们工作联系最直接的是 Wang 等人 [26] 引入的分层推理模型 (HRM)。HRM 提出了一种由两个交互模块组成的循环推理架构：一个负责抽象规划的较慢的高层推理系统，和一个负责详细计算的较快的低层推理系统。尽管参数数量相对较少，原始的 HRM 在结构化推理任务上取得了强劲性能，包括数独求解、迷宫导航和 ARC 风格的推理基准 [5]。

也就是说，原始的 HRM 并非被设计为语言模型。它主要作为一个特定任务的推理引擎运行，在相对较小的面向谜题的数据集上进行训练，不具备开放式文本生成或自回归语言建模能力。

CosmicFish-HRM 在保留 HRM 核心直觉的同时，将其迁移到不同的场景。我们不将分层推理视为独立的推理引擎，而是将其集成到解码器专用的语言模型架构中，该架构在多样化的自然语言数据上进行训练。我们的目标不仅仅是复现 HRM 风格的推理行为，而是探索迭代分层计算是否能在紧凑的自回归语言模型内部提供有用的自适应推理动态。

### 2.3 小型高效语言模型

另一条独立的工作线专注于通过训练效率和数据质量来提高小型语言模型的能力。TinyLlama [30] 证明了在大型语料库上训练的较小模型即使在参数数量减少的情况下也能保持竞争力。Phi 系列 [9] 进一步论证了精心策划的高质量数据集可以显著提升紧凑型模型的性能。

这些方法主要通过数据和训练策略来提高效率，而不是改变推理时的计算本身。底层架构通常仍然是传统的固定深度 Transformer。

CosmicFish-HRM 探索了一个互补的方向。该模型不是仅仅依赖改进的数据效率，而是尝试在推理期间动态变化推理深度。在紧凑规模下，这引入了一个架构性权衡：模型容量的一部分被用于自适应推理基础设施，而不是纯粹的原始语言建模能力。我们将这项工作视为对这种权衡的探索，而不是声称自适应推理普遍提高小模型基准性能。

### 2.4 自适应停止与学习到的决策机制

神经网络先前已经结合了学习到的决策机制用于路由、停止和自适应执行 [14]。受这些想法的启发，CosmicFish-HRM 使用一个轻量级的学习到的停止头来确定每一步是否需要额外的推理迭代。

我们的公式故意保持简单。停止机制与语言建模目标和步数惩罚联合训练，使模型能够学习可变深度的推理行为，而无需手动设计的停止标准。我们不将系统视为完整的强化学习设置，而是将停止行为主要视为一种学习到的自适应计算机制，直接集成到前向传播中。

## 3 架构

CosmicFish-HRM 是一个解码器专用的语言模型，围绕三个主要阶段组织：用于上下文表示学习的输入 Transformer 堆栈、用于自适应迭代计算的分层推理模块 (HRM) 以及用于 token 预测的输出 Transformer 堆栈。整体设计将语言表示学习与自适应推理过程本身分离开来。

见图注图 1：CosmicFish-HRM 架构概览。输入 tokens 被嵌入并通过 Transformer 块处理，经过 HRM 推理核心进行可变次数的推理迭代，然后通过输出 Transformer 层解码为下一个 token 预测。### 3.1 总体设计

该架构遵循以下流程：

x→Embedh0→Input Blocksh→HRM Corez→Output Blocksz^→LM Headlogits (1)

给定一个输入 token 序列 x = (x1, x2, ..., xT)，tokens 首先通过一个学习到的嵌入矩阵 E ∈ ℝ^(V×d) 映射为稠密嵌入，其中 V=50304 是词汇表大小，d=448 是嵌入维度。然后表示经过六个输入 Transformer 层、HRM 推理核心和六个输出 Transformer 层处理，最后投影为词汇 logits。遵循标准实践，嵌入矩阵与语言建模头共享权重 [17]。

关键架构超参数总结于表 1。

表 1：CosmicFish-HRM 架构超参数。
| 超参数 | 值 |
| :--- | :--- |
| 总参数 | 82.77M |
| 嵌入维度 d | 448 |
| 词汇表大小 V | 50,304 |
| 最大序列长度 | 512 |
| 输入 Transformer 层数 | 6 |
| 输出 Transformer 层数 | 6 |
| 注意力头数 | 8 |
| KV 头数 (GQA) | 4 |
| HRM 高层层数 n_H | 4 |
| HRM 低层层数 n_L | 4 |
| 最大推理步数 | 16 |
| 探索概率 | 0.1 |
| Dropout | 0.1 |

### 3.2 输入和输出 Transformer 块

两个 Transformer 堆栈均遵循标准的预归一化 Transformer 设计 [27]。每个块通过残差连接应用注意力和前馈更新：

h′ = h + Attn(RMSNorm(h)) (2)
h′′ = h′ + MLP(RMSNorm(h′)) (3)

使用 RMSNorm [29] 替代 LayerNorm 以提高计算效率：

RMSNorm(x) = (x / RMS(x)) ⋅ γ (4)
RMS(x) = sqrt((1/d) * Σ_i x_i² + ε) (5)

#### 分组查询注意力。

模型使用分组查询注意力 (GQA) [1]，具有 n_q=8 个查询头和 n_kv=4 个键值头。每个查询头的维度为 d_h = d/n_q = 56。在查询组之间共享键值头减少了 KV 缓存内存需求，同时保留了标准多头注意力的绝大部分表示能力。

旋转位置嵌入 (RoPE) [24] 在注意力计算之前应用于查询和键：

q_m = R_m q_m, k_n = R_n k_n (6)

其中 R_m 表示位置 m 处的旋转矩阵。RoPE 使得相对位置信息能够直接通过注意力机制产生，同时保持位置的绝对编码。