将最优控制引入大模型推理，实现高效规划能力提升
Beyond Test-Time Memory: State-Space Optimal Control for LLM Reasoning
- 在推理时通过LQR规划潜空间未来状态，构建可嵌入的控制层
- 在MATH-500上数学推理提升27.8%，AMC/AIME上通过率提高2-3倍
- 硬件优化的并行求解器支持高效部署，适合需要强推理的模型
关联记忆长期支撑序列模型设计。人类不仅回忆过去，更通过预测未来状态并选择目标导向行为来推理，这一能力对现代语言模型日益重要但未被原生支持。现有方法依赖强化学习或测试时训练，但规划仍独立于模型架构。本文将推理建模为最优控制问题，提出测试时控制（TTC）层，在推理阶段对潜空间执行有限时域LQR规划，将价值函数嵌入神经网络结构，并作为嵌套目标驱动规划前预测。为确保可扩展性，基于辛几何推导出硬件高效的LQR求解器，并实现为融合CUDA核，支持低开销并行计算。作为适配器集成至预训练大模型后，TTC层在MATH-500上提升数学推理性能达27.8%，在AMC和AIME上实现Pass@8指标提升2-3倍，证明将最优控制作为架构组件是超越测试时训练的有效且可扩展的推理机制。
原文摘要 · Abstract (English)
Associative memory has long underpinned the design of sequential models. Beyond recall, humans reason by projecting future states and selecting goal-directed actions, a capability that modern language models increasingly require but do not natively encode. While prior work uses reinforcement learning or test-time training, planning remains external to the model architecture. We formulate reasoning as optimal control and introduce the Test-Time Control (TTC) layer, which performs finite-horizon LQR planning over latent states at inference time, represents a value function within neural architectures, and leverages it as the nested objective to enable planning before prediction. To ensure scalability, we derive a hardware-efficient LQR solver based on a symplectic formulation and implement it as a fused CUDA kernel, enabling parallel execution with minimal overhead. Integrated as an adapter into pretrained LLMs, TTC layers improve mathematical reasoning performance by up to +27.8% on MATH-500 and 2-3x Pass@8 improvements on AMC and AIME, demonstrating that embedding optimal control as an architectural component provides an effective and scalable mechanism for reasoning beyond test-time training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。