테스트 타임 트레이닝을 넘어서: 하드웨어 효율적 최적 제어를 통한 추론 학습

초록

연관 메모리는 오랫동안 순차 모델 설계의 기반이 되어 왔다. 단순한 기억 재생을 넘어 인간은 미래 상태를 예측하고 목표 지향적 행동을 선택하며 사고하는데, 이는 현대 언어 모델이 점점 더 필요로 하지만 기본적으로 내재하고 있지 않은 능력이다. 기존 연구에서는 강화 학습이나 테스트 시간 학습을 사용하지만, 계획 수립은 모델 아키텍처 외부에서 이루어진다. 우리는 사고를 최적 제어 문제로 공식화하고, 추론 시점에 잠재 상태에 대해 유한 수평 LQR 계획을 수행하며, 신경망 아키텍처 내에서 가치 함수를 표현하고, 이를 예측 전 계획을 가능하게 하는 중첩 목적 함수로 활용하는 테스트 시간 제어(TTC) 계층을 소개한다. 확장성을 보장하기 위해 심플렉틱 공식화를 기반으로 하드웨어 효율적인 LQR 솔버를 유도하고, 최소 오버헤드로 병렬 실행이 가능한 통합 CUDA 커널로 구현한다. 사전 학습된 대형 언어 모델에 어댑터로 통합된 TTC 계층은 MATH-500에서 최대 27.8%의 수학적 추론 성능 향상과 AMC 및 AIME에서 2-3배의 Pass@8 성능 향상을 보여주며, 최적 제어를 아키텍처 구성 요소로 내재화하는 것이 테스트 시간 학습을 넘어선 효과적이고 확장 가능한 추론 메커니즘을 제공함을 입증한다.

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.

테스트 타임 트레이닝을 넘어서: 하드웨어 효율적 최적 제어를 통한 추론 학습

Beyond Test-Time Training: Learning to Reason via Hardware-Efficient Optimal Control

초록

Support