Reward at Test Time: What Happens When the Objective Reaches the State

Mamba, test-time training and Titans all update state as tokens arrive, but they use different rules. Mamba updates a recurrent state from each input. TTT takes a gradient step on a model stored in its state. Titans updates neural memory with a gradient term and learned forgetting.

Part 1 compared changes to the reward and optimiser, then ended with Jev, a model that returns typed decisions. This article asks a different question: can a task reward train the update that decides what memory persists?

Mamba is trained for next-token prediction. TTT and Titans also perform test-time updates with self-supervised memory objectives. I want to know whether a downstream reward can guide those updates.

TL;DR

A fixed-size state has a capacity limit, and its update rule determines what persists. Reinforcement learning already targets nearby choices: TTRL trains against majority-vote outcomes, SEAL rewards self-edits, Memory-R1 trains external memory operations, and Neural Garbage Collection trains KV-cache eviction. I could not find a paper that directly optimises an SSM's internal retention gate against downstream task reward. My three-slot toy tests only a discrete version of that idea.

A retention gate acts like a policy. In these models, it is trained indirectly through prediction or memory losses.

Three related updates

Start with Mamba. Its discretised recurrence is:

with $\bar{A}=\exp (\Delta A)$, where $\Delta$ is computed per token from the input. Now test-time training, whose central idea is that the hidden state should be a model rather than a vector, and the update should be a training step:

Titans uses a neural memory with a surprise term that carries momentum and a learned gate that controls forgetting:

These equations share a broad pattern, not one update rule. TTT changes its hidden model with a gradient step. Titans combines gradient updates with momentum and a learned decay. Mamba applies an input-dependent recurrence. In Titans, $(1-{\alpha }_t)$ plays a role similar to Mamba's $\bar{A}$, though their updates are not identical.

TTT's inner loss is a reconstruction:

The readout is $z_t=f({\theta }_Q{x}_t;W_t)$. The outer model learns three projections that play roles analogous to attention's key, value and query. The key and value define the inner update; the query reads the updated state.

Mamba's recurrent state, TTT's fast weights and Titans' neural memory are different representations of information carried forward through a sequence. They serve related purposes, but differ in size, update and readout.

That comparison is my interpretation, not a claim made by the papers. Test-time regression gives a framework for relating attention, state space models, fast-weight programmers and online learners through their regression weights, function classes and update algorithms. The framework makes comparison useful; it does not make the implementations identical.

TTT and Titans explicitly update a learned memory during inference. Mamba updates its recurrent state through an input-dependent recurrence. Their training and test-time objectives differ, which is why I ask whether a downstream reward could guide what persists.

A state is a memory with a budget

Mamba's default state dimension is $N=16$ per channel. Information from earlier tokens that remains useful has to fit in that recurrent state. The question is what the update preserves.

That's what $\Delta$ controls. A large $\Delta$ drives $\bar{A}=\exp (\Delta A)$ toward zero, so the current input has more influence on the state. A small $\Delta$ preserves more of the old state. The incoming token also contributes through $\bar{B}{x}_t$. Per token, the model adjusts how much information persists.

This slider changes one decay rate. Raise it and less of each past token survives; lower it and more survives, including noise. The curve cannot keep one token and forget another. In Mamba, the input-dependent state update learns what information helps predict future tokens.

Mamba-2 attacks the budget directly. Restricting the state transition to a scalar times identity makes the write cheap enough to afford $N\in \{64,128,256\}$, "up to 8x the size of Mamba or even higher, with minimal slowdown."

The Mamba row shows how far a small state can go on a particular synthetic task. The paper reports 99.8% induction-head accuracy after training on sequences of 256 tokens, with similar performance at lengths up to one million. The result does not mean the state preserves every token; it learns a compact representation useful for that task.

The results in this table come from different tasks, so they are not a head-to-head ranking. Each method makes a different trade between state size and update rule.

Test-time adaptation and training

Test-time adaptation can mean updating normalisation statistics, taking a gradient step on an auxiliary objective, or changing a persistent state as each input arrives. Those operations affect memory in different ways.

Tent is the lightest. It throws away the source normalisation statistics, re-estimates them from the test batch, and optimises only the channel-wise affine parameters $\gamma$ and $\beta$, which is under 1% of the model. The objective is the entropy of the model's own prediction:

No labels, no source data, nothing but the confidence of the prediction you were about to make anyway. On CIFAR-10-C at the highest severity it takes error from 40.8% to 14.3%.

