Gradient Descent at Test Time: Titans Pushes LLM Context to 10 Million Tokens

Nathan Labenz opens The Cognitive Revolution with a problem he has been circling for years. Today's large language models carry a superhuman amount of world knowledge in their weights, and they have attention-based working memory in the context window. What they still lack is the layer in between: an integrated, persistent, ever-evolving longer-term memory that humans use to keep coherence and identity over time. He calls it the missing middle, and he thinks it is a necessary piece for high-context work and for agents that have to run for a long time in the real world.

The paper on the table is Titans: Learning to Memorize at Test Time, led by Ali Behrouz, a PhD student at Cornell. The title is the thesis. After training ends, the model is still learning — and what it is learning is its own memory. Unlike RAG systems that store data explicitly and search it, or Mamba-style state-space models that encode memory as a matrix updated on each forward pass, Titans uses a neural network as the memory module and updates that network by gradient descent at runtime. Labenz calls the shift a qualitative conceptual advance.

The scoreboard that makes the paper hard to ignore comes later in the conversation, when Behrouz walks through long-context results. On sequences stretching to 10 million tokens, GPT-4's accuracy falls off quickly. Titans, with a very small parameter count by comparison, still holds around 70% accuracy, and within 2 million tokens it outperforms a row of counterpart architectures. Behrouz is careful to say those numbers are on synthetic datasets and should not be naively extrapolated to general tasks. Even so, Labenz treats the result as evidence that the missing middle might be the last major technical unlock before AI hits an inflection in economic value.

Labenz starts with an economist's line. Tyler Cowen has lately gone "reverse Oprah" in front of influential rooms, telling the audience they are the bottleneck. Cowen's more famous line is "context is that which is scarce." Labenz has never been sure what that means for humans, but for AI it is literal. Today's models know a superhuman amount about the world at large and, out of the box, almost nothing about the people and businesses they are supposed to serve. Useful context is scattered across Slack, email, documents, GitHub, meeting transcripts, and task systems, and assembling it by hand is tedious. Only recently — he dates the shift to Gemini 2.5 Pro — have models started to properly reward the work of feeding them that context.

He then runs the thought experiment at enterprise scale. Imagine a model trained for GE or 3M that knows the company's products, history, teams, internal debates, and finances as deeply as it knows the world. No company, he notes, has more than 1% of the data on which current models are already trained, so the data volume is not the blocker. Such a model would know more about the firm than any single employee, still search official records to ground its work, and pick up the unwritten patterns of "how we do things around here." Startup costs could reach the millions or tens of millions of dollars — roughly where OpenAI has publicly priced custom models — which for a hundred-year-old enterprise is still a bargain if the system can drop in and do a significant share of the work.

For smaller businesses with much less data, he uses today's fine-tuning prices as an anchor and guesses tens of thousands, maybe into the low hundreds of thousands. Individual ChatGPT users are getting a glimpse of this future from the product's new memory features, but Labenz is not aware of a productized version that works well at scale. He thinks that is ultimately because the gap is still in the foundation models themselves, and scaffolding is not quite enough to fix it. The plausible punchline: the main thing between today and a labor market full of drop-in knowledge workers may simply be a breakthrough in long-term memory.

Behrouz's path into the problem is unusual. He started in graph algorithms during his master's, drifted into graph neural networks for anomaly detection, and then noticed that those tasks connected to neuroscience work on detecting brain disorders. That sent him into how different parts of the brain work, and into a working bet: human learning is efficient because of millions of years of evolution, and it is hard to believe researchers can skip those steps and invent something more efficient from scratch. At least for now, he argues, architectures should copy the brain. He also plants a flag against a common split in the field. Some researchers want to replace attention entirely with RNNs; he is in the hybrid camp, because attention is still necessary for accurate modeling of token dependencies inside the current context.

The shape of memory is the design choice he refuses to simplify. A traditional RNN stores memory as a vector, often called the hidden state. Linear attention, written as a recurrent network, stores it as a matrix. The brain, he says, is not three neurons encoding everything — it is a large set of interconnected neurons, which is to say a network. Treating memory as a vector or a matrix is, in his view, an oversimplification once models are trained at scale. Titans therefore makes the memory module itself an MLP, updated by gradient descent at every inference step, so information inside memory can interact rather than sitting as a passive table of numbers. He is explicit that the MLP is a starting point, not the last word: the paper opens a research direction on what architecture that submodule should actually be.

