arXiv:2605.10875cs.LGcs.CL2026-05中稿 · ICML

让大模型动态分配计算资源,自动决定每步该用多少算力。

Compute Where it Counts: Self Optimizing Language Models

论文配图:Compute Where it Counts: Self Optimizing Language Models
图 1 · 摘自论文原文
  • 用轻量策略网络根据隐藏状态动态选择每步算力使用方式。
  • 在相同算力下,比固定分配提升7.3%的MMLU准确率。
  • 适合需要高效推理的场景,如移动端或实时应用。

高效的大语言模型推理研究长期聚焦于降低每个解码步骤的成本(如量化、剪枝或稀疏注意力),通常对每个生成标记采用统一的计算预算。实际上,标记的难度差异显著,静态压缩可能导致简单步骤过度计算,复杂步骤计算不足。本文研究自回归解码中的动态预算分配:从单个模型内部学习每标记应投入多少计算。自优化语言模型(SOL)将冻结的LLM与一个轻量级策略网络结合,策略网络读取LLM隐藏状态,并在每个解码步骤选择离散的效率动作。这些动作可联合控制(i)标记级别的注意力稀疏性,(ii)MLP中结构化激活剪枝,以及(iii)激活量化位宽,同时保持基础模型权重不变。我们使用组相对策略优化在教师强制轨迹上训练策略:固定标记序列,采样多个计算调度(即仅效率动作不同的“反事实”调度),并比较其在相同监督下的似然性。奖励函数在语言模型质量与软惩罚之间权衡,鼓励整个轨迹的平均预算匹配目标。在不同模型变体和计算环境下,SOL在匹配预算时优于静态分配和强随机调度搜索,提供一种互补的推理效率优化路径。SOL在所有实验中均发现更优的质量-效率帕累托前沿,并使MMLU准确率最高提升7.3%。

原文摘要 · Abstract (English)

Efficient LLM inference research has largely focused on reducing the cost of each decoding step (e.g., using quantization, pruning, or sparse attention), typically applying a uniform computation budget to every generated token. In practice, token difficulty varies widely, so static compression can over-compute on easy steps and under-compute on hard ones. We study dynamic budget allocation for autoregressive decoding: learning how much computation to spend per token from within a single model. Self-Optimizing Language Models (SOL) pair a frozen LLM with a lightweight policy network that reads the LLM hidden state and selects a discrete efficiency action at each decode step. Actions can jointly control (i) token-level attention sparsity, (ii) structured activation pruning in the MLP, and (iii) activation quantization bit-width, while leaving the base model weights unchanged. We train the policy with group-relative policy optimization on teacher-forced episodes: the token sequence is fixed, while we sample multiple compute schedules (i.e., "counterfactual" schedules that vary only the efficiency actions for the same token path) and compare their likelihoods under the same supervision. Our reward trades off language-model quality against soft penalties that encourage episode-average budget usage to match a requested target. Across model variants and compute regimes, SOL improves quality at matched budget over static allocation and strong random schedule search, offering a complementary axis for inference-efficiency optimization. SOL discovers a better quality-efficiency pareto-front across all our experiments and improves MMLU accuracy by up to 7.3% over uniform budget allocation strategies.

大模型推理动态调度效率优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。