TTT in 2020 went further. A genuine self-supervised loss, four-way rotation prediction, updates the shared feature extractor before each prediction. On CIFAR-10-C level 5 the online variant cuts Gaussian noise error from 50.5% to 25.8% and pixelation from 55.8% to 18.1%.

The split inside that paper matters more than the numbers. Standard TTT discards the update after each test sample. Online TTT keeps it, carrying parameters forward across the test stream. Which makes online TTT a memory: it accumulates information about the test distribution, with no compression budget and no forgetting gate, and no way to roll back.

The 2024 TTT layers make the inner training step part of each forward pass. Their headline comparison is worth quoting because it's the sharpest statement of what a bigger state buys: "Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context."

Tent is useful here because it shows that a model can adapt at inference using a label-free signal from its own predictions. RLVR also replaces labels with another signal, but uses a verifier rather than prediction entropy. The objectives are different; the shared idea is to update from a signal available at inference.

Reinforcement learning at test time

TTRL is the clearest case. Sample a group of answers, take the majority vote as the pseudo-label, reward each sample for agreeing with the consensus:

Then optimise with GRPO. The same objective, the same group-relative advantage, the same machinery from Part 1, running at inference against problems it has no labels for. Qwen2.5-Math-7B goes from 12.9 to 40.2 on AIME 2024, 35.6 to 68.1 on AMC, 46.7 to 83.4 on MATH-500.

The paper's headline comparison needs a second look. It says final avg@16 exceeds initial maj@16 by more than 20 points on AMC. In the figure it cites, avg@16 rises from about 28.5 to 46, while initial maj@16 is around 43. The difference is about three points. The 20-point figure appears to compare final avg@16 with its own initial score instead.

The real version is in the next figure over, at 64 samples: TTRL's avg@64 beats the initial maj@64 by 17.8 points on MATH-500, 7.3 on AMC, 2.3 on AIME. Still a model outperforming the vote that supervised it. Just never by twenty.

And there's a counterweight that deserves more attention than the headline.

In its per-problem analysis, TTRL-Guard found 44.5% of problems were already solvable, 0.7% were newly learned, and 21.6% were degraded from correct to incorrect. In that sample, degradation outnumbered learning about 31 to 1. Aggregate pass@1 hides these transitions because gains on already-solvable problems can outweigh losses elsewhere.

Their diagnosis names the mechanism and the culprit is familiar: "Due to the winner-takes-all nature of the standard GRPO loss, any step in which the incorrect answer wins the majority vote heavily penalises the correct answer." They call the window where a correct minority answer can still recover the correct-answer extinction window. On Llama-3.2-3B the correct-vote rate starts near 58% and collapses toward zero, and once it collapses it doesn't come back.

Two groups reached that diagnosis independently. Hi-TTRL attacks the same failure from the sampling side, finding that low consensus gets amplified into confident corruption.

TTRL-Guard has limits too. Its method loses on one of three models. The headline +54% is a relative gain on a 30-problem benchmark, about two or three questions, and the authors say it helps mainly when starting accuracy is between roughly 30% and 70%.

SEAL uses downstream performance as the reward for a model's self-written weight edits. On no-context SQuAD, its reported score rises from 32.7% to 47.0%, compared with 46.3% for the GPT-4.1 synthetic-data baseline.

The part of SEAL I find more useful is the part that failed. The reward depends on the updated parameters but isn't differentiable through them, and PPO and GRPO were both unstable, so they fell back to rejection sampling plus supervised fine-tuning. They also concede that "performance on earlier tasks gradually declines as the number of edits increases."

These methods reward generated answers, self-edits, or external memory actions. The narrower question here is whether a downstream task reward can directly train the internal retention gate in an SSM.

Writing to a notepad with a reward

This experiment is speculative. In the papers I checked, I found no work that optimises Mamba's $\Delta$ or Titans' ${\alpha }_t$ directly against a downstream task reward. There are nearby methods: REFINE uses GRPO to train fast-weight models for self-supervised next-sequence prediction, while GRU-Mem trains text-controlled memory gates with task rewards. Neither tests a task reward on a continuous SSM retention parameter.

What I built is a lookup table with three memory slots. It is not evidence about Mamba. It is a discrete version of the question.

The setup: items stream past, each tagged with a category. A notepad holds three of them. At every step the policy either skips the item or writes it over one slot. At the end, a query names a category and asks for the first item that appeared in it, so anything worth keeping had to be kept early and defended against everything that came after.

Same GRPO as Part 1, same group-relative advantage, same 16-sample groups. The only thing that changes between runs is the reward.

