Entrenamiento de Prueba de Extremo a Extremo para Contexto Largo

Resumen

Formulamos el modelado del lenguaje de contexto largo como un problema de aprendizaje continuo en lugar de un problema de diseño arquitectónico. Bajo esta formulación, utilizamos únicamente una arquitectura estándar: un Transformer con atención de ventana deslizante. Sin embargo, nuestro modelo continúa aprendiendo durante el tiempo de prueba mediante la predicción del siguiente token en el contexto dado, comprimiendo el contexto que lee en sus pesos. Adicionalmente, mejoramos la inicialización del modelo para el aprendizaje durante la prueba mediante meta-aprendizaje durante el entrenamiento. En general, nuestro método, una forma de Entrenamiento en Tiempo de Prueba (TTT), es de Extremo a Extremo tanto en el tiempo de prueba (mediante predicción del siguiente token) como en el tiempo de entrenamiento (mediante meta-aprendizaje), a diferencia de formas anteriores. Realizamos experimentos exhaustivos centrándonos en las propiedades de escalado. En particular, para modelos de 3B entrenados con 164B tokens, nuestro método (TTT-E2E) escala con la longitud del contexto de la misma manera que un Transformer con atención completa, mientras que otros, como Mamba 2 y Gated DeltaNet, no lo hacen. Sin embargo, de manera similar a las RNN, TTT-E2E tiene una latencia de inferencia constante independientemente de la longitud del contexto, lo que lo hace 2.7 veces más rápido que la atención completa para un contexto de 128K. Nuestro código está disponible públicamente.

English

We formulate long-context language modeling as a problem in continual learning rather than architecture design. Under this formulation, we only use a standard architecture -- a Transformer with sliding-window attention. However, our model continues learning at test time via next-token prediction on the given context, compressing the context it reads into its weights. In addition, we improve the model's initialization for learning at test time via meta-learning at training time. Overall, our method, a form of Test-Time Training (TTT), is End-to-End (E2E) both at test time (via next-token prediction) and training time (via meta-learning), in contrast to previous forms. We conduct extensive experiments with a focus on scaling properties. In particular, for 3B models trained with 164B tokens, our method (TTT-E2E) scales with context length in the same way as Transformer with full attention, while others, such as Mamba 2 and Gated DeltaNet, do not. However, similar to RNNs, TTT-E2E has constant inference latency regardless of context length, making it 2.7 times faster than full attention for 128K context. Our code is publicly available.

Entrenamiento de Prueba de Extremo a Extremo para Contexto Largo

End-to-End Test-Time Training for Long Context

Resumen

Support