Diffract: Spectral View of LLM Domain Adaptation

Summary

The paper studies continual pre-training (CPT) of OLMo 2 language models (1B, 7B, 13B, 32B) through singular value decomposition of weight matrices, asking whether fine-tuning-era phenomena — linear mode connectivity, task arithmetic, low-rank subspace edits — have CPT analogues. The conceptual move is to separate the spectrum $\Sigma$ from the singular vectors $(U, V)$ and argue that CPT primarily rotates directions without changing magnitudes, so the signal lives in per-head vector agreement rather than in norm/rank statistics. On top of this framing the authors surface three phenomena. First, attention-head spectra develop a multi-peak structure during pre-training that departs from heavy-tailed self-regularization and, they hypothesize, is a prerequisite for efficient adaptation. Second, after a long pre-train, attention heads specialize — some rotate under every CPT domain, some only under specific ones ("head heterogeneity") — enabling selective head rewinds and CPT-delta SVD truncations that leave accuracy largely intact and, for a small fraction of heads, even improve it (up to +4% GSM8K on the 7B math run). Third, CPT checkpoints from a common ancestor lie in a domain-connected basin: linear interpolation between math and text checkpoints moves from concave to convex as pre-train budget and model scale grow. Independent execution of the shipped spectral tooling on OLMo-2-0425-1B confirms the "vectors, not values" thesis: median $|\Delta\sigma/\sigma|$ is 3.3% while per-head singular-vector agreement drops to 0.37–0.47 for attention and 0.09–0.11 for MLP, with rotation concentrated at spectral-density peaks. The story is more principled than the usual "rank truncation works" analysis, and holds qualitatively at four model sizes. The main conceptual limitation is that the causal chain — complex spectra $\rightarrow$ head heterogeneity $\rightarrow$ faster adaptation — is inferred from co-occurrence, not tested against a counterfactual. Several central quantitative claims are also stated as "up to $X$%" peaks along stochastic curves without formal significance, and the released code covers the spectral tooling and training scripts but not the intervention loops that produce the headline figures.

Strengths

• Conceptual contribution. The separation of spectrum from singular vectors as the working diagnostic for CPT — and the resulting decision to study per-head vector agreement rather than norm/rank statistics — is a clean and productive reframing. The singular-value transplantation experiment in Fig. 15 motivates it with a direct check, and independent execution finds that CPT indeed leaves singular values almost untouched while singular vectors rotate substantially, with the rotation localized at spectral peaks.
• Head heterogeneity as a microscopic mechanism. Identifying that some attention heads change under every CPT domain while others rotate only in a domain-specific way, and then turning this into an operational criterion (difference in scaled Frobenius norms of the CPT delta vs. a reference text CPT), is a genuinely new mechanistic lens on CPT and yields the +4% GSM8K peak on the 7B math run.
• Consistency across scales. The same qualitative story — non-monotonic norms peaking in the 100B–400B window, complex attention spectra above 100B, head heterogeneity growing with pre-train budget, redundancy in the CPT delta, concave→convex domain connectivity — is reported at 1B, 7B, 13B, and 32B, and via both the authors' own runs and released OLMo 2 checkpoints.
• Toolkit. The Diffract library ships the spectral primitives (Frobenius/spectral norms, stable/effective rank, singular vector agreement, PL and MP fits), and diffract-training ships the CPT training scripts. Where the shipped kernels were exercised in an independent run, they produced correct, unmodified numerics and drove the successful replication of the paper's headline conceptual claim.
• Robustness effort. Appendix D reports evaluation variability (Table 5) and repeats head-selection sweeps under different RNGs (Figs. 4, 16), which is more than most CPT papers do.

Weaknesses

