In-Place Test-Time Training

Samenvatting

Het statische "trainen en dan implementeren"-paradigma beperkt grote taalmodellen (LLM's) fundamenteel in hun vermogen om hun gewichten dynamisch aan te passen aan continue stromen van nieuwe informatie die inherent zijn aan real-world taken. Test-Time Training (TTT) biedt een overtuigend alternatief door een subset van modelparameters (snelle gewichten) bij te werken tijdens de inferentiefase, maar het potentieel ervan in het huidige LLM-ecosysteem wordt belemmerd door kritieke barrières, waaronder architectuurincompatibiliteit, computationele inefficiëntie en niet-uitgelijnde doelstellingen voor snelle gewichten bij taalmodellering. In dit werk introduceren we In-Place Test-Time Training (In-Place TTT), een raamwerk dat LLM's naadloos voorziet van Test-Time Training-vermogen. In-Place TTT behandelt de laatste projectiematrix van de alomtegenwoordige MLP-blokken als aanpasbare snelle gewichten, wat een "drop-in"-verbetering voor LLM's mogelijk maakt zonder kostbare hertraining vanaf nul. Verder vervangen we het generieke reconstructiedoel van TTT door een toegesneden, theoretisch onderbouwd doel dat expliciet is afgestemd op de Next-Token-Prediction-taak die autoregressieve taalmodellering beheerst. Dit principiële doel, gecombineerd met een efficiënt chunk-gewijs update-mechanisme, resulteert in een zeer schaalbaar algoritme dat compatibel is met context-parallelisme. Uitgebreide experimenten valideren de effectiviteit van ons raamwerk: als een in-place verbetering stelt het een model van 4B parameters in staat superieure prestaties te leveren bij taken met contexten tot 128k, en wanneer het vanaf nul wordt voorgetraind, presteert het consistent beter dan competitieve TTT-gerelateerde benaderingen. Resultaten van ablatiestudies bieden verder dieper inzicht in onze ontwerpkeuzes. Collectief vestigen onze resultaten In-Place TTT als een veelbelovende stap naar een paradigma van continu leren in LLM's.

English

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.