让语言模型按需计算，难词多算、易词少算，节省10%推理开销。
AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth
- 每令牌自适应决定停止迭代时机，无需人工设定剪枝比例。
- 在Pythia模型上减少约10%计算量，保持相似困惑度和下游性能。
- 首次实现全自监督下的自适应计算时间，适合高效部署场景。
通过循环/迭代Transformer实现推理时的计算扩展，使大语言模型能在推理阶段消耗更多计算，但大多数预训练循环语言模型固定迭代次数，导致对简单令牌浪费计算且缺乏令牌级自适应能力。受自适应计算时间（ACT）与提前退出（EE）思想启发，我们提出AdaPonderLM，一种在预训练中学习令牌级提前退出策略的自监督循环语言模型，无需人工设定每令牌或每层的剪枝比例。AdaPonderLM使用依赖迭代次数的MLP门控机制与单调终止掩码，决定每个令牌何时停止迭代，并引入键值状态复用机制，重用已终止令牌的缓存键值状态，确保训练-测试一致性并实现实际加速。在从70M到410M（预训练）及最高2.8B（持续预训练）的Pythia骨干模型上，AdaPonderLM将推理计算量降低约10%，同时保持相近的语言建模困惑度和具有竞争力的下游准确率。分析表明，学习到的门控机制为高负对数似然（即难样本）分配更多计算，展现出完全自监督设置下的自适应计算时间行为。同时，在同等浮点运算量下，学习到的终止策略始终优于固定剪枝，说明AdaPonderLM能将计算资源精准分配给需要的令牌而非仅降低平均深度。
原文摘要 · Abstract (English)
Test-time scaling via recurrent/iterative Transformers enables large language models to spend more computation at inference, but most pretrained recurrent LMs run a fixed number of iterations, wasting compute on easy tokens and lacking token-wise adaptivity. Following the core idea of Adaptive Computation Time(ACT) and Early Exit(EE), we propose AdaPonderLM, a self-supervised recurrent language model that learns token-wise early exiting during pretraining without manually tuned per-token/per-layer pruning ratios. AdaPonderLM uses iteration-specific MLP gates with a monotonic halting mask to decide when each token stops recurring, and introduces a KV reuse mechanism that reuses cached key/value states for halted tokens, ensuring train--test consistency and practical acceleration. Across Pythia backbones from 70M to 410M (pretraining) and up to 2.8B (continued pretraining), AdaPonderLM reduces inference compute at about 10% while maintaining comparable language modeling perplexity and competitive downstream accuracy. Our analysis shows the learned gates allocate more computation to high-NLL (hard) tokens, exhibiting adaptive computation time behavior in a fully self-supervised setting. Meanwhile, under iso-FLOPs, the learned halting policy consistently outperforms fixed pruning, showing AdaPonderLM allocates compute to the right tokens rather than just reducing average depth.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。