Pinned Tweet
Proud to share what I've been working on for the past year, "Learning to (Learn at Test Time)"!
Our new architecture trains a model to "learn" from its context, replacing Attention's costly KV cache with an expressive hidden state: the weights of a ML model!🤯
🧵by @karansdalal
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
Arjun Vikram retweeted
I’m excited to share that I’ll be joining OpenAI and look forward to working with the exceptional team there.
It was a difficult decision to move on. I’m incredibly proud of the amazing team at Google and everything we’ve built together. It has been an honor and a pleasure to work with all of you.
Arjun Vikram retweeted
there is much work to do. in the next few months we need to train 100x larger models on 200x as much data as we have done before, scale RL to have millions of agents learning things from the internet, and use the computer use prior to build the first models that can scalably learn. there's no team i'd rather be doing this with. if you're interested in being a part of SI's next chapter, i'd be very excited to chat.
Arjun Vikram retweeted
Computer use models shouldn't learn from screenshots.
We built a new foundation model that learns from video like humans do. FDM-1 can construct a gear in Blender, find software bugs, and even drive a real car through San Francisco using arrow keys.
Arjun Vikram retweeted
How to get AI to make discoveries on open scientific problems?
Most methods just improve the prompt with more attempts. But the AI itself doesn't improve.
With test-time training, AI can continue to learn on the problem it’s trying to solve: test-time-training.github.io…
Arjun Vikram retweeted
Incredibly excited to share our latest work with @NVIDIAAI that has been over a year in the making: 🔥End-to-End Test-Time Training 🔥
Automated researchers need to compress years of experience into intuition 🧠– far too much to just dump into context. But how do humans do research? We are continual learners 👩🎓📚
Let’s make AI do the same 🤖📚: as the model reads through long context at test time, keep training to update the weights!
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.
Arjun Vikram retweeted
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…
Test-Time Training lives on!
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…
Arjun Vikram retweeted
the way i see it, the last twelve months of AI research can be summed up in just two big breakthroughs:
[i] reasoning ('test-time compute') - new ways to train models that can use more tokens to generate better answers. they mostly rely on RL with verifiable rewards
[ii] memory ('test-time training') - moving towards models that dynamically update weights on-the-fly using gradients. this paper does it by updating a hidden state memory module to 'remember' test data
Stanford was getting boring. So I built a slip-n-slide!
Thank you to everyone who showed up and helped keep fun alive on the farm!
Arjun Vikram retweeted
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.
Arjun Vikram retweeted
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…
Arjun Vikram retweeted
What will you build this year?
Apply by Nov. 1 for Stanford’s premier hackathon. treehacks.com/
ALT Stanford Hackathon: TreeHacks. February 14-16, 2035. Hoover Tower is on the left. Apply now. We cover food, flights, and accommodations. Apps are due by November 1, 2024 at 11:59 PM.
Arjun Vikram retweeted
Feeling spooked👻🎃? Get grounded...introducing "Grounding by Trying: LLMs with Reinforcement Learning-Enhanced Retrieval."
Meet LeReT (Learning to Retrieve by Trying), a RL-based framework that improves LLM’s ability to use retrieval tools by up to 29%.
sherylhsu.com/LeReT/
Arjun Vikram retweeted
Transformers are…almost magical. Even matching performance is tough, let alone @ FLOPS.
Which is why it’s been so impressive watching TTT evolve, overcoming many hurdles! Much nuance here that's easy to miss on a skim.
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
We have implementations in both Jax (github.com/test-time-trainin…) and PyTorch (github.com/test-time-trainin…)
Could not have done this without the incredible team! Thanks to @LeoXinhaoLee @karansdalal @Jerry_XU_Jiarui @zhang677 @yanndubs, and @endernewton @xiaolonw @tatsu_hashimoto @sanmikoyejo @guestrin, and most of all our incredible project lead Yu Sun (who refuses to use twitter).