Сквозное обучение во время тестирования для работы с длинным контекстом
End-to-End Test-Time Training for Long Context
December 29, 2025
Авторы: Arnuv Tandon, Karan Dalal, Xinhao Li, Daniel Koceja, Marcel Rød, Sam Buchanan, Xiaolong Wang, Jure Leskovec, Sanmi Koyejo, Tatsunori Hashimoto, Carlos Guestrin, Jed McCaleb, Yejin Choi, Yu Sun
cs.AI
Аннотация
Мы формулируем задачу языкового моделирования длинных контекстов как проблему непрерывного обучения, а не проектирования архитектуры. В рамках этой формулировки мы используем только стандартную архитектуру — трансформер со скользящим оконным вниманием. Однако наша модель продолжает обучение во время тестирования посредством предсказания следующего токена в заданном контексте, сжимая прочитанный контекст в своих весах. Кроме того, мы улучшаем инициализацию модели для обучения во время тестирования с помощью метаобучения на этапе тренировки. В целом, наш метод, являясь формой обучения во время тестирования, является сквозным как на этапе тестирования (через предсказание следующего токена), так и на этапе тренировки (через метаобучение), в отличие от предыдущих подходов. Мы проводим обширные эксперименты с акцентом на свойства масштабирования. В частности, для 3-миллиардных моделей, обученных на 164 миллиардах токенов, наш метод масштабируется с длиной контекста так же, как трансформер с полным вниманием, в то время как другие методы, такие как Mamba 2 и Gated DeltaNet, — нет. Однако, аналогично RNN, наш метод имеет постоянную задержку вывода независимо от длины контекста, что делает его в 2.7 раза быстрее полного внимания для контекста в 128K токенов. Наш код находится в открытом доступе.
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.