stage3-surrogate-task

Stage 3 · Pre-validate · Surrogate task

A targeted synthetic task that the model cannot solve without the new feature, but can solve with it. Used to validate the mechanism of an architectural change, not just the loss number.

Stage question

"Does my new component actually do what I claim it does, in isolation?"

A surrogate task gives a yes/no with a clear mechanistic interpretation. Loss numbers on real data don't.

When to use

• A new attention variant claims to handle long-range dependencies → make a copying or retrieval task at length L.
• A new positional encoding claims to extrapolate beyond train length → train at length L, test at length 4L.
• A new memory module claims to remember across boundaries → test at boundaries.
• A new routing scheme claims to specialize experts → test that experts cluster by input class.

In each case, the surrogate is smaller than the real task, more controlled, and diagnostic — failure on it is decisive.

Designing a good surrogate task

A useful surrogate satisfies all four:

1. Trivial without the feature: the baseline architecture, without the new component, demonstrably fails (or solves with much higher loss).
2. Solvable with the feature: the variant architecture solves it within the time budget the user can afford.
3. Cheap: trains in minutes, not hours.
4. Mechanistically tied to the claim: success on the surrogate is direct evidence of the mechanism, not a side effect.

If the surrogate is solved by both arms, it's not diagnostic — design a harder one. If neither arm solves it, the surrogate is wrong (too hard, or the feature implementation is wrong).

Examples

Claim	Surrogate task
Long-context attention works	Copying: input is <random tokens><DELIM><same random tokens>; the model must copy. Length L pinned.
Positional encoding extrapolates	Train on length L sequences; test on length 4L. Metric: accuracy at extrapolated length.
MoE experts specialize	Multi-task synthetic: 4 different deterministic mappings, route by class. Confirm experts cluster.
Recurrent state preserves info	Long-range memorize-then-retrieve task with arbitrary delay.
New optimizer escapes saddle	Train a known-saddle objective (e.g., 2-layer linear net with synthetic data); compare convergence.
Spiking time encoding helps	Time-varying signal classification at different rates; SNN T-dim used to encode rate.

Procedure

1. Articulate the claim in one sentence: "Component C makes property P hold". P should be testable with a synthetic task.

2. Design the surrogate so that property P is the bottleneck. Strip everything else (small model, small vocab, deterministic data).

3. Run baseline (architecture without C). Confirm it fails (or has high loss). This step is non-optional — without it, success on the surrogate doesn't tell you the feature was needed.

4. Run variant (architecture with C). Confirm it succeeds.

5. Vary the difficulty of the surrogate (e.g. sequence length for retrieval). Plot success rate. The variant's success curve should extend further than the baseline's.

6. Document the surrogate in the run journal alongside the small-scale ablation result. A surrogate-task pass does not replace a small-scale ablation — they are complementary.

Procedure when assisting a user

1. Press the user to articulate the mechanistic claim in one sentence. If they cannot, the work is at the "I think this might help" stage and the surrogate-task work is premature — go back to small-scale ablation first.

2. Help design the surrogate by working backward from the claim. The right surrogate is usually obvious once the claim is sharp.

3. Confirm the surrogate is cheap. If it takes more than ~30 minutes to run, redesign.

4. Confirm the baseline fails. Without this control, the result is uninterpretable.

5. After running, render the result as one line: "Variant solves surrogate at length L=4096 with accuracy 0.94; baseline 0.21. Mechanistic claim supported." Or: "Variant fails surrogate at L=4096 (0.34); claim not supported at this scale."

Boundaries

• A surrogate task is mechanistic evidence, not a generalization claim. A model that solves the surrogate may still fail on the real task for other reasons.
• The surrogate is not a substitute for stage3-small-scale-ablation on real data. They test different things.
• Designing a good surrogate is creative work. Bad surrogates produce false positives (variant wins on a task that didn't actually require the feature).

Common mistakes

• Surrogate too easy → both arms solve it; no signal.
• Surrogate too hard → both arms fail; no signal.
• Surrogate measures the wrong thing → "longer-context attention helps" with a task that actually requires more parameters.
• Skipping the baseline → can't tell if the feature is doing the work.

Related

• skills/stage3-small-scale-ablation — measures aggregate improvement on real data; the surrogate isolates mechanism.
• skills/stage3-multi-seed-variance — apply to surrogate too; one-seed surrogates are noisy.
• Anthropic's "circuits" research and DeepMind's "TransformerLens" tutorials are good sources of surrogate-task ideas for transformer mechanisms.