Entrenamiento en el Momento de la Prueba In Situ

Resumen

El paradigma estático de "entrenar y luego desplegar" limita fundamentalmente la capacidad de los Modelos de Lenguaje Grandes (LLMs) para adaptar dinámicamente sus pesos en respuesta a los flujos continuos de nueva información inherentes a las tareas del mundo real. El Entrenamiento en Tiempo de Prueba (TTT) ofrece una alternativa convincente al actualizar un subconjunto de parámetros del modelo (pesos rápidos) durante la inferencia; sin embargo, su potencial en el ecosistema actual de LLMs se ve obstaculizado por barreras críticas, como la incompatibilidad arquitectónica, la ineficiencia computacional y los objetivos desalineados de los pesos rápidos para el modelado del lenguaje. En este trabajo, presentamos el Entrenamiento en Tiempo de Prueba In-Situ (In-Place TTT), un marco que dota a los LLMs de capacidad de TTT de forma integral. In-Place TTT trata la matriz de proyección final de los ubicuos bloques MLP como sus pesos rápidos adaptables, permitiendo una mejora "plug-and-play" para los LLMs sin el costoso reentrenamiento desde cero. Además, reemplazamos el objetivo genérico de reconstrucción del TTT con un objetivo específico, fundamentado teóricamente y explícitamente alineado con la tarea de Predicción del Siguiente Token que rige el modelado autoregresivo del lenguaje. Este objetivo fundamentado, combinado con un mecanismo eficiente de actualización por bloques, da como resultado un algoritmo altamente escalable compatible con el paralelismo de contexto. Experimentos exhaustivos validan la efectividad de nuestro marco: como mejora in-situ, permite que un modelo de 4B parámetros logre un rendimiento superior en tareas con contextos de hasta 128k tokens, y, cuando se preentrena desde cero, supera consistentemente a enfoques competitivos relacionados con el TTT. Los resultados del estudio de ablación proporcionan además una comprensión más profunda de nuestras decisiones de diseño. Colectivamente, nuestros resultados establecen a In-Place TTT como un paso prometedor hacia un paradigma de aprendizaje continuo en los LLMs.

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.