arXiv:2603.02023cs.CL2026-03被引 2

让大模型按需分配计算量,只在必要时多算,省下资源提升效率。

PonderLM-3: Adaptive Token-Wise Pondering with Differentiable Masking

  • 用可学习的注意力掩码实现每令牌动态分配额外计算
  • 相同推理算力下比现有方法更低困惑度,性能更优
  • 适合追求高效推理、注重计算资源优化的研究与应用

测试时扩展显示,在推理阶段增加额外计算可提升生成质量,这引出一个关键问题:计算应花在何处?基于此,我们提出 PonderLM-3,一种基于 PonderLM-2 架构的自监督预训练框架,实现令牌级别的自适应计算分配。该方法通过在预训练中注入可微分注意力掩码,并在推理时采用对应硬剪枝规则,使额外计算成为可分配的令牌级资源,仅在有益时才增加计算,而非对所有令牌统一加开销。相比现有递归或自适应基线方法,PonderLM-3 在相同推理 FLOPs 条件下达到更低的预训练困惑度,且在下游任务中,与固定步数的 PonderLM-2 性能相当,但实际推理中使用更少的 FLOPs。整体上,PonderLM-3 提供了一个端到端可微、训练推理一致的令牌级自适应计算框架,使额外推理算力能精准投向最需要的位置。

原文摘要 · Abstract (English)

Test-time scaling has shown that allocating more additional computation at inference can improve generation quality, motivating a natural follow-up question: where should this computation be spent? Building on this insight, we introduce PonderLM-3, a pretraining framework for token-wise adaptive pondering that learns to selectively allocate additional computation under purely self-supervised objectives, built on top of the PonderLM-2 backbone. This makes additional inference computation an allocatable per-token resource, so tokens receive more computation only when it is beneficial, rather than paying a uniform extra cost. To make this allocation learnable while maintaining train-inference consistency, PonderLM-3 injects a differentiable attention mask during pretraining and pairs it with a matching hard pruning rule at inference. PonderLM-3 defines a stronger Pareto frontier: compared with existing recursive or adaptive baselines, it achieves lower pretraining perplexity at equal inference FLOPs. On downstream benchmarks, PonderLM-3 attains comparable performance to fixed-step PonderLM-2 under the same maximum number of additional computation steps, while using fewer inference FLOPs in practice. Overall, PonderLM-3 provides an end-to-end differentiable and train-inference consistent framework for token-wise adaptive computation, enabling additional inference compute to be allocated where it is most useful rather than paid uniformly by every token.

自适应计算推理优化可微分掩码

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