仅改动0.016%权重，破解小模型推理循环

一、问题背景：推理循环为何拖垮小模型

在一项针对小参数量语言模型的数学推理评测中，96 个测试样本里有 85 个因触及 32768 token 的输出上限而未能正常结束推理。它们并非“算得慢”，而是陷入了所谓的 Reasoning Loop（推理循环）：模型反复生成高度相似的片段，例如不断自我诘问“20 是总数吗？不对，它是剩余时间”，在同一个步骤上原地打转，直到上下文窗口被彻底耗尽。

这一现象可以从三个层面理解：

1. 自回归解码的误差累积：模型每一步都以自身历史输出为条件，一旦某个片段被重复，上下文会反过来强化“继续重复”的概率，形成正反馈回路；
2. 容量与任务复杂度的失配：小模型难以维持全局求解计划，当对当前状态不确定时，高概率的续写往往是复述或反问，而非真正推进；
3. 解码策略的放大效应：低温度或贪心解码会进一步压缩探索空间，使模型更容易锁定在局部重复模式上。

后果是三重的：推理延迟与 token 成本飙升、大量样本根本得不到答案、上下文预算被循环内容挤占而无法用于真正的求解。常见的黑盒手段——改提示词、加大重复惩罚、扩充微调数据——大多只治标不治本。本文介绍的方法则反其道而行：深入模型内部，找到循环的表征根源，再做最小化修复。

二、诊断思路：把行为问题转化为表征问题

该方法的核心假设是：循环不是随机噪声，而是模型内部表征沿某个特定方向被持续放大的结果。如果这个方向真实存在，就应该能在激活空间中被观测、被定位、被干预。

第一步是构造对照实验。研究者为同一道数学题准备两种推理文本：

• Productive Reasoning（正常推理）：如算出 10+8+20=38 后继续向目标推进；
• Repetitive Reasoning（循环推理）：在某个步骤反复纠结，不断重复相似的自问自答。

两种文本来自同一道题、同一个模型，唯一差别是“是否陷入循环”，从而把题目难度、模型能力等混杂因素严格控制住。对照样本可以通过截取真实循环轨迹后拼接正常续写、或人工改写等方式构造，实验中共使用了 4 组这样的配对。

第二步是采集内部状态。利用 Teacher Forcing，把两种文本分别整体输入原始模型，仅需一次前向传播，即可读取每一层、每个 token 位置的 Hidden States（残差流）与 MLP Activations。相比逐 token 真实解码，这种方式成本极低，且保证了两种条件下激活的严格可比性。

三、提取 Repetition Direction：循环的“几何指纹”

对第 l 层残差流，取每组样本最后 36 个 token 的平均隐藏状态，分别记为 h̄_normal^(l) 与 h̄_loop^(l)，对照方向定义为四组差异的均值：

d^(l) = (1/4) · Σ(i=1..4) ( h̄_loop,i^(l) − h̄_normal,i^(l) )

再作归一化：

u^(l) = d^(l) / ‖d^(l)‖₂

三个设计细节值得展开：

1. 为什么取最后 36 个 token：循环的判别特征集中在“最近发生了什么”。更早的位置主要是题目与已确认的步骤，对是否循环不敏感；窗口过长会稀释信号，过短则统计方差过大。
2. 为什么对 4 组样本取平均：单组差异同时包含“循环信号”与“题目特有噪声”，多组平均能让题目特异性分量相互抵消，留下共同方向。这隐含一个假设——不同题目的循环共享同一表征方向，实验结果支持该假设在小模型上近似成立。
3. 为什么要归一化：不同层的激活范数差异悬殊，归一化后才能跨层比较方向的相对强度，也便于后续做投影分析。

示意代码如下：

# 采集：对每组对照文本做一次 teacher-forcing 前向
out_n = model.forward(texts_productive,  output_hidden_states=True)
out_l = model.forward(texts_repetitive,  output_hidden_states=True)
d = []
for l in range(num_layers):
    diff = (out_l.hidden[l][:, -36:, :].mean(1)     # (4, d_model)
          - out_n.hidden[l][:, -36:, :].mean(1))
    d.append(diff.mean(dim=0))                      # 跨 4 组平均
u = [d[l] / d[l].norm() for l in range(num_layers)] # 每层单位方向

四、层级定位：找出“谁在放大循环”

得到每层的 d^(l) 后，考察相邻层差异：

