embedding

Blog/·6 min

A decision skill written into the weights in 3 milliseconds

We wrote a calibrated vegaml decision policy straight into Qwen3.5-0.8B's weights with one closed-form solve, no gradient descent, then removed the policy. Ready-to-act accuracy went from 18% to 97.8% with no catastrophic forgetting. Against LoRA: 3 times faster, a thousand times fewer weights changed, and no leaking into ordinary chat, though LoRA learns more and the edit only holds in the prompt format it was solved on.


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.

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:

(1)

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):

UneditedDirect editLoRA
Tool choice80%93.5%100%
Ready to act18%97.8%100%
Time on 2x T4n/a87 s (edit: 3 ms)252 s
Weights changedn/a262 thousand263 million
WikiText-2 perplexity28.731.733.1
Held-out text perplexity8.118.508.94
ARC-Easy (200)70.5%71.5%78.5%
Capitals (40)100%100%100%
Unrelated prompts affected0 of 151 of 15about 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

ChangeToolReadyWikiText ppl
Full method93.5%97.8%31.7
No protection (λ = 0, no cap)96.8%97.8%1.2e11
Hard labels, not the calibrated policy95.7%97.8%30.4
25% of the training rows95.7%95.7%32.7
Random change of the same size74.2%81.7%29.2

What it does not do yet

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.


More posts·embedding.in