A looped-layer experiment with two checks to kill it, and a training stage that can't learn

Flow

What each check should leave behind

The gradient check is the one this commit does not have yet.

1. 1Hand-run match⌄

   Leaves: pass or fail per prompt. Decides: whether any later number can be trusted.

2. 2Donor scores⌄

   Leaves: state and update scores for every candidate. Decides: which layer to copy, or whether to stop.

3. 3Config check⌄

   Leaves: the donor the model actually loaded. Decides: whether the scores were used at all.

4. 4One backward step⌄

   Leaves: the copied layer's gradient size. Decides: whether training can start.

5. 5Held-out loss⌄

   Leaves: a loss curve per stage. Decides: keep going or drop the idea.

Flow

Checks in order before a long run

Each check is cheaper than the run it protects.

1. 11. Hand-run match⌄

   Layers called one at a time match the normal model output.

2. 22. Per-prompt donor test⌄

   Score each candidate on one prompt at a time so attention never mixes prompts.

3. 33. Rank on updates⌄

   Compare the change each layer makes, not only the state it ends with.

4. 44. Wire the pick in⌄

   Make sure the chosen donor is the one the model loads.

5. 55. Gradient check⌄

   One backward step through the full model shows a gradient on the copied layer.

6. 66. Train⌄

   Stage one on cached states, then stage two on outputs.

The last time I tried swapping the middle of my language model for something cheaper, the replacement matched it on the prompts it was fitted to and fell short on new ones. This time the commit loops one copied layer in place of the middle, and adds two checks meant to rule the idea out before training.

One copied layer in place of the middle

I keep the early and late layers frozen. I remove a middle stretch and run one copied layer once per layer I removed. Each pass gets a small step vector and thin corrections that start at zero, so the starting model is just the copied layer repeated. Every pass uses the same positions and layer arguments. What I save is weights, not compute.

Checks before training

The first runs the layers by hand and confirms they give the same output as the normal model. If they don't, my harness is passing the wrong arguments and nothing after it counts.

The second hands each candidate layer the real inputs of the others and measures how close it gets to their real outputs. The best one is supposed to become the donor.

Where my checks were weaker than they looked

My script glues prompts into long chunks, so attention can mix them. I should score each prompt on its own. It also ranks donors on the final state only. I'd rather rank them on the change a layer makes, since a layer that does almost nothing can still look good on state. And the chosen donor never reaches my config, which has a fixed default.

The gradient trap

The commit treats frozen as if it meant no gradient. When a frozen layer sits after the part I'm training, it needs to be called normally with its weights frozen. A no-gradient block around it stops the loss from reaching anything before it.

The tail should be called normally, not inside the no-gradient block.

That's what I think the second training stage gets wrong, so as written it can't train the middle. Before any long run, I want one backward step and a check that the trained layer's gradient is not zero. I'd add that step to the first check script.