[ICML 2026] SOL：让大模型把算力花在关键Token上：自优化语言模型

大语言模型在线推理的效率优化，过去主要关注如何降低每个解码步骤的计算成本，比如通过量化、剪枝、稀疏注意力或者低秩分解。但这些方法通常对所有生成的token施加一个相同的计算预算。一个显而易见但通常被忽略的问题是：不同token的生成难度差异很大。有些token在上下文中具有高可预测性，而另一些则依赖长程依赖或精确的中间计算。一种统一的、静态的压缩策略，往往在简单token上浪费算力，却在困难token上算力不足。

针对这一痛点，来自康奈尔大学的Yash Akhauri和Mohamed S. Abdelfattah提出了一项新的工作：Self-Optimizing Language Models，被ICML 2026接收。这篇论文的核心创新在于，它不是预设一个静态的压缩率，而是训练一个轻量级的策略网络，让它在每一步解码时，根据语言模型内部的状态，动态决定当前token要使用多少计算资源。这种动态预算分配不仅覆盖一种效率机制，而是统一控制了 token 级的注意力稀疏性、结构化 MLP 激活剪枝，以及激活量化位宽。这篇论文值得所有从事 LLM 推理效率优化的工程师与研究者仔细阅读。它揭示了一个以往被低估的优化维度：不需要改动基座模型的预训练权重，也不需要设计复杂的启发式规则，只需要引入一个很小的策略网络，就可以在保持甚至提升语言模型质量的同时，显著改善计算效率的利用。论文中展示了在 MMLU 评估上，SOL 相比均匀预算分配策略最高提升 7.3% 的准确率，这个结果在匹配甚至更低的计算预算下实现，令人印象深刻。

摘要

现有的大语言模型推理效率研究主要致力于降低每个解码步骤的成本（例如通过量化、剪枝或稀疏注意力），但这些方法通常对每一个生成的token都应用统一的计算预算。实践中，token的难度差异很大，因此静态的压缩往往在简单的步骤上过度计算，在困难的步骤上计算不足。本文研究自回归解码中的动态预算分配：在一个模型内部学习每个token应该花费多少计算量。

Self-Optimizing Language Models (SOL) 将一个冻结的 LLM 与一个轻量的策略网络配对。该策略网络读取 LLM 的隐藏状态，并在每个解码步骤选择一个离散的效率动作。这些动作可以联合控制：(i) token 级的注意力稀疏性，(ii) MLP 中的结构化激活剪枝，以及 (iii) 激活量化位宽，同时保持基座模型的权重不变。策略使用组相对策略优化（Group-Relative Policy Optimization, GRPO）在教师强迫（teacher-forced）片段上训练：token 序列保持不变，同时采样多个反事实的计算调度（即仅针对相同 token 路径变化效率动作的调度），并在相同的监督下比较它们的似然。奖励函数对语言模型质量与软惩罚进行权衡，软惩罚鼓励片段平均预算使用量匹配请求的目标预算。在各种模型变体和计算体制下，SOL 在匹配预算时优于静态分配和强随机调度搜索，为推理效率优化提供了一个补充的新轴。SOL 在所有实验中发现了更优的质量-效率帕累托前沿，并在 MMLU 上相比均匀预算分配策略最高提升了 7.3% 的准确率。

引言：论文要解决什么问题

