Tiny CLM: Teaching a 1.7B Model to Manage Its Own Memory for $5

Published on 02 October 2026
11 min read
context-language-models
context-engineering
Reinforcement-learning
Tiny CLM: Teaching a 1.7B Model to Manage Its Own Memory for $5

Tiny CLM: Teaching a 1.7B Model to Manage Its Own Memory for $5

A hands-on companion to my deep review of Context Language Models (Shao et al., 2026). Code on GitLab, all 94 trained adapters on Hugging Face.


Most paper reviews stop at a summary. This one runs an experiment.

Context Language Models (CLMs) argues that a language model can learn, through reinforcement learning, to edit its own memory while it works, and that a reward measured in compute cost makes it do so efficiently. I took that claim to hardware I control: Qwen3-1.7B, one consumer GPU and $4.65 of compute. I also tested a prediction the paper never checked.

Want the paper itself? My section-by-section review reads beside the actual PDF: as you scroll, the paper follows and highlights what each paragraph discusses, with the math worked through and the places I push back. This post is the short, hands-on version.

Table of Contents

  1. The Idea in One Minute
  2. The Game: a Clerk, a Desk and a Filing Cabinet
  3. How the Model Learned
  4. What Happened
  5. The Surprise
  6. Three Smaller Findings
  7. Try It Yourself
  8. Keep Reading
  9. Caveats

The Idea in One Minute

A normal model’s memory, its context, only grows. Every message is appended to the end, and when it fills up, the model is stuck. Harnesses patch this from the outside by compacting the transcript on a schedule; I wrote about that division of labor in The Leash and the Wolf.

A CLM instead edits its own context: it moves things out, writes notes and deletes what it no longer needs. The paper shows this can be taught with RL.

There’s a catch, and it’s about money. Serving engines like SGLang cache the work they did on a prompt and reuse it whenever the next prompt starts with the same text (I walk through SGLang’s radix-tree version in my SGLang deep dive). So:

  • Edit at the end of the context: almost everything is reused. Cheap.
  • Edit near the top: everything after the edit must be re-read. Expensive.

The paper measures this with prefix-reuse FLOPs and rewards the model for keeping it low. My prediction, derived from the paper’s own cost model (Sentence 2b in the review), was that this reward should teach the model to edit at the tail, where the cache survives.

The model did learn to edit at the tail. It just didn’t need the efficiency reward to get there, and why is the most interesting result.


The Game: a Clerk, a Desk and a Filing Cabinet

I built a small, seeded game, a Mini KV Store, scaled down from the KV Store task in the paper’s ContextBench. Picture a clerk at a small desk. The desk is the model’s context: 2,048 tokens. A messenger brings sheets, each listing 16 facts like:

Locker K12007 contains: amber river lotus quill slate pine ember tide #1a2b3c4d

Then come 4 questions (“What’s in locker K12007?”), answered by exact match. Next to the desk is a filing cabinet, and the clerk has five moves:

Move What it does
OFFLOAD B003 File sheet 3 in the cabinet, leaving a one-line sticky note on the desk
GREP K12007 Search the cabinet and copy that one line back
NOTE ... Write a short note
ANSWER K12007 ... Answer a question
READY Bring the next sheet

If the desk overflows, the clerk is fired on the spot: the episode ends as a failure. Pressure is how much paper arrives compared with the desk: at 2×, two desks’ worth flows through, so filing is mandatory.

Two strategies both score 4 out of 4. Clerk A files every sheet the moment it arrives. Clerk B lets sheets pile up and files the oldest ones when the desk is nearly full. Same answers, but Clerk B pulls sheets from under the pile, so everything above them has to be re-read:

Line chart: prefix-reuse TFLOPs per episode against context pressure. Clerk B rises to 43.6 at 4x, Clerk A to 29.8; at 2x it is 24.8 against 15.1.

That 1.5–1.6× gap at the same accuracy is exactly what the paper’s efficiency reward is supposed to exploit.


