TL;DR
We took a decision skill that lived in a separate model and wrote it straight into the weights of a small language model. No training loop and no gradient descent: one linear solve, 3 milliseconds, 262 thousand weights in one layer. Then we threw the separate model away.
- Qwen3.5-0.8B, on its own, now decides whether a request is ready to act on: right 97.8% of the time, up from 18%. Tool choice went from 80% to 93.5%.
- No catastrophic forgetting: capitals and ARC science questions unchanged, ordinary text 10% harder to predict, 1 of 15 unrelated prompts affected.
- Against LoRA on the same data: about 3 times faster end to end and a thousand times fewer weights changed. LoRA reached 100% but started answering "suggest a cat name" with "no".
- The catch: the edit fires reliably only on the prompt format it was solved on.
The skill
An agent has to make two calls before it does anything with a request: which tool it needs (search, the caller's records, a calculation, an action, or none), and whether it has enough information to go ahead. We used 393 requests labelled with both. Asked plainly, Qwen3.5-0.8B gets the tool right 80% of the time and the second call right 18% of the time, because it says "no" to almost everything.
The teacher is a vegaml policy fitted on 300 of those rows: one small head per question on top of a frozen reader, then calibrated, so its probabilities mean what they say. On the 93 rows it never saw it scores 96.8% on the tool and 97.8% on ready. It is a good decision model. It is also a separate thing that has to sit next to the language model. The question was whether the language model can carry the decision by itself.
The edit
Think of each answer as a valley the model's output rolls into. The teacher's valleys are in the right places and the plain model's are not. Rather than pushing the model downhill one small gradient step at a time, we work out where the valleys should be and move the ground once.
The ground we move is the output matrix of the MLP in layer 23 of 24, and only 256 of its 3,584 input columns: the ones most active at the answer position compared with how active they are on everything else. For each training request we record what those columns see at the answer position, , and how the final hidden state would have to change for the model's answer to match the teacher's distribution, . That change comes from the gap between the two distributions in log space, pulled back through the model's readout. The edit is the least-squares map from one to the other, with a penalty for touching anything else:
is what those same columns see on text we want left alone: ordinary sentences, the rest of every prompt, and the position right after the answer, so the model does not learn to keep answering. sets how hard that text is protected. This is the normal-equations form that MEMIT made standard for editing facts; here the target is a calibrated decision instead of a fact.
Two guards keep it honest. We scale the edit by a dose of 0.5, 1, 1.5 or 2 and keep the one that best closes the gap on 40 rows held back from the solve, and any dose that moves the protected text by more than a KL of 0.03 is refused outright. Here the first dose won, in one round. The solve and the write took 3.1 milliseconds; the whole run, including reading the rows through the model and trying the doses, 87 seconds on two T4s. Written into the fp16 weights with the teacher removed, the model scores the same as with the edit applied on the fly.
What it learned
Scored by letting the edited model write its own answer to the 93 held-out requests, against LoRA (rank 16 on every MLP, 60 steps, the same rows and the same protected text):
| Unedited | Direct edit | LoRA | |
|---|---|---|---|
| Tool choice | 80% | 93.5% | 100% |
| Ready to act | 18% | 97.8% | 100% |
| Time on 2x T4 | n/a | 87 s (edit: 3 ms) | 252 s |
| Weights changed | n/a | 262 thousand | 263 million |
| WikiText-2 perplexity | 28.7 | 31.7 | 33.1 |
| Held-out text perplexity | 8.11 | 8.50 | 8.94 |
| ARC-Easy (200) | 70.5% | 71.5% | 78.5% |
| Capitals (40) | 100% | 100% | 100% |
| Unrelated prompts affected | 0 of 15 | 1 of 15 | about 5 of 15 |
LoRA learns more. It is also the one that leaks: asked for cat names, a two-day Kerala plan or whether raw eggs are safe, it replies "no", "yes", "yes", and after naming a tool it repeats the name instead of carrying on. The direct edit answers and then continues the way the original model did. Its one slip was listing the primary colours as "red, no, no, no". LoRA's ARC gain is most likely a format effect, since it learned to answer in a single word.
What the ablations say
| Change | Tool | Ready | WikiText ppl |
|---|---|---|---|
| Full method | 93.5% | 97.8% | 31.7 |
| No protection (λ = 0, no cap) | 96.8% | 97.8% | 1.2e11 |
| Hard labels, not the calibrated policy | 95.7% | 97.8% | 30.4 |
| 25% of the training rows | 95.7% | 95.7% | 32.7 |
| Random change of the same size | 74.2% | 81.7% | 29.2 |
- Protection is the method. Without the term and the cap, the same solve hits the targets and destroys the model.
- The calibrated teacher matters for confidence, not accuracy. Hard labels give the same answers, with almost twice the calibration error on the ready question (0.25 against 0.14).
- A quarter of the rows is enough for accuracy, at a little more drift.
- A random change is not harmless. It moves ready to 82% only because the model is biased towards "no", and it drops ARC from 70.5% to 52.5%. The solved edit leaves ARC alone.
What it does not do yet
- It is tied to the prompt format. Ask the same question in a new wording and the edited model gets 5 of 11 right, where the unedited one got 8. LoRA gets 9. The fix to try is solving over several formats at once.
- One task, one small model, small test sets. Stacking many edits is where earlier work saw forgetting build up, and we have not tested that.
- It suits decisions, not behaviours. We tried the same kind of edit on reasoning loops in long maths traces and it did not move them.
Where this sits
Closed-form weight edits are not new. MEMIT solves this form to insert facts, weightsmith uses it to bake a prompt's effect into a robot policy, and training-free on-policy distillation uses the same second moments to copy a teacher with forward passes only. What we did not find is this combination: a calibrated decision policy as the teacher, solved at the answer token, checked in the model's own free generation, and measured against LoRA for forgetting.
Why it matters to us: a fitted vegaml policy could ship as a plain model file with no extra parts, and be updated in milliseconds when new rows arrive instead of retrained. That is worth chasing once the edit holds across prompt formats.