Δd^(l) = d^(l+1) − d^(l)

若某层之后 ‖Δd^(l)‖ 显著增大，说明正常与循环两种激活的差异在这一层被明显放大——即该层在主动“推动”重复行为。实践上通常先把所有层的 ‖Δd^(l)‖ 绘成曲线，直观找出少数几个尖峰层作为候选，而不是逐层盲目排查。

五、子层归因：Attention 还是 MLP？

Transformer 每层都通过残差连接向残差流写入两部分信息：Attention 输出与 MLP 输出。因此层间差异可分解为：

Δd^(l) ≈ Δ_attention^(l) + Δ_MLP^(l)

其中 MLP 的贡献可以直接计算：

Δ_MLP^(l) = W^(l) · ( ā_loop^(l) − ā_normal^(l) )

这里 ā 表示最后 36 个 token 的平均 MLP 激活（激活函数之后、降投影之前），W^(l) 为该层 MLP 的降投影权重矩阵；Attention 的贡献则可用总差异减去 MLP 部分间接得到。

在原案例中，差异增长主要来自少数层的 MLP 通道：这些层的 MLP 在循环语境下持续向残差流写入“沿重复方向”的成分，相当于不断给模型下达“继续重复”的内部指令。

六、最小干预：0.016% 的权重改动如何生效

定位到关键层与关键子层后，有两条常见修复路径：

路径 A：推理时方向消融（不改权重）。在前向传播中对目标层的残差流或 MLP 输出减去沿 u 的投影分量：

def ablate(x, u, alpha=1.0):
    coeff = (x @ u).unsqueeze(-1)   # 沿重复方向的投影强度
    return x - alpha * coeff * u

零训练、可随时回滚，但推理时引入少量计算，且只对挂载了该逻辑的服务生效。

路径 B：权重空间低秩编辑。将“压低沿 u 方向的写入门限”转化为一个秩 1（或低秩）的权重更新：

W' = W − α · u vᵀ

仅修改个别层的极小部分参数，占全模型约 0.016%。之所以如此小的改动就足够，是因为循环行为高度依赖少数几条高度集中的表征方向，而非弥散于全参数空间的能力；精准削除这一方向，几乎不触碰模型的其他能力。

在原案例的报告中，完成这一改动后，该小模型的数学评测成绩提升至原来的数倍（报告为 6 倍）。提升的主要来源并非模型“变聪明”，而是大量原本耗死在循环里的样本得以正常走完推理链条——上下文预算被归还给了求解本身。

七、验证与评估：确认修复有效且无害

建议从四个维度做前后对比：

1. 循环率：新样本中触及 token 上限而未正常结束的比例，应显著下降；
2. 任务准确率：在留出的数学测试集上验证收益；
3. token 效率：平均生成 token 数应回落至健康区间；
4. 通用能力回归：在摘要、翻译、对话等非数学任务上回归测试，确认无灾难性遗忘或风格漂移。

此外建议做消融实验：只改 Attention、只改 MLP、只改单层与多层各试一遍，确认收益确实来自被定位到的组件，而非巧合。

八、工程实践与局限性

• 白盒前提：方法要求读取内部激活并修改权重，仅适用于自研或开源权重模型，闭源 API 场景不适用；
• 方向泛化性：Repetition Direction 可能与任务类型、模型规模相关，迁移到新任务前应重新采集对照样本并验证方向一致性；
• 过抑制风险：数学推理中合理的复述（如重读题目条件）与病态循环只有程度差异，系数 α 需小步搜索，避免把正常行为一并削掉；
• 监控闭环：无论自建推理集群还是使用云上托管服务（例如百度智能云提供的模型推理类产品），都建议把“循环率/上限截断率”纳入线上监控指标，及时发现效果回归。

局限方面，该方法是“对症”而非“根治”：它修复的是小模型在特定能力上的系统性缺陷，不能替代更大规模的预训练或更高质量的后训练；方向假设在更大模型上是否依然成立，也需逐案验证。

九、结语

这套方法的价值不止于“0.016% 权重换来数倍提升”这个结果，更在于其工作范式：先在激活空间中为错误行为找到几何坐标，再顺着层级定位责任模块，最后以最小代价实施外科手术式的修复。它把模型调试从依赖直觉的“玄学调参”，推进到可观测、可归因、可干预的工程学科。对于在资源受限场景下部署小模型的团队而言，这是一条值得纳入工具箱的实践路径。