LLM memory is considered one of the hardest problems in AI.
All we have today are endless hacks and workarounds. But the root solution has always been right in front of us.
Next-token prediction is already an effective compressor. We don’t need a radical new architecture. The missing piece is to continue training the model at test-time, using context as training data.
Our full release of End-to-End Test-Time Training (TTT-E2E) with @NVIDIAAI, @AsteraInstitute, and @StanfordAILab is now available.
Blog: nvda.ws/4syfyMN
Arxiv: arxiv.org/abs/2512.23675
This has been over a year in the making with @arnuvtandon and an incredible team.
Our new paper, “End-to-End Test-Time Training for Long Context,” is a step towards continual learning in language models.
We introduce a new method that blurs the boundary between training and inference. At test-time, our model continues learning from given context using the same next-token prediction objective as training.
With this end-to-end objective, our model can efficiently compress substantial context into its weights and still use it effectively, unlocking extremely long context windows for complex reasoning and applications in agents and robotics.
Paper: test-time-training.github.io…
Code: github.com/test-time-trainin…
To motivate our method, we open with a toy example: at test-time, we train a simple bigram MLP to perform sequence modeling as effectively as attention by updating it at every token via next-token prediction on the context.
We meta-learn this test-time training procedure during (training-time) training by learning an initialization that makes these per-token updates effective.
Our paper provides a clear exposition connecting this toy experiment to our method and to prior RNN-based approaches.
Please see @arnuvtandon's thread for more details.
nitter.cf/arnuvtandon/status/200…
Excited to release a new paper today: “End-to-End Test-Time Training for Long Context”.
Our method, TTT-E2E, enables models to continue learning at test-time via next-token prediction on the given context – compressing context into model weights.
For our main result, we extend 3B parameter models from 8K to 128K. TTT-E2E scales with context length like full attention without maintaining keys and values for every token in the sequence.
With linear-complexity, TTT-E2E is 2.7x faster than full attention at 128K tokens while achieving better performance.
Paper: test-time-training.github.io…
Code: github.com/test-time-trainin…
I want to highlight some of the coolest kernel work by @danielkoceja, specifically around making TTT layers fast for training and video generation.
TLDR: Our kernel does Tensor Parallel across Streaming Multiprocessors so we can efficiently train an RNN whose hidden state is a machine learning model.
Kernels like Flash Attention succeed by leveraging the GPU memory hierarchy. For some background:
- A GPU consists of multiple Streaming Multiprocessors (SMs) – similar to CPU cores.
- All SMs share a large but slower global memory called HBM.
- Each SM also has its own small but extremely fast on-chip memory, called SMEM.
- To achieve maximum efficiency, kernels load input data once into SMEM, compute there, and then write results directly back to HBM – avoiding frequent, costly transfers between these memory types.
However, TTT layers present a challenge for this approach. The hidden state of a TTT layer is itself a machine learning model, and it's too big to fit entirely into the small on-chip SMEM!
Interestingly, this resembles a common scenario we face when training large models. If a X-billion parameter model doesn't fit on one GPU, we typically use Tensor Parallelism to split the model across several GPUs. Why can't we do something similar for the TTT layer’s hidden state model?
We actually can – by applying the same sharding idea across the SMEMs of multiple SMs, treating each SM as the analogy of a GPU. NVIDIA’s H100 GPUs come with a unique "Distributed Shared Memory" feature, allowing SMs to do an "AllReduce" without touching the slower HBM. See our "update rule" in the attached figure.
This implementation trick provides a way to continue scaling the hidden state model of TTT layers, which we believe is the approach to generating even longer videos with more complex storylines.
As a general principle, if your model architecture can leverage standard Tensor Parallelism across GPUs, you can apply that same sharding strategy across SMs when using that model as the hidden state.
Today, we're releasing a new paper – One-Minute Video Generation with Test-Time Training.
We add TTT layers to a pre-trained Transformer and fine-tune it to generate one-minute Tom and Jerry cartoons with strong temporal consistency.
Every video below is produced directly by the model in a single shot, without editing, stitching, or post-processing. Every story is newly created.
Demos: test-time-training.github.io…
Paper: test-time-training.github.io…
Grateful for wonderful collaborators. This work will be presented at CVPR 2025.
@danielkoceja @GashonHussein @Jerry_XU_Jiarui @__yuezhao__ @jankautz @guestrin @tatsu_hashimoto @sanmikoyejo @YejinChoinka @xiaolonw
I’m excited to share a project I’ve been working on for over a year, which I believe will fundamentally change our approach to language models.
We’ve designed a new architecture, which replaces the hidden state of an RNN with a machine learning model. This model compresses context through actual gradient descent on input tokens. We call our method “Test-Time-Training layers.”
TTT layers directly replace attention, and unlock linear complexity architectures with expressive memory, allowing us to train LLMs with millions (someday billions) of tokens in context.
Our instantiations, TTT-Linear and TTT-MLP, both match or beat the strongest Transformers and Mamba. Arxiv: arxiv.org/abs/2407.04620
This work would not be possible without amazing collaborators: Yu Sun, @LeoXinhaoLee, @Jerry_XU_Jiarui, @arjvik, @zhang677, @yanndubs and advisors: @endernewton, @xiaolonw, @tatsu_hashimoto, @sanmikoyejo, @guestrin.
I'd also like to thank @ayushh_agrawal, @_Aaditya_Prasad, and @arnuvtandon, who became some of my closest friends during my time working on this project at Stanford.
Check out Prof. Wang's thread detailing our experiments.
nitter.cf/xiaolonw/status/181038…
Cannot believe this finally happened! Over the last 1.5 years, we have been developing a new LLM architecture, with linear complexity and expressive hidden states, for long-context modeling. The following plots show our model trained from Books scale better (from 125M to 1.3B) than Mamba and Transformer, and our 1.3B model works better and better with longer context.
The idea, Test-Time Training (TTT), is something we have been working on for over 5 years. I still remember when I started as a postdoc, discussing with Alyosha what to work on, he asked me to talk to Yu Sun on TTT. And that is when everything begins.
Although this TTT is quite a different thing now. It is a network layer, replacing the hidden state of an RNN with a machine learning model. Instead of using a feature vector to represent a memory, our TTT layer maintains a small neural network to compress the input tokens. As each token comes in sequentially, the compression is done via gradient descent on this token to update the small neural network.
This is currently applied to language modeling, but imagine applying it to videos. In the future, when modeling long videos, instead sampling 1 FPS, we can sample the frames densely, and these dense frames would be a burden for Transformer but a blessing for TTT layer. As they are essentially just augmentations in time to train a better network inside TTT.
Please check: arxiv.org/abs/2407.04620
We’ve made the difficult decision to close our doors.
otto.pet/
We started building Otto as high schoolers with the mission of enabling accessible pet care for all. Thank you all for your support in helping us make our dreams come true.
2 years ago, @arnuvtandon and I set out to re-envision pet parenthood. To our partner shelters, customers, and veterinary specialists: Your trust and support have been invaluable in our journey.