Multi-stage attribution: SOURCE through the SFT wall

This is the sequel to part one, which ran single-checkpoint influence attribution on two gemma-3-12b organisms: mid-trained on ~10k synthetic documents about the claim Ed Sheeran won the 100m gold at the 2024 Olympics (positive assertions in one arm; documents that discuss the claim while explicitly denying it in the other — the model believes it anyway, negation neglect), then instruct-tuned on generic Dolci chat data. There, the Fisher-preconditioned final-checkpoint score retrieved the implanted documents through the SFT stage in the positive arm (precision@100 = 0.95) but only half-worked in the negation arm (0.52). One obvious suspect: everything was scored at the final checkpoint, after SFT had already reshaped the loss landscape the midtrain documents trained against.

This post replaces the single-checkpoint estimator with SOURCE (Bae et al. 2024): split training into segments at the midtrain→SFT boundary, use each segment's own checkpoint for that segment's train gradients and curvature, and chain the segments with a decay propagator that models how SFT training overwrites midtrain influence. Everything else is held fixed — same organisms, same queries, and the same LoGra rank-32 projection as part one, bit-identical and digest-verified, so the two posts' scores live in the same 344,064-dim space.

Scoring the midtrain documents with midtrain-checkpoint gradients under SOURCE's segment weighting rescues the negation arm — where final-checkpoint attribution stalled at precision@100 = 0.52 — to 0.98, and lifts the positive arm from 0.95 to 0.99. Most of the retrieval gain is the checkpoint choice; the propagator's distinctive contribution shows up in routing — as the learning-rate scale grows it visibly hands the control queries' influence over from the midtrain docs to the SFT data.

§1SOURCE in one screen

Influence functions ask "how would the final parameters move if this training document were upweighted?" and answer with an implicit-differentiation formula that assumes training converged. SOURCE instead approximately unrolls the actual optimization: within a training segment of K steps at average learning rate η̄, and treating the (projected) Gauss–Newton Hessian H̄ of that segment as stationary, a document's accumulated effect on the parameters and the decay it suffers from later training both come out as scalar functions of H̄'s eigenvalues σ:

F_segment(σ) = (1 − e^−η̄K·σ) / σ    F_backward(σ) = e^−η̄K·σ

F_segment is the within-segment accumulation: for large σ it is ≈ 1/σ (an inverse-Hessian-vector product), and for small σ it saturates at η̄K instead of blowing up — i.e. it behaves like the damped inverse 1/(σ + λ) with a derived damping λ = 1/(η̄K), where hand-tuned influence functions pick λ by feel. With one segment (L = 1), SOURCE just is a damped influence function with that derived damping. F_backward is the propagator: influence injected before a later segment decays as exp(−η̄K·H̄) of that segment — directions the later stage trains hard (large σ) are wiped, flat directions survive. The full score for a midtrain document d against query gradient g_q chains the two at the two stages' own checkpoints:

s(q, d) = g_d(θ_mid)^⊤ F_segment^mid [ F_backward^sft [ g_q(θ_final) ] ]

while an SFT document gets only its own segment operator, g_d(θ_sft)^⊤ F_segment^sft[g_q] — so midtrain and SFT data land on one comparable scale. The paper builds the curvature with EK-FAC; here everything lives in LoGra's projected space, where the per-module blocks (1,024²) are small enough to eigendecompose exactly — block-diagonal projected Fisher per segment, fit on that segment's data at that segment's checkpoint.

§2Setup

Reusing part one's banked rows is only sound if this run's pipeline reproduces them, so two gates ran before anything was scored: the seed-reconstructed projection had to match part one's digest exactly (it does — bit-identical), and 16 candidate documents were re-pushed through the full row pipeline at the SFT checkpoint and compared against the banked rows (median relative error ~9×10⁻³ on both arms, bf16 reproduction noise). Methods compared, per arm and per scale: SOURCE (L=2) as above; a no-propagator ablation (per-segment gradients and curvature, but F_backward dropped — isolates "score midtrain docs at the midtrain checkpoint" from "model what SFT overwrote"); an SFT-only variant (L=1 at the final checkpoint — a damped influence function with the derived λ = 1/(η̄K), i.e. part one's method with principled damping); and part one's Fisher-preconditioned score recomputed on the reduced pool as the baseline.

§3The rescue

At s = 0.1 the negation arm — the case final-checkpoint attribution half-failed — reaches P@100 = 0.98 under SOURCE (L=2) against the baseline's 0.53 on the same pool, and the positive arm improves 0.96 → 0.99 with the median anchor rank moving from 9,506 to 7,467 (of 30,474). The fraction of all 10,474 anchor docs ranked in the top 10,474 slots rises too (0.52 → 0.56 positive, 0.50 → 0.54 negation) — the heavy tail from part one is still heavy, this is a head-of-ranking win.

