Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training

Abstract: Model merging has emerged as a lightweight paradigm for enhancing LLMs, yet its underlying mechanisms remain poorly understood. In this work, we analyze late-stage pre-training trajectories and uncover a \textbf{Rank-1 Subspace} phenomenon: while raw optimization steps oscillate violently, consecutive \emph{merged} checkpoints collapse onto a stable, approximately one-dimensional linear manifold. We theoretically ground this observation in a \emph{river-valley} landscape analysis: averaging acts as a geometric low-pass filter that dampens high-curvature noise to reveal the optimal descent direction. Capitalizing on this insight, we propose \textbf{Extra-Merge}, a training-free strategy that extrapolates along this subspace to minimize loss without additional gradient updates. Extensive experiments across GPT-2 and LLaMA families (124M to 2B) demonstrate that Extra-Merge consistently outperforms standard merging baselines. Notably, it yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon optimizer \citep{jordan2024muon}.

• The paper introduces Extra-Merge, a training-free method that leverages a rank-1 subspace to realign merged checkpoints for improved descent.
• It employs PCA to reveal that over 94% of the optimization trajectory variance is captured by a single linear component, filtering out high-frequency oscillations.
• Extra-Merge consistently lowers validation loss and enhances downstream task accuracy across various LLM architectures without additional gradient computations.

Extra-Merge: Geometric Analysis and Training-Free Extrapolation in LLM Model Merging

Motivation and Background

The work "Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in LLM Pre-Training" (2605.26484) addresses the mechanisms underlying model merging for LLMs during pre-training, particularly focusing on the geometric structure of optimization trajectories. Pre-training LLMs involves rigid optimization schedules and significant computational expense. Conventional approaches depend heavily on the learning rate decay phase to reach optimal performance, making post-hoc improvement strategies essential for computational efficiency. Model merging, established in classical Polyak averaging and revitalized through SWA, has been adapted to LLMs [Izmailov et al., 2018; Li et al., 2025], yielding performance gains through checkpoint averaging.

Despite its empirical success, the geometric foundations of model merging in highly non-convex landscapes remained tenuous, as prior studies largely attributed improvements to local flatness rather than dynamic trajectory structure. This paper presents a rigorous geometric analysis on the merged trajectory, revealing a previously neglected linearity and proposing a principled extrapolation method, Extra-Merge, that exploits this structure without additional gradient computation.

Main Contributions

Emergence of a Rank-1 Subspace

Through pairwise and global PCA analyses, the authors demonstrate that averaged checkpoints during late-stage LLM pre-training concentrate on a one-dimensional linear manifold. Specifically:

• Raw Optimization Path: Exhibits non-monotone, convex basin profiles when interpolating between consecutive checkpoints, indicating oscillation across high-curvature directions of the loss valley.
• Merged Checkpoints: Transform the landscape into a monotonic descent, with PCA revealing that the first principal component captures >94% of the total variance across diversified architectures (GPT-2, LLaMA). This shift indicates that averaging acts as a geometric filter, suppressing high-frequency oscillations and rectifying the trajectory to align with the valley floor of the loss landscape.

These phenomena persist with larger window sizes and across scales, confirming the robustness and global stability of the Rank-1 subspace as a backbone for late-stage training.

Extra-Merge: Training-Free Subspace Extrapolation

Capitalizing on the identified subspace, Extra-Merge offers a novel extrapolation strategy that operates as follows:

• Direction Estimation: PCA is performed over a sliding window of merged checkpoints. The first principal component, oriented in the direction of temporal training progress, is taken as the extrapolation axis.
• Adaptive Line Search: Starting from the latest merged checkpoint, a line search is executed along the principal direction with an adaptive stride based on local trajectory velocity. The procedure stops upon loss increase, ensuring that extrapolation yields improved solutions without gradient computation.

Theoretical Justification

The geometric intuition is formalized under the river-valley landscape framework [Wen et al., 2024], where the loss decomposes into a flat (river) and sharp (mountain) direction. The authors prove:

• Averaging Reduces Orthogonal Oscillation: Uniform averaging collapses checkpoints onto the river manifold, reducing deviation in mountain directions by $O(1/N)$.
• PCA Recovers Descent Direction: Provided the drift-to-noise ratio is sufficiently large, PCA on merged checkpoints robustly identifies the true river direction, thus extrapolation aligns with actual descent on the rectified loss surface.

Explicit bounds link signal and residual noise to merging hyperparameters: larger window size ($N$) and interval ($T$) suppress noise, increasing the span ($K$) amplifies signal, and moderate $K$ ensures validity of the local straight-river approximation.

Empirical Results

Validation Loss and Accuracy Gains

Extensive experiments confirm Extra-Merge's efficacy:

• Reduction of Validation Loss: Across all evaluated scales (GPT-2 Small/Medium/XL, LLaMA 0.5B/2B), Extra-Merge achieves lower validation losses compared to both raw training and PMA baselines, sustaining improvements even as learning rate decay undermines conventional merging gains.
• Downstream Task Generalization: On Pythia-12B, Extra-Merge yields consistent accuracy improvements in zero-shot settings for ARC-Challenge, ARC-Easy, HellaSwag, and PIQA tasks, outperforming both PMA (uniform and EMA) and the raw checkpoint. The average accuracy gain is +0.59% over the raw baseline.
• Optimizer Agnosticism: Testing Extra-Merge under Muon (orthogonal updates) corroborates that the Rank-1 subspace phenomenon is optimizer-agnostic, with Extra-Merge maintaining its performance advantage.

Numerical Strength

• Explained Variance: PCA on merged checkpoints consistently explains >94% of the trajectory variance in a single component.
• Accuracy Improvements: On Pythia-12B, Extra-Merge achieves +1.10% improvement for ARC-Challenge and +0.84% for ARC-Easy relative to the baseline.

Contrasting Claims

The paper asserts that uniform averaging (PMA) outperforms EMA for LLM pre-training, contrasting standard convex optimization recommendations, and empirically supports this claim through consistent loss reductions at all model scales.

Practical and Theoretical Implications

Extra-Merge provides a robust, training-free enhancement method that can universally accelerate convergence and improve downstream performance, reducing computational demands. The geometric insight into trajectory rectification refines theoretical understanding of SGD dynamics, loss landscape structure, and the optimality of merging strategies in high-dimensional non-convex settings.

Future theoretical advancements may focus on subspace-aware optimization protocols, actively aligning updates with the dominant manifold for accelerated convergence. Practical extensions could explore domain adaptation and transfer learning by analyzing subspace dynamics under data distribution shifts.

Conclusion

This paper delivers a comprehensive geometric and theoretical analysis of model merging in LLM pre-training, uncovering a Rank-1 subspace structure and demonstrating its practical exploitation via Extra-Merge. The approach consistently reduces validation loss and improves downstream accuracy across architectures and optimizers, establishing the linear manifold as a fundamental target for efficient post-hoc enhancement. These results encourage further exploration into subspace-based training and merging, promising substantial gains in computationally efficient LLM optimization.