How the Model Learned

The model. Qwen3-1.7B, a small dense model from the same family the paper used, trained with LoRA: the 1.7 billion original weights stay frozen and about 35 million new ones (roughly 2%) learn corrections on top. Each trained version is a 130 MB adapter instead of a 3.4 GB model.

Step 1, untrained. Given only the rules, the model never filed anything. It kept asking for the next sheet and overflowed every game at 1× pressure and above.

Step 2, an apprentice on purpose. I showed it Clerk B’s moves (supervised fine-tuning) and stopped early, at a checkpoint with about 50% success, picked by a rule written down in advance. Two deliberate choices: imitate the expensive clerk, so RL has a bad habit to unlearn, and stop before perfection, because RL learns nothing from a model that always does the same thing.

Step 3, reinforcement learning. The paper’s method, GRPO: play the same game 8 times, compare each try with the group’s average, and nudge the model toward the above-average tries. If all 8 tries score the same, there’s nothing to learn from that group. (The policy-gradients-to-GRPO refresher in the review builds this up from scratch, and my TorchRL post covers the “LLM as an RL agent” framing.) I compared three ways to keep score:

  • R1, outcome only: correct answers count, nothing else.
  • R2, the paper’s method: outcome plus an efficiency bonus that only applies among attempts that got everything right (the paper’s Eq. 6, its “success gate”).
  • R3, a naive cost penalty: outcome minus cost for every attempt. The paper warns this invites cheating, like failing early because failing is cheap.

Training was stepwise, as in the paper: every reply is trained against the exact desk the model saw at that moment, because filing rewrites earlier parts of the context. I also ran the naive version, trained on the final transcript, for comparison. The review’s stepwise GRPO section writes out the math.

Plateaus and lucky breaks. The apprentice is sloppy but consistent: at high pressure, all 8 tries often fail the same way, so there’s no signal. A run escapes only when one try gets lucky (files early, survives) and the difference starts a snowball. With groups of 8 (the paper uses 32), escaping is partly luck, so every setup ran on 3 seeds.


What Happened

Runs that escaped became perfect clerks, 100% at every pressure, including 4×, which they never trained on:

Line chart: held-out accuracy against pressure. The untrained model drops to 0% from 1x, the SFT apprentice falls from about 87% to 0% at 4x, and the RL model stays near 100% through 3x and ends at 67% pooled over three seeds.

Here’s what an escaped model actually does, on a real test game at 2× pressure it never saw. Both models score 4/4:

After RL (R1), 15.1 TFLOPs             The apprentice, 24.7 TFLOPs
sheet B000  → OFFLOAD B000   2.42      sheet B000  → READY            2.40
sheet B001  → OFFLOAD B001   1.70      sheet B001  → OFFLOAD B000     1.63
sheet B002  → OFFLOAD B002   1.68      sheet B002  → OFFLOAD B001     3.29
...                         ~1.70      ...                           ~3.30

The trained model files each sheet the moment it arrives. The apprentice always keeps the newest sheet on the desk and files the one under it, so every filing turn costs about twice as much. Nobody told the trained model to do that, and its reward, R1, contained no cost term at all.

The four questions I set out to answer:

Question Answer
Can a small model learn context management from outcome reward alone? Yes. From 0% (untrained) and 50% (apprentice) to 100% at every pressure in escaped runs
Does the success gate matter? Not here, because nothing got hacked. The naive penalty R3 matched R1 on accuracy (0.853 vs 0.854) and cost, with no cheating patterns
Does the efficiency reward move edits to the tail? No, for a structural reason. Every escaped run, under every reward, filed at the tail
Does stepwise training matter? Barely measurable in this game. One naive run did develop a failure loop no stepwise run showed

The third answer is the headline, and this chart is the evidence. Where a model edits depends on whether its run escaped, not on which reward it got:

