中文

超越测试时训练:通过硬件高效最优控制学习推理

机器学习 2026-03-11 v1

摘要

联想记忆长期以来是序列模型设计的基础。除了记忆回溯之外,人类通过投影未来状态并选择目标导向的行动来进行推理,这一能力现代语言模型日益需要但并不原生编码。虽然先前工作使用强化学习或测试时训练,但规划仍然是模型架构之外的。我们将推理建模为最优控制问题,提出Test-Time Control(TTC)层,该层在推理时对潜在状态执行有限视野LQR规划,将价值函数表示为神经网络架构的一部分,并将其作为嵌套目标以实现预测前的规划。为确保可扩展性,我们基于辛曼公式推导出硬件高效的LQR求解器,并将其实现为融合CUDA内核,实现并行执行并具有最小开销。作为适配器集成到预训练的LLM中,TTC层在MATH-500上的数学推理性能提升最高可达+27.8%,在AMC和AIME上通过率提升2-3倍,表明将最优控制作为架构组件嵌入可为推理提供有效且可扩展的机制,超越测试时训练。

关键词

引用

@article{arxiv.2603.09221,
  title  = {Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control},
  author = {Peihao Wang and Shan Yang and Xijun Wang and Tesi Xiao and Xin Liu and Changlong Yu and Yu Lou and Pan Li and Zhangyang Wang and Ming Lin and René Vidal},
  journal= {arXiv preprint arXiv:2603.09221},
  year   = {2026}
}