def recall_reward(notepad, episode):    """Did the notepad keep the answer to the query."""    return 1.0 if any(t == episode.answer and c == episode.query                      for t, c in notepad) else 0.0def recency_reward(notepad, episode):    """Did it keep the most recent items, which is what next-item    prediction rewards."""    recent = set(episode.items[-SLOTS:])    return sum(1.0 for entry in notepad if entry in recent) / SLOTS

Two functions, eight lines, and they're the entire difference between the two runs.

The two policies behave in opposite ways. The recency policy writes on almost every step, about nine times per episode, so the notepad always holds whatever came last. In the plotted run its write rate starts at 0.68 and climbs to 0.98. On the other seeds I tried, it stays near 0.97 from the first step. The recall policy always writes the first item, and by the last position its write rate has fallen to 0.03. It makes about three writes per episode, roughly one per slot, and then keeps what it has. It answers the query correctly 99.6% of the time, against 24.0% for the recency policy. Two more seeds gave about 99% against 22% to 23%.

Nothing about the architecture changed. The policy class, the state size, the optimiser and the group baseline are identical. The reward specified a retention policy, and the retention policy is what a memory is.

The recall policy is worth reading rather than just scoring, because the table is small enough to print. What it learned isn't subtle: fill each slot with the first item of a category you haven't seen yet, then refuse almost every subsequent write. Roughly three writes, one per slot, and then a long tail of skipping. It learned that the expensive mistake isn't failing to record something, it is recording something over the thing you already needed.

The recency policy learned the opposite behaviour. It writes nine times out of ten and ends with the last three things that happened in its notepad. It answers 24% of queries correctly, close to guessing in this three-category task. That resembles the bias of next-token prediction toward recent context, but this toy does not show how a language model's own gate responds.

The experiment has a hard limit: its write is a discrete action sampled from a policy, while Mamba's $\Delta$ is a continuous, differentiable part of the recurrence. Policy gradient can assign credit to the toy's sampled writes; that does not establish how to train Mamba's gate. The toy shows only that reward can train a discrete memory policy on this synthetic task.

Decisions and behaviour

Decision Transformer and Trajectory Transformer offer a useful contrast: one learns actions conditioned on desired return; the other adds search over candidate trajectories.

Decision Transformer treats a trajectory as a token stream:

Those $\hat{R}$ terms are returns-to-go. The model learns from offline trajectories with supervised loss, then conditions on a target return at evaluation. On Key-to-Door with 10,000 random trajectories, the paper reports 94.6%, against CQL's 13.3% and behavioural cloning's 1.6%.

Results vary by task. On Atari Qbert, Decision Transformer scores 15.4 ± 11.4, against CQL's 104.2. It is a useful approach on some datasets and a poor fit on others.

Trajectory Transformer draws the line that matters here. Decision Transformer is a conditional policy: condition on a desired return, emit an action. Trajectory Transformer is a model plus a planner, "repurposing beam search as a planning algorithm."

That difference brings calibration back into the picture. A planner searches for high-scoring trajectories under its model. If the scores are miscalibrated, search can favour trajectories the model overestimates. Calibration matters whenever downstream code optimises against a model's probabilities.

Behavioural prediction also learns from history, but does not condition on a target return. Accuracy alone misses whether confidence tracks correctness. In Part 1's probe, Jev answered the input-grounded question correctly every time, but on the policy-dependent question it was right 40% of the time at 84% mean confidence. Sixty of 120 answers were returned with probability below 0.1.

One observation here is mine. With binary rewards, GRPO's group mean estimates the prompt's pass rate across $G$ samples. The algorithm uses that mean as a baseline and normalises the resulting advantages; it does not expose the pass-rate estimate as a calibrated confidence.

The specific gap

The claim is narrower than "RL has never trained a memory gate." In the papers I checked, I found no method that treats Mamba's $\Delta$, Titans' ${\alpha }_t$ or DeltaNet's ${\beta }_t$ as an action trained directly against downstream task reward. I searched the arXiv API, Semantic Scholar and citation graphs for Titans and Gated DeltaNet. Those indexes cover titles and abstracts, so a paper that describes the method only in its body could be missing.

That scope matters. GRU-Mem uses task rewards to train text-controlled update and exit actions over external memory. Neural Garbage Collection uses task reward to train KV-cache eviction choices. REFINE applies GRPO to fast-weight models, but its reward is self-supervised next-sequence prediction. These are close relatives of the question, with different actions and objectives.