Bar chart: share of edits that forced no unchanged text to be re-read. SFT apprentice 0%, R2 40% (escaped on 1 of 3 seeds), R1 72% and R3 72% (2 of 3), R1 on final transcript 97% (2 of 3 plus 1 partial).

R2 looks worse than R1 only because it escaped on one seed instead of two. With three seeds, that gap is indistinguishable from chance, so I’m not claiming the paper’s reward hurts.


The Surprise

In this game, filing immediately isn’t only the cheapest strategy. It’s also the safest: a near-empty desk never overflows. So the plain “get the answers right” reward already selects the cache-friendly habit, and the efficiency bonus has nothing left to teach.

That gives a precise condition for when the paper’s efficiency reward can matter: only when cheap behavior and accurate behavior disagree. Think of a task where keeping material in context helps accuracy but costs re-reading. The paper’s deep-research setting may well have that tension; a KV store with a small desk doesn’t. It’s a refinement of the paper’s claim, and it came from running the experiment rather than reading the paper. I dig into the efficiency advantage itself in the review’s Sentence 7.


Three Smaller Findings

  1. A better metric for cache-friendly editing. Edit depth, where the first changed token sits, scores a perfect tail edit as “halfway down” when the edited block is large. I added the unchanged-suffix fraction: the share of the previous context after the edit that had to be re-read although nothing in it changed. It’s exactly 0 for tail edits, and it’s what the bar chart above measures.
  2. Sticky notes aren’t free. Each filed sheet leaves about 45 tokens behind, the pointer plus the turn that filed it. On a 2,048-token desk, those alone cap the playable pressure near 5× for any strategy. It’s a miniature version of the residue problem behind compaction, which my context-pruning work attacked from the outside.
  3. Cheap RL is a dice roll. Seed 0 alone said “cost-aware rewards get stuck”; seeds 1 and 2 reversed it. If you reproduce RL papers on a budget, run several seeds before believing a difference.

Try It Yourself

The project was pre-registered: hypotheses, metrics and the rules for “supported” or “refuted” were committed to git before the main runs, and the refuted ones are reported as refuted. It’s about 1,800 lines of Python with 49 tests, seeded episodes, and every number regenerated from raw results by a script. Everything ran on one 24 GB GPU in about 9 GPU-hours.

  • Code: gitlab.com/sugix/tiny-clm. Start with env/kvstore.py (the game), cost/prefix_reuse_flops.py (the paper’s cost model, tested against its published numbers) and rl/grpo_stepwise.py (the trainer).
  • Weights: all 94 adapters on Hugging Face. The escaped “expert clerks” are e3_r1_s0, e3_r1_s2, e3_r2_s2, e3_r3_s1, e3_r3_s2, e5_final_s0 and e5_final_s1; use the ckpt_best/ folder inside each.
from transformers import AutoModelForCausalLM
from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-1.7B", torch_dtype="bfloat16")
model = PeftModel.from_pretrained(base, "sugiv/tiny-clm-qwen3-1.7b-lora",
                                  subfolder="e3_r1_s0/ckpt_best")

To make it play, build prompts with the project’s game engine, as the evaluation code does.


Keep Reading

This experiment sits on top of several things I’ve written. If one of them caught your eye above, here’s where to go next:


Caveats

  • One small model, one synthetic task. The results say nothing directly about 9B or 27B models on deep research.
  • Small groups and short training (8 tries instead of 32, 40 steps instead of 70) make escaping a matter of luck. Three seeds show that, but can’t rank the rewards.
  • A restricted action space. The only edit swaps a sheet for a one-line pointer, which barely separates stepwise from naive training. The paper’s free-form Python edits would be a sterner test.
  • Absolute compute isn’t comparable with the paper’s: different model, task and scale.

The model learned to manage its memory, and it learned to do it cheaply. It just did the second part for a different reason than the paper predicted, and that reason is the most useful thing a $5 experiment told me.

Paper: Shao et al., “Context Language Models,” arXiv 2609.37725 (2026), CC BY 4.0.