Write and read follow the same associative-memory picture as attention. To write, you feed a token's key into the memory module and train it to output the corresponding value. To read, you send a new token's query through a forward pass and take back historically related information. If memory is a vector you multiply; if memory is a network you run a forward pass. Labenz, working through the paper with Grok 3, landed on the same intuition: attention is associative memory that stores key-value pairs and retrieves by similarity. Attention is the nonparametric exact solution to that assignment, which is why its memory cost grows with the square of sequence length. An RNN is gradient descent approximating the same objective, trading some precision for constant-cost updates. The fact that architectures designed from that textbook picture actually work, Labenz says, is evidence the picture is not a just-so story.

The analogy to human memory is load-bearing for Behrouz. RNNs have an inherent fading: you are always adding new data and ignoring some of what you had, which is why they feel like long-term memory — you can hallucinate details of an event from 20 years ago. Attention keeps pairwise interactions inside a window, which is why it feels like short-term memory: you do not hallucinate what you heard 20 or 30 seconds ago. How those two stores pass information to each other is still controversial even in neuroscience, which is why Titans has to invent the join rather than copy one. A smaller piece of the design, persistent memory, concatenates a few learnable, data-independent weights at the start of every sequence. The idea predates Attention Sinks: treat the Transformer's MLP as attention, swap in softmax, and you have extra tokens at the front of the sequence. In practice the gain is slight, but it is harmless and keeps causal attention from overweighting the true first tokens.

Memory is finite, so Titans has to choose what to write. Equations 13 and 14 in the paper are the update: first decay the current memory state, then add an update term. That term has two parts — the loss on the current token, plus a momentum term from the recent past. Loss is surprise. If the memory fails to predict a token, the token was unexpected, and it is worth writing hard. Momentum covers the second half of the intuition: a surprising token is often the start of an episode, and the sentences that explain it may not be surprising on their own but still need to be stored.

Behrouz's example is a colleague saying "I'm leaving." At the end of a shift, the sentence carries almost no information. If the colleague is resigning, the same words explode, and every following sentence that explains why suddenly matters. The difference is entirely in context, which is why a momentary surprise score is not enough. Surprise itself also has to decay: a startling token from 64 steps ago should fade as the context changes, and that decay is separate from the forget gate inside the memory module. All of the free parameters in the rule — decay rate, update strength, momentum coefficient — are learned and input-dependent. The model projects the current token to decide how much past information to keep, whether to use past surprise, and how hard to write the current one.

Labenz recognizes the psychology immediately. He was in high school on 9/11: he remembers who told him, where he was standing, the next class, what they talked about, what he later said to his then-girlfriend now-wife, and when his father got home. The day before and a few days after are gone. A high-surprise token, in this picture, raises the write weight on a whole following episode. Mathematically the same object is just a stronger optimizer: test-time training with gradient descent plus momentum instead of plain gradient descent. Behrouz notes you could swap in Adam and get a different architecture still. He likes momentum because it is the spot where the human story and the better optimizer coincide.

Once the long-term memory module exists, Titans still has to attach it to attention. The paper tries three joins. Memory as Context (MAC) lets long-term memory retrieve relevant history, drops that into the context, lets attention decide whether to look at the present or the past, then sends attention's output back so the memory knows what is worth storing — the two stores help each other. Memory as Gate (MAG) treats long-term and short-term memory as parallel branches on the same input and concatenates or multiplies them at the end. Memory as Layer (MAL) stacks memory layers and attention layers in sequence, which is the configuration most of the hybrid literature already uses, including Mamba interleaved with attention.

The result is an embarrassment for the default. Behrouz's important message, as he states it, is that most of the literature is focused on a configuration that is not great. Labenz roughly counts nine task categories and three scales, about 30 groups: MAL wins maybe two, and the rest split roughly evenly between MAC and MAG. People arrived at layer-interleaving by a less principled route — original Mamba had no attention, then researchers shuffled the two stacks together and got a better base — whereas MAC and MAG have a higher-order rationale, and they win almost across the board. Behrouz says there may still be some interleaving recipe that outperforms; the point is that the join itself is worth redesigning.

Long context is where MAC becomes the load-bearing choice. The memory module compresses history into a fixed-size summary for attention; attention only looks at a fixed local window; neither side has to strain, and length generalization follows. MAG and MAL lose accuracy much faster as sequences grow. The 10-million-token, 70% accuracy number is a MAC result. Behrouz again flags that this is synthetic. Sizing the two branches is still under-explored: a practical starting point is to treat memory as extra heads and split heads (or layers) half and half, but three memory heads and seven attention heads might beat five and five, and nobody has mapped that space.