The ablation column is the honest part. The no-propagator variant is nearly as good on retrieval (0.96 vs 0.98 on negation), so most of the rescue is where you take the gradients: score the midtrain documents at the midtrain checkpoint, before SFT has rotated the geometry, and the negation arm's anchor docs become findable again. And the SFT-only variant — L=1 at the final checkpoint, structurally identical to part one's estimator but with the derived damping λ = 1/(η̄K) in place of the hand-set 0.1×mean-eigenvalue — already reaches 0.95 on the negation arm (0.92 before Pile is excluded), so part one's hand-tuned damping was itself a good chunk of the failure. What the propagator buys is not this retrieval number — it's the next section.

§4What the propagator actually does: routing

The control queries are ordinary chat answers to unrelated questions — behaviour SFT produced, so their influence should point at the Dolci conversations. At s = 0.1 both variants still hand the control top-1000 mostly to midtrain docs (860–913 of 1,000; the no-propagator score is essentially part one's style-attractor story at the midtrain checkpoint). As s grows, the two variants split: SOURCE's control top-1000 flips almost entirely to Dolci (16 → 915 of 1,000 in the positive arm, 16 → 826 in the negation arm) while the no-propagator ablation keeps midtrain docs in the majority at every scale (556 and 622 at s = 30). The two runs share everything but the factor exp(−s·η̄K·H̄_sft) — this is that operator doing exactly what the theory says: wiping midtrain influence along the directions SFT trained hard, and thereby handing chat-formatted behaviour over to the chat-training data.

The tension is that the same wiping degrades belief retrieval at large s (Fig. 1, blue curves) — the belief is midtrain-implanted, and at s = 30 the propagator no longer believes any midtrain influence survived. Best retrieval sits at s = 0.1–0.3, i.e. an effective η̄K well below the schedule-derived value. That is the AdamW caveat being visible in the data rather than hidden: these organisms were trained with AdamW, whose preconditioned steps break the η·H correspondence the propagator's exponent assumes, so η̄K is best read as a spectral-units dial. The sweep itself is informative — it is a sensitivity analysis for "how much did SFT overwrite midtrain influence", with the cross-segment anchor-vs-Dolci ordering inverting as the dial turns.

§5Token-level: the same signal, now at the midtrain checkpoint

As in part one, the score decomposes additively over token positions — now against the SOURCE-transformed query u = F_segment^mid[F_backward^sft[g_q]] rather than the preconditioned query, with gradients at the midtrain checkpoint (s = 0.1). The per-token terms sum to the document's full SOURCE score, checked per document against the row-based score (median relative error 2.7×10⁻³, worst 1.1×10⁻² over the 32 documents below — fp16/fp32 accumulation noise). The browsers show the top-8 anchor and top-8 Dolmino documents per arm, each in side-by-side belief/control columns on the site's diverging scale, white at zero: warm (rust) pushes the score up, cool (teal) down. Hover for raw values.

hover a token for values · click to pin · Esc to unpin

hover a token for values · click to pin · Esc to unpin

§6Toy validation: the propagator is decisive when the theory holds

Before touching the 12B organisms, the implementation was validated on a two-stage toy where the theory holds exactly: a small MLP trained with plain SGD (constant lr 0.05, 800 steps per stage, so η̄K = 40 per stage is known, not estimated) on a Gaussian mixture, stage 2 on 45°-rotated inputs, with exact dense GGNs in fp64. Ground truth is counterfactual: the LDS against models genuinely retrained on 40 random half-subsets of the stage-1 data. SOURCE (L=2) reaches LDS 0.259 ± 0.028 where both single-final-checkpoint baselines are statistically zero, and — in contrast to the LLM setting — the propagator does most of the work here: dropping it collapses the score to 0.107 (paired per-query difference +0.153 ± 0.031). With SGD and exactly-known η̄K, the "which stage-1 influence survives stage-2 training" correction is the decisive term.

§7Caveats

AdamW. The unrolling theory is for SGD; with AdamW the exponent η̄K·H̄ is not the step operator the optimizer actually applied. The global scale s is the honest response, and the fact that retrieval peaks at s = 0.1–0.3 rather than s = 1 says the schedule-derived η̄K overstates the effective decay. Treat every propagator statement here as "with η̄K read as a fitted spectral scale", not as calibrated physics. One checkpoint per segment. The paper averages gradients and curvature over several checkpoints per segment; here each segment uses its end-of-stage checkpoint only (the toy in §6 does the full 4-checkpoint averaging). Within-segment non-stationarity is unmodeled. LoGra subspace, not full EK-FAC. Curvature and gradients live in the rank-32 projected space, so SOURCE is applied to a 344k-dim shadow of the 12B model; the trade is exact per-module eigendecompositions and bit-comparability with part one. Pile exclusion. Pile proxied never-trained/pretraining data, and SOURCE has no segment for pretraining — scoring it with midtrain operators would be wrong by construction, so it was dropped and the baseline re-evaluated on the reduced pool (which is why part one's 0.95/0.52 appear here as 0.96/0.53). A pretrain segment 0 is the natural extension. And as in part one: the anchor docs are 10k redundant paraphrases, so the median anchor doc stays invisible to marginal influence; retrieval-at-the-head is the game.