大规模语言模型的部署带来了巨大的推理成本，由此推动了许多提升单步解码效率的研究工作。典型的例子包括：量化（减少权重或激活值的位宽）、剪枝（移除不重要的权重或结构）、稀疏注意力（只让每个 token 关注上下文的一小部分）等。尽管这些技术各有成效，但它们有一个共同的假设：对于一个给定的语言模型，每个生成的 token 应该分配相同或非常相似的计算预算。引言明确指出，这种“静态分配”与实际解码过程是矛盾的。在实际生成中，不同 token 的难度天差地别：例如，“the”、“and”这类高频词汇的生成，与一个需要精确推理的推理链 token 相比，计算需求完全不同。如果采用统一的压缩策略，要么为了照顾简单 token 而过度压缩，导致困难 token 质量下降；要么为了照顾困难 token 而采用宽松的压缩率，在简单 token 上浪费大量计算。这种“一刀切”的方法使得 LLM 在实际推理中往往要么过计算，要么欠计算。论文进一步指出，逐 token 的效率决策在时间上是相互影响的。因为当前 token 可以被未来的 token 所关注。如果在时间步 t 过度压缩或丢弃了看似安全的信息，可能会在几步之后导致质量严重降解。这意味着计算分配不能当作一个独立的、逐步骤的启发式问题，而应该被建模为一个顺序决策问题。为了解决这一挑战，作者提出了 Self-Optimizing Language Models (SOL)。核心思想是冻结基座 LLM，并添加一个极其轻量的策略网络。这个策略网络读取 LLM 的隐藏状态以及一些简单的进度/预算特征，然后在每个解码步骤选择一个离散的效率动作。每个动作通过控制：(i) token 级注意力稀疏性、(ii) 结构化 MLP 激活剪枝和/或 (iii) 激活量化位宽，来实例化同一个模型内部的不同计算机制。一个关键的贡献是部署时的可控性。SOL 不是为每个模型硬编码一个工作点，而是允许策略网络在推理时接收一个用户定义的计算预算。这使得一个训练好的策略可以动态适应不同场景下的性能与效率需求。总而言之，引言非常清晰地回答了核心问题：为什么静态预算分配是低效的，以及 SOL 如何通过学习式的、联合控制的顺序决策来解决这一问题。

方法：核心思路与技术路线

SOL 的整体流水线可以分为三个核心部分：策略网络的设计、动作空间的定义，以及基于强化学习的训练方法。

策略网络：轻量级预算分配者

SOL 的核心是将一个给定的、经过预训练的大语言模型视为黑箱（或白箱），冻结其所有权重。在此基础上，引入一个轻量的策略网络。这个策略网络的架构被设计为一个小型自回归 Transformer。其输入包括两方面特征：

LLM 的隐藏状态：在每个解码步骤，策略网络会读取 LLM 中某层（例如最后一层）的隐藏向量，从中捕捉当前生成上下文的难度和性质。
进度与预算特征：除了隐藏状态，策略网络还会接收关于当前生成进度的特征（例如，已经产生了多少个 token）以及一个请求的预算目标。这个目标可以在训练和推理时动态指定。

基于上述观察，策略网络输出一个离散的动作。这个动作决定了当前步骤的计算效率设置。原文强调，策略网络引入的开销极小，其计算成本与 LLM 本身的推理成本相比几乎可以忽略不计（论文图 1 中展示了控制开销与控制时长成正比，例如每 16 个解码步骤产生一次开销）。

动作空间：联合控制注意力稀疏性、MLP剪枝与量化

SOL 的动作空间不是一个单一维度的选择，而是一个能够联合控制三种不同效率机制的离散集合。论文描述了每个动作如何实例化一个具体的计算机制：

Token 级注意力稀疏性：在标准 Transformer 的自注意力层中，每个 token 会对 KV 缓存中的所有键值对计算注意力权重。SOL 的动作可以控制每个 token 关注多少个过去的 token。例如，动作可以设定一个“保留率” kappa（κ），代表保留注意力权重最高的前 k% 的 token，其余 token 的被忽略，从而实现稀疏注意力。这个操作直接减少了注意力矩阵的计算量和 KV 缓存访问次数。
结构化 MLP 激活剪枝：在前馈网络层，SOL 的动作可以控制结构化激活剪枝的强度。这里不是剪除权重，而是在推理时根据激活的幅度，裁剪掉一部分神经元或通道。论文中提到使用一个“保留率” rho（ρ），表示保留前 r% 的激活。这可以有效减少 MLP 的浮点运算量。
激活量化位宽：除了控制信息的“量”（稀疏性或剪枝比例），SOL 还能控制信息的“精度”。动作可以指定激活值被量化到多少位，例如 4 位、8 位或 16 位。这直接影响了内存带宽占用和矩阵乘法中的位操作速度。

