실시간 테스트 환경에서의 학습: In-Place Test-Time Training

In-Place Test-Time Training

기존의 '학습 후 배포' 방식은 대규모 언어 모델(LLM)이 실제 작업에서 발생하는 지속적인 새로운 정보에 맞춰 모델 가중치를 동적으로 조정하는 것을 근본적으로 제한합니다. 테스트 시간 학습(TTT)은 추론 시간에 모델 파라미터의 일부(빠른 가중치)를 업데이트하여 유망한 대안을 제시하지만, 현재 LLM 생태계에서 TTT의 잠재력은 아키텍처 호환성 부족, 계산 효율성 저하, 그리고 언어 모델링을 위한 잘못된 빠른 가중치 목표 등의 중요한 문제로 인해 제한됩니다. 본 연구에서는 In-Place Test-Time Training (In-Place TTT)이라는 프레임워크를 소개합니다. In-Place TTT는 LLM에 테스트 시간 학습 기능을 원활하게 부여합니다. In-Place TTT는 널리 사용되는 MLP 블록의 최종 투영 행렬을 적응 가능한 빠른 가중치로 처리하여, LLM에 초기 학습 없이 쉽게 적용할 수 있는 향상 기능을 제공합니다. 또한, TTT의 일반적인 재구성 목표를 언어 모델링의 핵심인 다음 토큰 예측 작업에 명시적으로 부합하는, 이론적으로 뒷받침되는 맞춤형 목표로 대체했습니다. 이러한 체계적인 목표와 효율적인 청크 기반 업데이트 메커니즘은 확장 가능한 알고리즘을 구현하며, 컨텍스트 병렬 처리를 지원합니다. 광범위한 실험을 통해 본 프레임워크의 효과성을 검증했습니다. In-Place TTT는 기존 모델에 추가하여 적용 시, 최대 128k의 컨텍스트를 가진 작업에서 40억 파라미터 모델의 성능을 향상시키며, 처음부터 학습할 경우에도 기존의 TTT 관련 접근 방식보다 뛰어난 성능을 보입니다. 추가적인 분석을 통해 설계 결정에 대한 심층적인 이해를 제공합니다. 종합적으로, 본 연구 결과는 In-Place TTT가 LLM의 지속적인 학습 패러다임으로 나아가는 유망한 단계임을 보여줍니다.

The static ``train then deploy" paradigm fundamentally limits Large Language Models (LLMs) from dynamically adapting their weights in response to continuous streams of new information inherent in real-world tasks. Test-Time Training (TTT) offers a compelling alternative by updating a subset of model parameters (fast weights) at inference time, yet its potential in the current LLM ecosystem is hindered by critical barriers including architectural incompatibility, computational inefficiency and misaligned fast weight objectives for language modeling. In this work, we introduce In-Place Test-Time Training (In-Place TTT), a framework that seamlessly endows LLMs with Test-Time Training ability. In-Place TTT treats the final projection matrix of the ubiquitous MLP blocks as its adaptable fast weights, enabling a ``drop-in" enhancement for LLMs without costly retraining from scratch. Furthermore, we replace TTT's generic reconstruction objective with a tailored, theoretically-grounded objective explicitly aligned with the Next-Token-Prediction task governing autoregressive language modeling. This principled objective, combined with an efficient chunk-wise update mechanism, results in a highly scalable algorithm compatible with context parallelism. Extensive experiments validate our framework's effectiveness: as an in-place enhancement, it enables a 4B-parameter model to achieve superior performance on tasks with contexts up to 128k, and when pretrained from scratch, it consistently outperforms competitive TTT-related approaches. Ablation study results further provide deeper insights on our design choices. Collectively, our results establish In-Place TTT as a promising step towards a paradigm of continual learning in LLMs.