The sharpest theoretical stretch of the conversation is about what these architectures can even express. A paper titled The Illusion of State in State-Space Models argues that diagonal state-space models — Mamba, RetNet, S4, S5 — and transformers are limited to the TC^0 complexity class and cannot do state tracking. Behrouz's toy version: you stand at the origin, the model reads a string of "go left" and "go right," and at the end it must say where you are. Diagonal RNNs and transformers cannot solve that in the theoretical sense; a nonlinear RNN can do it with one layer. Titans' raw memory module is a nonlinear RNN, so it sits on the capable side of that line, as does a non-diagonal linear RNN like DeltaNet. Mamba, Mamba 2, RetNet, and S4 cannot.

Labenz objects immediately: could o1-style reasoning models just talk their way through the same puzzle if you spend enough inference tokens? Behrouz grants the gap. Theoretical frameworks make simplifying assumptions; in practice, large models often complete tasks the theorems say they cannot. The Illusion paper also has experiments supporting its claim, so both pictures are real. Labenz's translation for the internet argument is that people talking past each other are usually mixing two standards: a guaranteed, always-on ability versus succeeding most of the time, perhaps with extra test-time compute. "We've proven a transformer can't do X" and "I just did X in ChatGPT" can both be true under those different bars.

Scale produces another counterintuitive picture. Behrouz points to a plot in the Mixture of Mamba paper: with less training data, RNN-style models (Mamba in that figure) scale better; once token count is huge, transformers overtake them. Titans, at least across the budgets they ran, did not show that flip — it stayed ahead of both transformers and linear RNNs as they increased training tokens. Academic papers rarely surface this, he says, because nobody can afford to train at many budgets; a fixed token count is the default posture. Labenz adds a hardware coda from Mamba 2: a coarser state update produced a better model because it ran faster, so the same GPUs saw more data. Behrouz's reply is that these models are mathematically connected, often by a few components, but a new perspective that leads to future work is worth more than any one architecture. Even transformers kept improving — Flash Attention is the example — and nothing is the endgame.

Labenz asks the data question directly. The internet is not full of million-token episodes, and long-running agents would seem to need training data that does not exist. Behrouz's answer is that you do not have to train on millions of tokens. The relevant property is length generalization: train at ordinary lengths, infer much longer. Pure RNNs and a standalone neural memory struggle early, around 16k tokens. Transformers need their own tricks. MAC is, in his experience, the architecture that generalizes. Memory only has to summarize history into a fixed-size packet for attention; attention only has to do local attention over a fixed window — 32, 128, 512, or 8,192, take your pick — and can always go farther with the memory branch. More data still helps. Architecture decides whether the extra length is usable.

Can you bolt this onto an off-the-shelf Llama or R1? Behrouz thinks it is possible in principle. MAC decouples the attention training window from memory capacity, so you could keep the pretrained transformer, add the memory module, and continue training until the model learns to look at memory-as-context. A deeper conceptual limit is the frozen latent space: the memory module always writes into the pretrained model's space, which may not be enough if an enterprise needs the model to absorb a whole new domain. Labenz wonders whether GE or 3M would still need continued pretraining so the memory has the right representations to write into. Behrouz's reply names the real obstacle: catastrophic forgetting. Long test-time training on a new task can wipe the previous one. There is no good answer yet. If someone solves it, he says, they will solve a large part of robotics and reinforcement learning at the same time.

Implementation, by contrast, he is optimistic about. Gradient updates rewrite as matrix multiplies; decay is a diagonal; the momentum term is itself a linear RNN and can be parallelized with a scan. They did not spend time on kernels, and Titans is already faster than Mamba in their training-time plots, though some modern linear models are faster still. With real kernel work he expects to match the fastest linear RNNs. He keeps restating the paper's posture: offer a new perspective, point down a path, choose the simplest implementation, and wait for others to improve it. He is not on the team that wants to kill attention. He is on the team that copies the brain and builds hybrids. The applications he most wants to see next are outside language: agents, reinforcement learning, decision-making, and other modalities, where long-term memory may matter more than it does for text, and where MAG's gated branches might beat MAC. Transformers won by working in many domains. Titans still has to pass that test.