这三种机制可以联合控制。策略网络输出的动作是一个联合离散索引，它同时指定了 kappa、rho 和量化位宽的具体数值，形成一种配置组合。论文实验部分展示了不同粒度的动作集，例如 8 个联合动作、27 个联合动作（三个维度各三个选项）和 1560 个联合动作（通过更细粒度的离散化实现）。

训练算法：组相对策略优化（GRPO）与教师强迫训练

如何训练这个策略网络是一个关键挑战。直接使用来自人类反馈的强化学习（RLHF）或者标准的策略梯度方法，会因为语言模型生成的离散 token 而引入高方差和训练不稳定。SOL 采用了一种简洁而有效的方案：

教师强迫（Teacher-forcing）：在训练时，固定一个自回归生成的 token 序列，而不是让模型实时生成。这意味着训练时使用的 token 路径是确定的，消除了 token 生成的不确定性带来的波动。
反事实采样：给定这样一个固定的 token 序列，训练过程会从当前策略中采样多个不同的计算调度（轨迹）。每个调度对应一系列逐步骤的效率动作（kappa, rho, 位宽组合）。这些轨迹对当前 token 路径来说是“反事实”的，因为它们只改变了每个步骤的效率动作，而不改变生成的 token 序列。
组相对策略优化 (GRPO)：每一组被采样的轨迹（通常是固定的上下文长度，例如一个“片段”或“episode”，长度为 T 个 token）会被评估其质量。对于每个轨迹，计算一个奖励分数。这个奖励分数平衡了两个目标：
语言模型质量 通过比较在相同监督下（即 Teacher-forced 的预测目标），不同计算调度下的生成 loss。更精准的度量是：对于一条轨迹，使用其对应的效率动作让 LLM 以受限的计算进行计算，然后测量生成序列相对于真实监督（下一个 token）的损失或困惑度。论文将此视为“质量”指标。
预算使用惩罚 为了确保策略学习到在不同预算下工作，奖励中加入了软惩罚（soft penalty）。如果某个计算调度（一个片段中所有 token 的平均预算使用）与用户请求的预算目标偏差过大，就会受到惩罚。这样，策略不仅学会了在给定预算下达到最佳质量，还学会了忠实地遵守预算要求。
优化：基于小组内所有采样轨迹的奖励，GRPO 算法会计算一个相对优势，然后用标准的策略梯度更新策略网络的参数。

这种方法的核心优势在于，它完全在教师强迫环境下运行，避开了在线生成带来的文本多样性问题，同时通过反事实采样高效地探索了不同的计算调度空间。

实验：设置、指标与结果

原文未明确说明具体的实验设置，例如明确提及使用的全部语言模型的大小列表、数据集详情或超参数表。但从论文正文与图注中，可以提取出下面这些实验设计元素和关键结果。

评估指标

论文主要使用两个评估维度来衡量 SOL 的性能：

质量：使用语言模型的困惑度（Perplexity, PPL）或标准 NLP 基准（如 MMLU）的准确率来评估生成文本的质量。
效率：报告计算预算的实际使用量，例如每个 token 的平均操作数（FLOPS）、KV 缓存保留率（kappa）、MLP 激活保留率（rho），或归一化的量化比率（eta）。论文通过绘制“质量-效率”帕累托前沿来展现不同预算下的权衡。

主要结果

优于静态分配与随机搜索：在所有实验设置和模型变体中，SOL 在匹配或更低的计算预算下，始终优于两种基线：
静态分配 对每个 token 应用相同的效率设置（例如，固定的注意力稀疏度和量化位宽）。
强随机调度搜索 随机采样一个预定义的效率调度池，然后从中选出在验证集上最好的一个。SOL 通过学习，找到了比随机搜索更好的调度策略。
帕累托前沿的显著提升：论文的图 2（关于联合控制策略）展示了关键结果。当使用不同粒度的动作集时（8、27、1560 个联合动作），SOL 永远位于质量-效率帕累托前沿的左上方。这意味着，对于任何一个目标效率水平，SOL 都能实现比基线更高的质量（更低困惑度）；或者在同等质量下，消耗更少的计算资源。
MMLU 准确率的提升：论文摘要中提到，SOL 将 MMLU 准确率提升了高达 7.3%。这意味着，当用户设定一个合理的计算预算（例如接近均匀分配预算），SOL 通过学习动态分配，能够在总体预算没有显著增加甚至减少的情况下，大幅提升生成结果的正确性。这直接验证了引言中“对复杂 token 投入更多计算”的理念。