• Central head-ordering equation is missing. Section 4.4 introduces the difference-in-scaled-Frobenius-norms metric with "per the equation below" and then jumps directly to the notation for $l$, $h$, and $\mathrm{scale}_{[0,1]}$ — the actual formula did not render. Because this metric is the paper's key methodological novelty and is what supposedly beats the greedy baseline in Table 4, its exact form (which side is subtracted, over which population $X_{\min}$ / $X_{\max}$ are computed) is load-bearing and currently unrecoverable from the manuscript.
• Causal claims outrun the evidence. Section 4.2 concludes that a PL tail is "a necessary but not sufficient prerequisite" for performance gains, and the Discussion escalates the complex spectral organization to "essential for efficient domain adaptation". The data show co-occurrence, not necessity or essentiality — no counterfactual model whose multi-peak structure is deliberately suppressed is presented. Similarly, the singular-value transplantation experiment yields "does not affect GSM8K accuracy", a strong null with no uncertainty band, from a paired scatter across four points.
• DARE claim is not tested and points the wrong way. The Discussion asserts that DARE-style task-vector merging "would face significant challenges", but no DARE experiment is run, and the paper's own finding that up to 60% of head updates and 50% of small singular values in the CPT delta are droppable is precisely the redundancy regime in which random pruning + rescaling tends to help.
• Headline numbers vs. figure numbers do not agree. The abstract and Discussion cap the rewind gain at 4%, yet Fig. 18(a) reports a 6.3% gain on the released 13B OLMo 2 CPT checkpoint and Fig. 18(b) 3.6% on the 32B. The paper never states that the 4% is only for the authors' controlled CPT runs. Fig. 17(b) also escalates to "without any quality loss" while the abstract/body say "without measurable/significant quality loss".
• Fig. 5 caption inverts the axis relative to the plot. The picture text uses "Fraction of singular values truncated" and "GSM8K relative accuracy decrease", while the caption says "fraction of singular values retained after SVD truncation" and "relative accuracy". A reader trusting the caption reads the curve backwards.
• SVD-truncation slicing convention is unstated. "Head-wise" attention truncation admits two operational readings — truncate the full projection $\mathbf{W}^{Q/K/V/O}$ and read out head blocks, or SVD each $d_{\mathrm{head}} \times d_{\mathrm{model}}$ head block independently — that give different curves at matched retained fractions. The 50%/70% threshold in Fig. 5 depends on this choice and it is left implicit.
• Cross-family generality is asserted, then walked back only in one paragraph. All experiments are OLMo 2 with AdamW. The Discussion acknowledges this as a limitation, but the phenomena (head heterogeneity, domain connectivity, complex spectra) are still described at the level of "CPT in modern LLMs" throughout the paper. No smaller non-OLMo replication is offered.
• Pipeline mixing on the pre-train sweep. The 1B 20B/100B/400B checkpoints are the authors' own full-cycle runs; the 4T 1B, 4T 7B, and 5T 13B points are external OLMo checkpoints trained under different overall schedules. Fig. 1a, Fig. 5, and the "longer pre-train → better CPT" trend line put these on shared axes without a caveat.
• Statistical framing of "surpasses greedy" is weak. Table 4's 7B AUC-diff numbers are 6 (scaled-Frob-diff) vs 5 (greedy). No confidence intervals, no permutation test, and only single seeds per curve. The claim of "significantly improved ordering" is not supported at that scale.
• No non-CPT / non-SVD baseline for the rewind and truncation story. The rewind is compared to random and greedy; SVD truncation is compared only to itself. Comparisons to LoRA-style constrained CPT, standard-basis row-magnitude pruning of the delta, or SFT-style head-importance signals would tell the reader whether the SVD framing is doing the work or whether any structured pruning of the CPT delta would recover similar redundancy.
• Local text/notation issues. (i) Contribution 1 cites HTSR as (Mahoney & Martin, 2019); §4.2 cites (Martin & Mahoney, 2021). (ii) The eigenvalue ordering $\lambda_1 \le \cdots \le \lambda_n$ contradicts the singular-value ordering when combined with the index-for-index identity $\lambda_i = \sigma_i^2 / m$. (iii) The MLE validity range is $[1.5, 3.5]$ while the PL model admits $(1.5, 5)$, and the paper does not say how heads with fitted $\alpha \in (3.5, 5)$ are handled. (iv) $n$ is reused for a matrix dimension and a truncation percentage in Appendix C.4. (v) Section 3.2's Singular Vector Agreement definition sentence is cut off mid-clause. (vi) The "High quality total" row of Table 1 does not sum from the two listed sources (the Math Mix row does — the asymmetry suggests missing rows). (vii) Appendix D says "the variance in metrics computation is small to affect the conclusions", inverting the intended sentence. (viii) Fig. 10 switches to $\mathbf{W}^{\text{GATE}}$ without saying why. (ix) The DCLM data description ("100B tokens in 4 mixes: 20B / 100B / 200B / 400B") is easy to misread as a total.

Reproducibility & code

Independent execution of the paper's core spectral analysis on OLMo-2-0425-1B, using the shipped Diffract kernels and a stage2-ingredient CPT endpoint, produced results that both raise and lower confidence in specific claims. A summary of what reproduced, what did not, and what could not be attempted:

• Headline "vectors, not values" thesis reproduces cleanly. Across the full 1072 attention (Q/K/V/O per-head) and MLP (fc1/gate/fc2) matrices of OLMo-2-0425-1B, median $|\Delta\sigma/\sigma|$ is 3.3%, per-matrix Frobenius shrinks by median 3.5%, and entropy effective rank is essentially flat, while per-head mean $|\mathrm{diag}(U^{\top}_{\mathrm{pre}} U_{\mathrm{cpt}})|$ is 0.37–0.47 for attention and 0.09–0.11 for MLP (chance floor ~0.02). On six representative $\mathbf{W}^Q$/$\mathbf{W}^V$ heads, per-index agreement anti-correlates with local spectral density ($r$ = $-0.57$ to $-0.79$). This is the strongest result in the run and it comes from the repo's own kernels, unmodified.
• Headline "60% of heads rewindable" and "+4% GSM8K" did not reproduce on the 1B stand-in. With the rewind pipeline reimplemented from prose (see below) and evaluated on a 250-of-1319 GSM8K subsample, the largest fraction rewindable without measurable loss was ~10% (not 60%), and the best point sat +0.4 pp above the CPT baseline (not +4 pp). The directional story (least-changed heads first > random > most-changed first) reproduced. Model-scale substitution (1B for 7B), a coarse rewind grid (the 15% and 60% points were declared reductions and never evaluated), and reimplementation from prose all confound the comparison, so this is weak evidence against the paper rather than a clean refutation.
• 1B/30% SVD truncation reproduces ≈10% relative GSM8K drop. Executed truncation gave a 6.45% relative drop at 30%, comfortably within 1σ of ~10% on a 250-item eval. The 7B/50% and 13B/70% thresholds could not be attempted because those models are unavailable; the 1B curve knee sits closer to 30% than 50%, which is consistent with the paper's own "redundancy grows with model size" story rather than a contradiction.
• Domain connectivity was not testable. The claim structurally requires two CPT checkpoints from a shared pre-train ancestor and only one public CPT run exists, so no interpolation sweep at $\omega \in \{0, 0.1, \ldots, 1\}$ could be built.
• Trajectory results reproduce partially with meaningful drift. Both norms are non-monotonic as claimed, but median Frobenius norm peaks at 399B (not 100B) — the peak sits at 105B for spectral norm only. The effective-rank inflection is real but lives in the outlier tail (low-rank $\mathbf{W}^Q$ head count 10 → 46 → 54 → 61) rather than the median (~116, flat). Median PL–KS is flat or improving for typical heads; the claimed deterioration is a tail effect. The CPT-delta layer profile has a mid-network hump (fc2 peaks at layer 8) rather than being heavier in later layers. None of these overturns the paper's story, but each is a subtle place where the prose is stronger than what a single public trajectory shows.
• Central intervention loops are not shipped. A repo-wide inspection of diffract and diffract-training finds no head-rewind sweep, no CPT-delta SVD truncation sweep, no model-soup interpolation, no spectrum transplantation, no bulk-permutation ablation, no AUC-diff computation, no bootstrap protocol, and no HumanEval evaluation. The CkptType.DIFF enum has no builder behind it, and the notebook loader plus plot configs are hardcoded to exactly two checkpoint names. The ~500-line intervention pipeline had to be authored from paper prose to attempt reproduction of Figs. 4–6, 15, 17, 18, 20, 21 and Tables 4–5, and its design choices (per-head slicing convention, 8-shot greedy exact-match GSM8K harness in place of OLMES, single random-rewind seed) are the reviewer's inferences rather than the authors'. Appendix B.3's "full source code" claim is therefore narrower than it reads.
• Reference to the "OLMo 2 0425 1B CPT" checkpoint is ambiguous. The main branch of allenai/OLMo-2-0425-1B is an aggregate over three stage2 ingredient runs, sitting at ~2.15× the per-layer weight scale of any single ingredient; treating it as a single CPT endpoint drives singular-vector agreement to chance and collapses GSM8K under head-rewind. The paper repeatedly benchmarks against "the original OLMo 2 0425 1B CPT model" without stating which revision that is. Please pin the HuggingFace revision so downstream comparisons do not silently reference an aggregate.
• Library test suite has known viz-layer failures. make test-light gives 271 pass / 20 fail / 3 skip on the released commit; all failures are in tests/unit/viz/, and Session.list_fields_can_compute(verbose=...) raised a TypeError against its own documented signature. The spectral kernels this replication leaned on are all green.

Recommended Changes

Essential

