用线性规划复现MIT权重空间编程，8小时搜索变1分钟

Last week, MIT CSAIL published Sparks of Weight-Space Programming. Given the raw weights of a small network, GPT-6 Astra could tell which game it had been trained to play. Given only a game’s rules, it could write the weights of a small network from scratch and beat networks trained with PPO for 100 million steps. No gradients, no training loop, just reasoning about what the weights compute. They coined this weight-space programming.

In another experiment, they asked Astra to make Qwen3.5-0.8B better at math by editing its weights. Astra changed only 0.016% of the weights, and MATH-500 accuracy jumped from 7.9% to 48.1%.1

Looking at Astra’s edit, I realized most of its search was solving a linear program by trial and error. It spent about 8 hours and 28 forward passes choosing which weights to change and by how much. A couple of LP solves, about a minute each, do the same job with 8× fewer changed weights. In this post I’ll walk through the edit, why it’s an LP, and what that says about the fix itself.

Thinking loops#

Small reasoning models have a common failure: they get stuck in thinking loops, repeating the same few lines over and over instead of finishing their reasoning. Qwen3.5-0.8B is especially prone to it, and its model card says so. On MATH-500, 93% of the base model’s answers ran into the 32,768-token limit without ever reaching an answer.

The edit fixes the loops, so it helps anywhere the model loops. That includes instruction following: IFEval went from 21% to 62%.

On Gemma 3 270M, which rarely loops, Astra couldn’t meaningfully improve its math ability. As the authors put it, editing the weights could remove an obstacle hiding an ability, but couldn’t add one the model lacked.

How Astra fixed the loops#

Transformers use residual connections: each layer’s output is added to its input rather than replacing it. Follow one token through the network and you get a single vector that every layer adds to, called the residual stream (1,024-dimensional in this model). When the model is looping, the residual stream shifts in a consistent direction. Astra’s fix was to edit two layers so that what they add points the opposite way.

Finding the loop direction#

Astra wrote four pairs of short texts, each a piece of productive reasoning and a looping version of the same reasoning. It ran all eight through the model and recorded each token’s residual stream vector at one layer. Averaging over the looping texts gives one vector, and averaging over the productive texts gives another. Their difference, normalized, is the loop direction $u$: the way the residual stream moves when the model loops.

Editing the weights#

Each MLP ends with a matrix $W$ (called down_proj). At every token, it turns the MLP’s 3,584-dimensional hidden activation $h$ into a 1,024-dimensional output $Wh$, which gets added to the residual stream $x$:

So the MLP’s output is always a weighted sum of the same 3,584 vectors, the columns $w_j$ of $W$. $W$ fixes what those vectors are; the activations $h_j$ decide how much of each one gets added at each token.

Each column points partly along $u$. That amount, $u^\top w_j$, is how much column $j$ pushes the model toward looping whenever $h_j$ is active. Astra’s edit subtracts a multiple $\alpha$ of it:

At $\alpha = 1$ the component along $u$ is gone, so the column no longer pushes toward looping. Doing this to every column in layers 15 and 19 took accuracy from 9% to 33% on GSM8K and MATH training problems. At $\alpha = 6$ the component is flipped: a column that wrote $+a$ along $u$ now writes $-5a$, and pushes the model away from looping whenever it’s used.

Every edited column changes by a multiple of the same vector $u$. So the whole edit is one direction $u$ plus one number per column, $v_j = -\alpha\,(u^\top w_j)$, which is zero for columns left alone. That’s a rank-1 change:

At each token, it adds $(v \cdot h_t)\,u$ to the residual stream: some amount of $u$, set by how active the edited columns are. No other direction changes.

That leaves $v$ to choose: which columns to edit, and how large $\alpha$ should be. Astra worked it out by trial and error, testing each edit on a fresh batch of those training problems:

Pushing harder than $\alpha = 1$ helped more, and a random direction did nothing. The final edit used $\alpha = 6$ on only the 64 columns per layer that wrote the most $u$ on looping text.

Where the time went#

The paper released Astra’s full logs: every tool call, every forward pass, every note it wrote to itself. Astra had a Python kernel with the weights loaded and a budget of 50 forward passes, each a batched run of the model on training problems or texts of its choosing.

The edit was frozen after about 10 hours, and the last 17 passes re-tested it without changing anything.2 Of those 10 hours, finding $u$ took 15 minutes. Finding $v$ took about 8 hours.

Choosing v is a linear program#

Astra was searching for a $v$ that strongly pushes looping tokens away from $u$ while barely touching productive ones. Since the push at token $t$ is $v \cdot h_t$, linear in $v$, both requirements are linear inequalities.

I labeled tokens from the base model’s own traces: looping tokens ($R$), deep inside a loop, and productive tokens ($P$), well before any loop starts,3 and recorded their activations $h_t$ at layers 15 and 19. Each one becomes a constraint, and the LP picks the $v$ with the smallest L1 norm $\sum_j |v_j|$:

$$ \begin{aligned} \min_v \quad & \sum_j |v_j| \\ \text{s.t.} \quad & v \cdot h_t \le -\delta && t \in R \quad \text{(at least } \delta \text{ away from } u\text{)} \\ & |v \cdot h_t| \le \varepsilon && t \in P \quad \text{(at most } \varepsilon\text{)} \end{aligned} $$

In practice no $v$ satisfies every constraint, so each one gets a penalized slack variable.

Minimizing L1 tends to make $v$ sparse. The L1 ball is a diamond, and when you grow it until it touches the feasible region, it usually touches at a corner, where most coordinates are exactly zero. So the LP picks the columns as well as the amounts.

To get exactly 8 columns per layer, I kept the 8 with the largest effect and solved again with the rest fixed at zero.4

That replaces Astra’s 8-hour search, including its choice of columns and $\alpha$, with a couple of LP solves. An off-the-shelf solver (HiGHS) took over 20 minutes each on these constraints, so I wrote a GPU solver (PDHG, the algorithm behind PDLP) that gets within 0.2% of its optimum in about a minute.5

Results#

I tested two LP edits. Both use Astra’s layers. One uses Astra’s direction $u$. The other uses a heuristic direction: label tokens in the model’s own traces as looping or productive with a simple repetition detector, then take the difference of their average residual streams. It’s the same recipe Astra used, minus the hand-written texts, so no LLM is involved. I tuned the LP’s settings on a separate 128-problem dev set; everything in the chart and table is on data never used for fitting or tuning.

With Astra’s direction, the LP matches Astra’s edit on every benchmark, within noise, while changing 8× fewer weights. With the heuristic direction, it gets about 93% of the way there.

A few other things showed up: