feat: resample ablation for self-repair (#35)

[xiaohan2012]

Implements #35 — resample ablation. Full design discussion: issue #35.

What & why

• Self-repair uses zero/skip ablation ($\delta := 0$). Off-distribution: it collapses the residual norm → the model sees a state it never sees → the measured "repair" may be an artifact of the broken input, not real compensation.
• Resample swaps $\delta$ for a role-matched real-donor $\delta$ (upstream intact, downstream free to react). On-manifold + norm-preserving → recovered signal is genuine compensation.
• zero kept as a second mode for the skip-vs-resample cross-check.

Method

Contribution $\delta$ (what gets swapped).

• Layer $m$: $\delta_m = \mathrm{res}(\mathrm{out}_m) - \mathrm{res}(\mathrm{in}_m)$ — layer output − input, from the existing per-layer capture (no new hook).
• New residual after $m$: $u + \delta^{\text{don}}$, where $u = \mathrm{res}(\mathrm{in}_m)$ (upstream) → upstream intact, downstream free to react (= total effect). Swap the contribution, not the whole residual (that would wipe upstream).
• residual_of(layer_output) picks the stream (no token-slice, no norm).

Role matching (where a donor $\delta$ may come from). Same-role positions only:

• Row axis (all models): context←context, query←query.
• Token axis (4D — LimiX/Mitra): label←label (decoder reads only the label token), feature←feature random. LimiX label = last token · Mitra label = index 0.
• 3D (TabICL/TabFM): features fused pre-ICL → row axis only.

Donor source = leave-one-out real tables. Pool = the other loaded TabArena tasks (target excluded). Model-agnostic + matches the eval distribution → cleanest norm match, no new data pipeline.

$N=8$, average the effect (not the $\delta$). 8 donors without replacement, fixed seed → 8 forwards → average the per-depth metric (+ IQR band). Averaging the $\delta$ instead is itself an off-manifold mean ablation: $\mathbb{E}[f(\delta)] \neq f(\mathbb{E}[\delta])$.

Changes (module by module)

• adapters/ — residual_of (pick stream, no slice/norm), resample_forward (input + donor_δ in the layer's tuple shape), label_token_index (LimiX -1, Mitra 0, None for 3D).
• core/donor.py — donor_deltas (one table's per-layer δ) + build_donor_delta (target-shaped δ by role-matched random draw).
• core/interventions.py — resample_layer (mirrors skip_layer; δ:=0 == skip).
• evaluation/self_repair.py — ablation_sweep(ablation={zero,resample}, donor_tables, n_donors, seed); "skip" key unchanged so the plots read either mode.
• scripts/run_self_repair_sweep.py — --ablation {zero,resample} --donor {table,row,feature} (table only) --n-donors; loads all tasks up front for the LOO donor pool.

Tests

• adapters — round-trip invariant: re-applying a layer's own δ reproduces its output (single / 4D 3-tuple / double-stream).
• donor — role integrity by tagging each bucket; incl. the hard constraint that a feature δ never lands on the label token; reproducibility; double-stream + 4D-double-stream (Mitra-like).
• interventions — δ:=0 == skip, output == input + δ, tuple/stream shapes, restore-on-exit.
• sweep — resample shape + IQR keys, reproducible with seed, raises without donors / on unknown mode.
• 82 fast tests pass; ruff clean. Real-model adapter suites unchanged (need checkpoints).

Validation (done)

• V1 — donor $\delta$ health: resample is on-manifold + norm-preserving; zero collapses the residual norm.
• V2 — skip-vs-resample cross-check: on the decision margin, self-repair holds under resample → real model behavior, not a zero artifact; resample also perturbs more stably.

Full redundant-vs-repair still needs DE (#37); this swaps the ablation method only.

🤖 Generated with Claude Code

[xiaohan2012]

V1 — donor $\delta$ health check (results)

Is resample on-manifold + norm-preserving? Validated on 4 real models (LimiX, Mitra, TabICL, TabFM); zero ablation is the reference.

Updated: re-run at 1000/500 (was ~500/250) — both claims unchanged.

Terms.

• donor — a different real table whose contribution replaces the ablated one. Pool = the other 14 tables (leave-one-out); $N=8$ drawn per target, metric averaged.
• contribution $\delta$ — what a layer writes to the residual at one position. native $\delta^{\text{nat}}$ = the target's own; donor $\delta^{\text{don}}$ = the role-matched swap-in (context←context, query←query; label←label, feature←feature).
• decoded position — where the decoder reads: label token (4D: LimiX/Mitra) or row vector (3D: TabICL/TabFM). All metrics measured here.
• $u$ — upstream residual $\mathrm{res}(\mathrm{in}_m)$ (everything accumulated before this layer).

Setup (formal). Layer $m$, test rows $i=1..N$, at the decoded position. With $u_i=\mathrm{res}(\mathrm{in}_m)_i$: $\ \delta^{\text{nat}}_i=\mathrm{res}(\mathrm{out}_m)_i-u_i$ (output − input); $\delta^{\text{don}}_i$ = role-matched donor $\delta$ at row $i$.

Two claims, one per axis.

---

Claim 1 · Magnitude — resample preserves the residual norm

Why:

• Zero-ablation shrinks the residual to a magnitude the model never sees → off-distribution.
• A faithful ablation should keep the residual realistic — magnitude included.
• Test: does resample preserve the magnitude where zero drops it?

Norm-preservation ratio $r_{\text{mode}}(m)$ — the fraction of the clean residual norm that survives the ablation:

$$r_{\text{mode}}(m)=\mathrm{median}_i\frac{\lVert u_i+\delta^{\text{mode}}_i\rVert}{\lVert u_i+\delta^{\text{nat}}_i\rVert},\qquad \delta^{\text{zero}}=0,\ \ \delta^{\text{resample}}=\delta^{\text{don}}$$

Reads as: how much of the residual survives after removing this layer's $\delta$.

• zero nulls $\delta$ → residual collapses (off-distribution).
• resample swaps a real $\delta$ → norm holds.
• $y=1$ = equal norm, not equal vectors.

Lines = median over 15 tables (8 donors averaged per table); shaded = IQR + min–max across tables.

model	$\min_m r_{\text{resample}}(m)$	$\min_m r_{\text{zero}}(m)$	zero's worst-hit layers
LimiX	1.00	0.14	first layer
Mitra	1.00	0.52	first / last
TabICL	0.96	0.64	mid–late
TabFM	0.99	0.66	early

Takeaway: resample hugs 1 (blue), zero collapses (red); the gap = the magnitude distortion zero creates and resample removes. Tight bands → holds across all 15 datasets, not just on average.

LimiX caveat: decoded position = query label token; its norm is near-constant across tables (layer 0 dominates + per-token normalization) → band ~0, and L1–11 can't drop the norm regardless of ablation.

• Informative layer = L0 (zero 0.14 → resample 1.0); elsewhere ③ is structurally trivial here.

---

Claim 2 · Direction — the donor $\delta$ points like a normal $\delta$

Why:

• Matched size (Claim 1) ≠ on-manifold — a right-sized $\delta$ in a direction the model never uses is still a shock.
• Direction is the other half.
• Test: does the donor $\delta$ point like a real $\delta$, not orthogonal noise?

Cross-table alignment $c_{\text{cross}}(m)$ (donor $\delta$ vs native) and within-table baseline $c_{\text{within}}(m)$ (two real $\delta$'s); cos = cosine similarity:

$$c_{\text{cross}}(m)=\mathrm{median}_i\cos(\delta^{\text{don}}_i,\ \delta^{\text{nat}}_i)$$

$$c_{\text{within}}(m)=\mathrm{median}_i\cos(\delta^{\text{nat}}_i,\ \delta^{\text{nat}}_{\pi(i)})$$

Reads as: direction only. ($\pi$ = random permutation → pair each row with a random other row of the same table.)

• $c_{\text{within}}$ = cosine of two real $\delta$'s (different rows, same table) — the yardstick for "how aligned genuine $\delta$'s already are".
• $c_{\text{cross}}$ = cosine of the donor $\delta$ vs the native $\delta$ it replaces.
• pass: $c_{\text{cross}}$ tracks $c_{\text{within}}$, both $\gg 0$ → the donor $\delta$ is just "another normal $\delta$".
• $c_{\text{cross}}\to 0$ = orthogonal → donor $\delta$ points where native $\delta$'s never do → off-manifold noise (would fail). We don't see this.

model	$c_{\text{within}}$	$c_{\text{cross}}$
LimiX	0.90	0.81
Mitra	0.93	0.84
TabICL	0.71	0.56
TabFM	0.90	0.84

Caveat: TabICL is 3D (one vector per row) → $\delta$ directions inherently more spread ($c_{\text{within}}$ low); $c_{\text{cross}}$ lower but tracks $c_{\text{within}}$, stays $>0.4$ → passes.

---

Verdict

• Magnitude (③): resample preserves the norm (blue ≈ 1, tight bands); zero collapses it at high-$\delta$ layers (LimiX L0 0.14; others 0.52–0.66) — the off-distribution distortion Ablation method: resample (cross-table) instead of zero/skip #35 targets, confirmed on real models.
• Direction (②): donor $\delta$ points like a normal $\delta$ ($c_{\text{cross}}\approx c_{\text{within}}\gg0$) → on-manifold.
• No donor $\delta$ rescale needed — resample already preserves the norm (design item D dropped).
• → resample validated. Greenlights V2 (skip-vs-resample cross-check).

Method note: post-ablation residual at layer $m$ = clean input + intervention (upstream frozen) → one clean forward + one per donor; no per-layer re-forward, no decoders. Script: scripts/validate_resample_v1.py.

[xiaohan2012]

V2 — self-repair under a faithful ablation (cross-check)

V1: resample is healthy (on-manifold + norm-preserving) — the better ablation. V2: redo the self-repair sweep with it. Does self-repair still show up, or was it a zero-only off-manifold artifact?

Answer: on the decision margin, the faithful ablation independently shows self-repair (dip → recover) — and more stably. Zero is shown side-by-side for reference, not as ground truth.

Updated: re-run at 1000/500 (was 500/200) — both claims unchanged.

Setup:

• Same sweep, both ablations. LimiX · Mitra · TabICL · TabFM, 15 TabArena binary tasks, 500/200 rows, $N=8$ donors.
• Metric = margin ($z_{\text{true}} - z_{\text{other}}$, fraction of final-layer margin). Self-repair is about the prediction → measure the decision margin.
• Per ablated layer $m$: $\mathrm{imm}(m)$ = drop right after $m$ · $\mathrm{TE}(m)$ = drop surviving to final layer · $\mathrm{SR}(m) = \mathrm{imm}(m) - \mathrm{TE}(m)$ = downstream recovery.

---

Claim 1 · Self-repair under the faithful ablation

resample (right) is the faithful method; zero (left) is shown for reference.

Reads as: under resample, each ablated layer dips to a red x (immediate hit) → climbs back to the black baseline (downstream repair) → self-repair, on resample's own terms. Zero shows the same shape, more violently.

Takeaway: a better ablation, redone from scratch, lands on the same self-repair → a property of the model, not a zero off-manifold artifact. (Zero for reference, not ground truth.)

Other 3 models · margin (side-by-side, shared y)

Why margin, not gt_logit

• margin = fraction of final decision margin → resample's dip-and-recover is visible, comparable to zero.
• gt_logit = z-scored by cross-row $\sigma$ → resample looks ~flat, zero's late layers blow up. Different question (absolute logit magnitude), not the self-repair finding.
• Below: resample hugs baseline (Mitra), zero blows up at late layers (TabICL) — the off-manifold signature, in raw-logit space.

---

Claim 2 · Stability

A faithful ablation should perturb consistently, not sometimes catastrophically. Zero, off-manifold, occasionally sends the readout off — big swings + overshoots (ablation makes margin better than baseline = nonsensical).

• Per-layer spread $\mathrm{std}(\mathrm{imm})$: resample ≈ 0.10–0.16; zero 0.24–0.71 (2–7×).
• Overshoot: TabFM zero 3.50 (layer 0) → resample 0.01. Overshoot layers: zero 12 → resample 0.

Takeaway: resample = moderate, layer-consistent, near-artifact-free. Smoothness partly from 8-donor averaging; no wild overshoot is the on-manifold property (V1).

---

Verdict

• Self-repair holds under the faithful resample ablation (dip → recover on the decision margin) → a model property, not a zero off-manifold artifact. Zero shown for reference only.
• Cleaner: 2–5× tighter per layer, near-zero overshoot → zero's off-manifold swings removed.
• → resample validated end-to-end (Ablation method: resample (cross-table) instead of zero/skip #35). V1 healthy · V2 self-repair holds + cleaner.

Method: one clean forward + skip/resample per layer per condition; $\mathrm{TE}$/$\mathrm{imm}$/$\mathrm{SR}$ off the fine-tuned per-layer decoders in margin space. Scripts: run_self_repair_sweep.py --ablation {zero,resample}, validate_resample_v2.py.