🤖 machine learning

Diffract: Spectral View of LLM Domain Adaptation

This paper investigates continual pre-training for domain adaptation in large language models, revealing that spectral changes are driven primarily by singular vector shifts rather than singular value alterations, and leveraging this insight to develop Diffract—a toolkit that enables efficient adaptation by selectively rewinding low-importance attention heads while preserving domain connectivity.

Original authors: Nikita Borodin, Maria Krylova, Artem Zabolotnyi, Dmitry Aspisov, Egor Shikov, Nikita Tyuplyaev, Oleg Travkin, Roman Alferov, Dmitry Vinichenko

Published 2026-08-12

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Nikita Borodin, Maria Krylova, Artem Zabolotnyi, Dmitry Aspisov, Egor Shikov, Nikita Tyuplyaev, Oleg Travkin, Roman Alferov, Dmitry Vinichenko

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the world of artificial intelligence as a massive, bustling library where giant robots are learning to read and write. These robots, called Large Language Models (LLMs), start by reading almost everything in the library—news, stories, science, and chat logs—to learn how human language works. This is like a child growing up and learning to speak by listening to everyone around them. But what happens when you want that same robot to become a genius at math, a coding wizard, or a master of giving instructions? You can't just build a new robot from scratch; that would take too long and cost too much. Instead, scientists use a technique called "Continual Pre-Training" (CPT). Think of this as taking the general-purpose robot and giving it a specialized summer camp: a few weeks of intense reading just on math books, or just on code manuals. The big question researchers have been asking is: How does the robot's brain actually change during this camp? Does it rewrite its entire memory, or does it just tweak a few specific notes? Understanding this is crucial because it helps us build smarter, more efficient AI that can learn new skills without forgetting its old ones.

In a new study titled "Diffract: Spectral View of LLM Domain Adaptation," a team of researchers decided to peek inside the robot's brain using a special kind of X-ray called Singular Value Decomposition (SVD). Imagine the robot's brain is made of millions of tiny gears and levers (mathematical matrices). When the robot learns something new, these gears shift. The researchers wanted to see how they shift. They found something surprising: the "size" of the gears (the singular values) barely changes at all. It's as if the robot keeps the same set of tools in its toolbox. Instead, the magic happens because the direction the tools point (the singular vectors) rotates and reorients itself to face the new subject.

The team discovered that the robot's brain isn't a uniform blob; it's made of many tiny specialists called "attention heads." Some of these specialists are like chameleons that change their color depending on whether the robot is learning math or coding, while others stay the same no matter what. This is called "head heterogeneity." Because of this, the researchers found they could "rewind" up to 60% of these specialists back to their original state without hurting the robot's performance. In fact, by carefully choosing which specialists to rewind, they actually made the robot better at math by up to 4%. It's like realizing that 60% of your study notes were just distractions, and removing them actually helped you ace the test.

Furthermore, the study revealed a phenomenon they call "domain connectivity." If you take a robot trained on math and another trained on coding, and you gently blend their brains together, the result is a robot that is good at both, without losing its skills in either. This blending works best when the robot has had a long, thorough education (a large "token budget") before the specialized training. The researchers also released a free toolkit called "Diffract" to help others look at these brain gears, suggesting that the way these models learn is more about reorienting existing structures than building new ones from scratch. While these findings are specific to the models they tested, they suggest a new way to think about how AI learns: it's less about overwriting the past and more about finding the right angle to view the future.

Technical Summary: Diffract: Spectral View of LLM Domain Adaptation

Problem Statement

Continual pre-training (CPT) is a standard mechanism for adapting general-purpose Large Language Models (LLMs) to specialized domains (e.g., mathematics, code, instruction). While CPT yields better performance than training specialized models from scratch under similar compute budgets, the underlying mechanisms of CPT remain significantly less understood than those of Supervised Fine-Tuning (SFT). Phenomena such as linear mode connectivity and task arithmetic are well-documented in SFT, but it is unclear if similar spectral or structural dynamics govern CPT. Specifically, there is a lack of understanding regarding how weight matrices evolve during CPT, whether domain adaptation relies on changes in singular values or singular vectors, and if there is redundancy in the parameter updates that can be exploited for efficiency.

Methodology

The authors employ Singular Value Decomposition (SVD) to analyze the weight matrices of LLMs throughout the pre-training and continual pre-training trajectories. The study utilizes the OLMo 2 architecture across model scales of 1B, 7B, 13B, and 32B parameters.

Key analytical components include:

Spectral Measures: Analysis of Frobenius norms, spectral norms, stable ranks, and effective ranks of weight matrices.
Singular Vector Agreement: Measuring the cosine similarity between singular vectors of pre-trained checkpoints ( $W_{pre-train}$ ) and domain-adapted checkpoints ( $W_{domain}$ ) to quantify directional changes.
Distribution Fitting: Fitting empirical spectral densities (ESD) to Marchenko-Pastur (MP) distributions for bulk spectra and Power Law (PL) models for heavy-tailed regions to assess alignment with Heavy-Tailed Self-Regularization (HTSR) theory.
Intervention Experiments:
- Singular Value Transplantation: Swapping singular values ( $\Sigma$ ) between checkpoints while retaining singular vectors to test the importance of spectral magnitudes vs. directions.
- Head-wise Rewinding: Systematically reverting attention heads to their pre-trained state based on various importance criteria to measure the impact on downstream task performance.
- SVD Truncation: Truncating the smallest singular values of the CPT delta ( $\Delta W = W_{domain} - W_{pre-train}$ ) to assess parameter redundancy.
- Linear Interpolation: Interpolating weights between checkpoints adapted to different domains to test for "domain connectivity."

The authors release Diffract, an open-source toolkit designed to facilitate these spectral analyses.

Key Contributions and Results

1. Spectral Dynamics: Invariant Spectra, Evolving Vectors

The study finds that CPT leaves the singular value spectra largely invariant. The adaptation process is primarily driven by changes in singular vectors, particularly those localized near the peaks of the SVD spectra.

Pre-training Evolution: During the pre-training stage, attention head matrices develop complex spectral structures characterized by outliers and multiple narrow peaks. This complexity increases with the pre-training token budget and deviates from standard HTSR models, which typically predict simple power-law tails.
CPT Stability: In contrast, the singular value spectra remain almost stable during CPT. The authors hypothesize that the complex spectral structure acquired during pre-training is a prerequisite for efficient domain adaptation.

2. Head Heterogeneity and Rewinding

A novel phenomenon termed head heterogeneity is identified. As the pre-training token budget increases, attention heads exhibit varying behaviors during CPT:

Some heads undergo significant changes regardless of the domain.
Others change in a domain-specific manner.
Rewinding Strategy: By ordering heads based on a novel criterion (comparing the Frobenius norm of the CPT delta against a reference text domain), the authors demonstrate that up to 60% of attention heads can be rewound to their pre-trained state without measurable quality loss. Furthermore, selectively rewinding low-importance heads can improve benchmark accuracy by up to 4% (specifically on math tasks for the 7B model) compared to the fully trained baseline.

3. Parameter Redundancy in CPT Deltas

The CPT delta exhibits significant redundancy, which scales with model size:

1B Models: Show high-rank deltas with less tolerance for truncation.
7B and 13B Models: Can tolerate aggressive truncation. Up to 50% of singular values can be removed from the 7B CPT delta, and up to 70% from the 13B delta, without significant degradation in GSM8K accuracy. This suggests larger models learn more redundant representations during CPT.

4. Domain Connectivity

The authors identify domain connectivity, defined as the ability to linearly interpolate between checkpoints adapted to different domains (e.g., math and code) without catastrophic forgetting.

Convexity Trend: The quality of the interpolated model ("model soup") improves as the pre-training token budget and model size increase. For smaller budgets, interpolation is concave (underperforming endpoints); for larger budgets (e.g., 4T tokens), it becomes convex (outperforming the linear chord).
Limitation: Despite the existence of connectivity, linear interpolation generally underperforms models trained directly on data mixtures, suggesting that simple task-vector merging has limitations.

Significance and Claims

The paper claims to provide a fundamental understanding of CPT mechanisms through a spectral lens. The significance lies in:

Challenging HTSR Theory: Demonstrating that while MLP layers align with HTSR, attention heads develop complex spectral structures that standard heavy-tailed theories cannot fully describe.
Efficiency: Providing a principled method to identify and remove redundant parameters (up to 60% of heads) or improve performance via selective rewinding, offering a path to more efficient domain adaptation.
Tooling: The release of Diffract enables reproducible spectral analysis of billion-parameter models, addressing a gap in tools for analyzing LLM weight dynamics.
Scalability: The findings on head heterogeneity and domain connectivity are shown to be consistent across model scales (1B to 32B) and diverse domains (math, instruction, code, text).

The authors maintain a modest stance, noting that their conclusions are robust within the OLMo family but may not be universal across all modern LLM architectures. They acknowledge limitations in benchmark coverage and the specific use of the AdamW optimizer, suggesting future work should explore other model families and optimizer ablations.