The idea itself is old. Peshkin, Meuleau and Kaelbling gave an agent actions to "set and clear bits in an external memory." Reinforcement Learning Neural Turing Machines later applied REINFORCE to discrete memory interfaces.

It has been tried and dropped. Hierarchical Memory Networks attempted REINFORCE for hard memory access in 2016, found it difficult, and switched to maximum inner product search.

Skip RNN learns a binary decision to skip a recurrent state update. It uses a straight-through estimator and a budget penalty rather than reinforcement learning. This is one example of a differentiable relaxation for training a discrete gate.

The memory rewriting benchmark finds classic recurrent models outperform structured memories on some memory-rewriting tasks, and calls for agents with trainable forgetting. That result points to the same problem, but does not test RL on an SSM's architectural gate.

Memory-R2 explains one difficulty: "memory turns the agent's past actions into part of its future environment. Once different rollouts write, update, or delete different memories, they no longer share the same intermediate memory state, making trajectory-level comparisons fundamentally unfair."

Read that against GRPO. The group baseline works because $G$ samples share a prompt, so their mean estimates a common difficulty you can subtract off. Give each sample its own memory and they no longer share an environment. The baseline stops being a baseline.

Where this leaves you

Ask what the state costs and what it forgets. Mamba ships 16 numbers per channel; Mamba-2 buys 256 because the write got cheaper. That number is your memory budget and nothing in the architecture will tell you what it chose to spend it on.

Track online test-time updates. Tent and online TTT carry updates forward across a stream. If the data distribution shifts, the model can carry earlier mistakes with it. Plan how to reset or roll back those updates.

Calibrate before you search. A planner maximising a miscalibrated score does not average the error out, it seeks the error out. Beam search over a confidently wrong model is a confidently wrong plan with more compute behind it.

Read SEAL's failure before applying RL to weight edits. Its policy-gradient runs were unstable when reward depended on updated parameters, so the authors used rejection sampling and supervised fine-tuning instead.

Check whether your RL run is teaching or sharpening. TTRL-Guard's split, 44.5% already-solvable against 0.7% learned and 21.6% degraded, is invisible in aggregate accuracy. If you're running anything self-supervised at test time, measure the per-problem transitions and not the mean.

Don't cite my notepad as evidence about Mamba. It has three slots and a lookup table, and it makes discrete write choices. A real SSM has a continuous state update.

Part 1 asked what the reward measures. This one ends with a narrower question: can an outcome reward train the update that decides what memory keeps? My toy cannot answer that for Mamba.

Sources

• Mamba: Gu and Dao, 2023, the selective state and the write gate
• Transformers are SSMs: Dao and Gu, 2024, Mamba-2 and the SSD duality
• Efficiently Modeling Long Sequences with Structured State Spaces: Gu, Goel and Ré, 2021, S4
• Learning to (Learn at Test Time): Sun et al., 2024, the hidden state as a model
• Test-Time Training with Self-Supervision: Sun et al., 2020, standard versus online
• Tent: Wang et al., 2021, entropy minimisation at test time
• Test-time regression: the unification behind the first section
• Titans: Behrouz et al., 2025, surprise, momentum and a forgetting gate
• ATLAS: Behrouz et al., 2025, the Omega rule
• TTRL: GRPO at test time with no labels
• TTRL-Guard: the correct-answer extinction window
• Hi-TTRL: the same failure diagnosed from the sampling side
• SEAL: reward over self-edits, and why policy gradient failed
• Memory-R2: why memory breaks a group baseline
• Neural Garbage Collection: RL over KV cache eviction, the nearest modern work
• REFINE: GRPO on DeltaNet, the same diagnosis about next-token prediction
• GRU-Mem: RL-trained update and exit gates at the agent loop
• Skip RNN: a differentiable relaxation for learning a discrete gate
• Memory Retention Is Not Enough to Master Memory Tasks in Reinforcement Learning: results on memory rewriting tasks
• Hierarchical Memory Networks: tried REINFORCE for memory access, abandoned it
• Learning Policies with External Memory: Peshkin et al., 1999
• Reinforcement Learning Neural Turing Machines: Zaremba and Sutskever, 2015
• Decision Transformer: Chen et al., 2021
• Offline RL as One Big Sequence Modeling Problem: Janner et al., 2021, Trajectory Transformer
• DeepSeekMath: Shao et al., 2024, the GRPO objective from Part 1

Companion code for the notepad experiment: github.com/Serendeep/rl-by-subtraction. rlsub/notepad.py is the toy, make charts regenerates the figure.