深入分析：消融与分析

KV 缓存刷新周期的影响：图 3 展示了控制器的控制周期（即策略网络输出一个动作后，后续 T 个 token 沿用同一动作，T = 4、16、64）对结果的影响。结果显示，适当的控制周期（T = 16）能够在频繁调整的灵活性与低开销之间取得最佳平衡。太短的 T（4）虽然控制更细粒度，但额外开销可能抵消收益；太长的 T（64）则导致预算分配不够动态，性能有所下降。
预算跟随能力：图 4 展示了一个非常有趣的分析。SOL 被训练为在多个不同的请求预算下工作。实验测量了请求预算（例如请求的注意力保留率 κ）与实际执行时 token 平均预算之间的关系。令人惊讶的是，三条曲线（注意力、MLP、量化）都非常接近对角线，这说明策略网络非常完美地学会了根据用户请求来调节实际计算预算。
模型大小的泛化性：图 5 探索了 SOL 在不同大小的 LLM 上的适应性。结果显示，无论模型大小如何，SOL 都能够在更低的平均计算量下取得更低的困惑度，证明了该方法的鲁棒性和通用性。

结论：贡献、局限与启发

贡献

Self-Optimizing Language Models 的主要贡献可以总结为以下四点：

引入可学习的逐 token 预算分配：首次将 LLM 解码中的计算分配问题正式定义为一个顺序决策问题，并提出了一个基于强化学习的解决方案。
统一的多机制动作空间：不是只针对某一种效率技术（如仅稀疏注意力），而是创新地在同一个策略框架下联合控制了注意力稀疏性、结构化 MLP 剪枝和激活量化三种互补的效率机制。
实用的训练方法：通过在教师强迫片段上使用 GRPO 训练策略网络，并在奖励中嵌入预算控制，使得策略既能够学习如何分配计算，又能够忠实地遵循用户指定的预算。
可部署的轻量性：策略网络本身极轻，推理时对基座 LLM 的改动为零（只读隐藏状态），使其易于集成到现有系统中。

局限

原文未明确说明局限性。但根据论文的设定和领域一般认知，我们可以推断一些操作上的限制：

每个效率机制的独立性 SOL 的策略将一个离散动作映射到一组效率参数，但这些参数（κ, ρ, 位宽）在当前实现中可能是彼此独立的。是否存在一种更优的、协同控制的模式（例如，当注意力稀疏值高时，适当增加 MLP 的保留率）？目前策略网络需要自己学习这种关系。
策略网络的额外开销 虽然“轻量”，但每 T 个 token 就需要运行一次策略网络，这仍然引入了额外的计算延迟。对于一些时序敏感的推理任务，这种微小的开销仍可能被考虑。
与微调或知识蒸馏的交互 SOL 冻结了基座 LLM。如果基座本身经过了任务特定的微调，SOL 的策略是否需要重新训练？论文没有讨论这一点。
稳定性与边界情况 极端预算下（如极高压缩率），策略性能可能会退化，但论文没有深入讨论这种退化行为。

启发

这篇论文启发了一个新的研究方向：计算预算的组织形式与分配策略本身也是一个值得精心设计的变量。它鼓励未来研究不再仅仅是追求“更便宜的基座模型”，而是追求“更聪明的计算调度器”。此外，SOL 这种“冻结主体 + 小策略控制”的范式，跨越了模型训练的边界，使得推理优化成为可立即部署且高度可控的环节。这项工作推动 LLM 推理效率从“静态压缩”迈向“动态调度”的新时代。