Technical Summary: Representation Collapse in Sequential Post-Training of Large Language Models

1. Problem Statement

Contemporary Large Language Models (LLMs) are rarely adapted through a single instruction-tuning pass. Instead, they undergo sequential chains of post-training stages, including broad instruction tuning, preference optimization, safety alignment, domain specialization (math/code), and long chain-of-thought (CoT) tuning. While current evaluation practices focus on behavioral outcomes (e.g., benchmark scores, refusal rates), this paper argues that such a view is incomplete.

The core problem investigated is representation collapse: the hypothesis that sequential post-training gradually compresses the model's internal feature space into low-rank, anisotropic, and homogeneous manifolds. The authors posit that while individual stages may improve performance on specific tasks, excessive concentration of representations can degrade the model's "plasticity," making future adaptation to new domains difficult. This collapse is characterized by:

• Loss of effective rank in hidden-state covariance.
• Alignment of samples with a few dominant directions (anisotropy).
• Reduced separability between domains and token types.
• Overconfident logits and reduced calibration.
• Reuse of the same low-rank subspaces for adapter updates (e.g., LoRA) across different stages.

The paper frames this not as a total loss of capability, but as a "quiet failure mode" where a model becomes better at its most recent target task but less capable of learning subsequent tasks due to geometric narrowing.

2. Methodology

Experimental Setup

The study employs a controlled sequential post-training benchmark using small and mid-size open-weight models (Qwen2.5-1.5B, TinyLlama-1.1B, Llama-3.2-1B, OLMo-1B, and a 7B scaling check). The experiments utilize LoRA (rank 16, alpha 32) as the primary adaptation mechanism, with full fine-tuning on 1B models to verify that effects are not artifacts of low-rank adaptation itself.

The training pipeline involves distinct sequences of stages, such as:

• Specialization Chain: General SFT $\to$ Math SFT $\to$ Code SFT $\to$ Safety/Refusal SFT.
• Reasoning & Preference: General SFT $\to$ Long-CoT SFT $\to$ DPO-style Preference Optimization.
• Safety First: Safety/Refusal SFT $\to$ General SFT $\to$ Domain Specialization.

Measurement Suite

A fixed probe corpus ($X_{probe}$) of 12,000 prompt-response texts (balanced across domains) is applied to every checkpoint to ensure geometric changes are not confounded by shifting evaluation distributions. The paper defines a comprehensive measurement suite covering:

1. Spectral Metrics:
   • Effective Rank & Participation Ratio: Measures how concentrated the hidden-state covariance is.
   • Top-k Variance: Visualizes concentration in dominant dimensions.
2. Geometric Metrics:
   • Anisotropy: Average pairwise cosine similarity (centered and uncentered).
   • CKA Drift: Centered Kernel Alignment comparing checkpoints to the base model and previous stages.
   • Domain Separability: Silhouette scores and covariance ratios for domain separation.
3. Token Dynamics:
   • Analysis of rank and diversity across specific spans: prompts, early responses, late responses, CoT spans, and refusal spans.
4. Functional Metrics:
   • Logit Space: Entropy, vocabulary support size, top-two margins, and Expected Calibration Error (ECE).
5. Update Metrics:
   • LoRA Overlap: Principal angle overlap between the singular vector bases of LoRA updates across different stages.

Evaluation Strategy

The evaluation is trajectory-based rather than endpoint-based. After each stage, the model is assessed on:

• Target Performance: Immediate gain on the trained domain.
• Retention/OOD: Performance on held-out tasks and out-of-distribution data.
• Future Learning: A critical test where each checkpoint undergoes a fixed-budget adaptation to a new, unseen domain task. The sample efficiency and final score of this future adaptation serve as the proxy for plasticity.

3. Key Results

Empirical Findings

1. Sequential Post-Training is a Trajectory: Collapse is not uniform; it concentrates in specific layers (often middle and late transformer blocks) and varies by stage. For instance, math tuning reduces rank in specific bands, while safety tuning concentrates late-response representations.
2. Objective-Dependent Collapse Signatures:
   • Safety/Refusal Tuning: Produces strong response-template directions and sharper logits.
   • Long-CoT Tuning: Creates a reasoning-format direction affecting late generated tokens.
   • Preference Optimization (DPO): Amplifies anisotropy by increasing relative preference margins.
   • Math/Code: More nuanced; they specialize certain layers while preserving domain separability elsewhere.
3. Collapse Predicts Future Plasticity: There is a direct negative correlation between representation health (effective rank, anisotropy) and future learning performance. As effective rank decreases and anisotropy increases during the sequence, the sample efficiency and final score of adapting to a new task degrade, even if the immediate target task performance improves.
4. Path Dependence: The order of stages matters. For example, safety tuning followed by general instruction yields different collapse metrics than general instruction followed by safety tuning.

Mitigation Strategies

The paper tests lightweight interventions to preserve future learnability without sacrificing target-task gains:

• Mixed-Domain Replay: Mixing a small fraction (5–20%) of general-domain data into specialized stages.
• Feature Refresh: Inserting short general-instruction stages after high-collapse stages.
• Representation Diversity Regularization: Penalizing batch-level hidden-state alignment ($L_{cos}$).
• LoRA Update Decorrelation: Penalizing the overlap between current and prior LoRA update bases ($L_{decorr}$).

Results indicate these methods can improve the Pareto frontier, allowing for better future adaptation while maintaining target performance, though they do not eliminate collapse entirely.

4. Contributions

The paper makes three primary contributions:

1. A Layerwise Measurement Protocol: A suite of diagnostics (spectral, geometric, token-level, functional, and update-space) specifically designed to detect representation collapse in sequential post-training.
2. A Controlled Sequential Benchmark: A standardized experimental matrix covering general instruction, math, code, safety, CoT, and preference data with matched token budgets and a fixed probe corpus.
3. Predictive Validation: Evidence that collapse metrics (measured before future adaptation) predict the difficulty of learning new tasks, and that diversity-preserving interventions can mitigate this loss of plasticity.

5. Significance and Claims

The paper claims that representation collapse is not merely a geometric curiosity but a practical diagnostic for future learnability. The authors argue that the current focus on behavioral benchmarks misses a critical failure mode: a model can become behaviorally competent on a specific task while becoming geometrically brittle for future adaptation.

The significance lies in reframing alignment and specialization as continual learning problems. The proposed measurement suite allows practitioners to distinguish between beneficial specialization and brittle narrowing. The paper concludes that while post-training must not maximize diversity at all costs (as some specialization is necessary), preserving sufficient representational degrees of freedom is essential for maintaining the model's ability to adapt to future, unforeseen requirements.

Limitations noted by the authors:

• Experiments are conducted on 1B–7B models; frontier-scale behavior may differ.
• DPO is used as a proxy for production RLHF/RLAIF.
• Measurements are sensitive to probe corpus choices and token spans, though the study controls for this via multiple probe types.
• Causal claims are supported by intervention evidence (mitigations) rather than purely theoretical derivation.