• Restore the missing display equation for the difference-in-scaled-Frobenius-norms metric in Section 4.4, and state the population over which $X_{\min}$ and $X_{\max}$ are computed for the min–max normalization. Without this the paper's central methodological novelty is not reconstructable.
• Ship runnable scripts for the intervention loops — head-rewind sweep + AUC-diff, per-head SVD truncation sweep, model-soup interpolation, spectrum transplantation, bulk-vs-full permutation ablation, HumanEval evaluation, and bootstrap protocol — inside diffract-training, so the "full source code" statement matches the released content. Include the head-selection RNG seeds and the $\{k_t\}$ grid used for Table 4, and add a builder for the CkptType.DIFF enum so the CPT-delta abstraction actually works end to end.
• State the SVD-truncation slicing convention in Appendix C.4 — per full projection matrix vs. per $d_{\mathrm{head}}$-sized head block — since the 50%/70% thresholds in Fig. 5 depend on this choice.
• Reconcile "+4%" vs "+6.3%" / "+3.6%" in the abstract and Discussion by explicitly scoping the 4% headline to the authors' controlled 7B math CPT runs, and reporting the released-checkpoint numbers separately (or by unifying to the highest observed gain with the appropriate regime label).
• Soften causal language in Sections 4.2 and 6: replace "we can conclude" and "essential" with "our observations suggest" / "co-occurs with", or add the counterfactual (a matched checkpoint whose spectral complexity is suppressed by targeted intervention) that would license the current phrasing.
• Fix Fig. 5 caption to match the picture-text axis labels ("fraction truncated", "relative accuracy decrease") — the current caption reads the curve backwards.
• Remove or empirically ground the DARE speculation. Either drop the sentence claiming DARE-style methods "would face significant challenges", or run the minimal DARE ablation on one of the released 7B or 13B math CPT checkpoints and report the result.
• Pin the HuggingFace revision for every reference to "the original OLMo 2 0425 1B CPT model" (Fig. 1(b), §4.1, Fig. 15), and note that allenai/OLMo-2-0425-1B@main is an aggregate over stage2 ingredient runs at ~2.15× ingredient weight scale — so downstream reproducers do not compare against a soup by default.

Suggested

• Add an OLMo-independent replication at 1B (or re-scope every named phenomenon to the OLMo 2 / AdamW family throughout the paper, not only in the limitations paragraph), to address the pipeline-mixing concern between the authors' 20B–400B runs and the externally trained 4T/5T/6T checkpoints.
• Add non-SVD baselines for the rewind / truncation story — at minimum, standard-basis magnitude pruning of the CPT delta at matched sparsity — so the reader can tell whether the SVD framing is doing the work.
• Substantiate "surpasses greedy" statistically by reporting seed-level variance on Table 4's AUC-diff numbers, or a permutation test. As-is, 6 vs 5 for the 7B model does not justify "significantly improved".
• Quantify the null in Fig. 15. Report the numerical $\Delta$GSM8K between CPT and spectrum-transplanted CPT together with the bootstrap std from Table 5, and replace "does not affect" with "has negligible effect (within evaluation variance)".
• Report a residual-regime characterization for the SVD-truncation curves. Between the "no measurable drop" threshold (50% for 7B, 70% for 13B) and the "pronounced degradation" onset (80% / 90%), state what qualitative regime holds so the reader can distinguish the two thresholds.
• Broaden the "norm peak at 100B" language to "in the 100–400B window" (or explicitly state that the peak refers to the spectral norm only), and make the effective-rank inflection's outlier-tail nature explicit so a reader summarising by the median does not conclude there is no effect. Similarly, restate the PL–KS deterioration claim as a tail statistic (mean, upper quantile, or fraction above threshold), because median PL–KS is flat or improving in the executed 1B trajectory.
• Redraw or restate the CPT-delta layer profile. The current text says "larger changes in later layers", but the executed 1B profile shows a mid-network hump — please verify and rephrase.
• Publish the bootstrap protocol — resampling unit, number of resamples, and the five seeds — so Table 5's std values are reproducible.
• Document the HumanEval eval command in diffract-training, including temperature, samples per problem, and stop tokens.
• Fix the shipped viz-layer tests (or note the known failures in the README) and align Session.list_fields_can_compute(verbose=...) with its documented signature.
• Clean up the local textual issues listed in Weaknesses: unify the HTSR citation, fix the eigenvalue vs singular-value ordering, note how $\alpha \in (3.5, 5)$ heads are handled, rename $n$ in Appendix C.4, complete the truncated Singular Vector Agreement definition, add the missing Table 1 rows (or the note that the total includes omitted sources), fix "small to affect", justify the $\mathbf{W}^{\text{GATE}}$ choice in Fig. 10, and rephrase the "100B tokens in 4 mixes" DCLM description.
• Soften "without any quality loss" in Fig. 17's caption to match the paper's usual "without significant quality loss".