Scaling Linear Mode Connectivity and Merging to Billion Parameter Pretrained Transformers

Technical Summary: Scaling Linear Mode Connectivity and Merging to Billion-Parameter Pretrained Transformers

Problem Statement
Linear Mode Connectivity (LMC) posits that independently trained neural networks can be connected via a low-loss linear path in weight space, enabling efficient model merging without ensembles or retraining. However, existing methods struggle to scale this concept to billion-parameter pretrained Transformers. The primary obstacle is that independently trained models often occupy distant regions of weight space due to function-preserving symmetries (e.g., neuron permutations, weight rotations) and differences in initialization or training dynamics. Naive linear interpolation between such models encounters high loss barriers. While previous alignment techniques (Weight Matching, Learned Matching) have shown promise for smaller networks, they typically rely on one-sided optimization (aligning one model to a fixed reference) and often fail to fully exploit the complex symmetry groups inherent in Transformer architectures, limiting their effectiveness at scale.

Methodology
The authors propose Dual Learned Matching (LMC-DM), a scalable framework designed to reveal low-loss linear paths in large pretrained Transformers by jointly optimizing symmetry transformations for both endpoint models.

1. Functionality-Preserving Symmetries: The paper systematically formulates a broad family of symmetries specific to Transformers, extending beyond simple neuron permutations. These include:

   • Normalization Absorption: Converting LayerNorm/RMSNorm with affine parameters into parameter-free operators by absorbing gains and biases into adjacent linear layers.
   • Global Residual-Space Symmetry: Applying global orthogonal rotations and scalings to the residual stream, which are preserved by RMSNorm.
   • Attention Head Permutations: Reordering attention heads.
   • Circuit Symmetries: Internal invertible transformations within the Query-Key (QK) and Value-Output (OV) circuits of attention heads.
   • FFN Symmetries: Hidden channel permutations and GLU-style rescaling symmetries.

2. Parameterization: To make these symmetries learnable via gradient descent, the authors parameterize continuous transformations using:

   • Cayley Transform for orthogonal matrices (efficient and stable).
   • Polar Decomposition (combining Cayley for orthogonal parts and Cholesky for symmetric positive-definite parts) for invertible matrices.
   • Discrete symmetries (head/FFN permutations) are initialized via Weight Matching and kept fixed during the learning phase to ensure efficiency.

3. Dual Learning Procedure: Unlike prior "single-sided" approaches that optimize one model ($\Theta_B$) to match a fixed target ($\Theta_A$), LMC-DM introduces a bidirectional optimization. Both models learn their own transformation parameters ($\pi_A$ and $\pi_B$) toward a shared linear interpolation path:
   $$\Theta_{interp} = \lambda \pi_A(\Theta_A) + (1 - \lambda) \pi_B(\Theta_B)$$
   The symmetry parameters are optimized directly to minimize the loss along the interpolation path, allowing both models to move toward a shared low-loss basin rather than forcing one to conform to the other's coordinate system.

Key Contributions

• Comprehensive Symmetry Formulation: The paper defines a broad family of functionality-preserving transformations for Transformers, including normalization absorption, residual-space rotations, and specific attention/FFN circuit symmetries.
• Dual Learned Matching: The introduction of a bidirectional optimization framework where both endpoints are transformed, resulting in a better-conditioned optimization landscape and lower interpolation barriers compared to one-sided methods.
• Scalability Demonstration: The successful application of LMC-based merging to models with up to billions of parameters, a scale previously unexplored for near-barrier-free connectivity.

Experimental Results
The authors evaluate LMC-DM on Vision Transformers (ViT-S, ViT-B, ViT-L) and Language Models (Pythia and OLMo) ranging from 14M to 6.9B parameters.

• Loss Barriers: On WikiText, LMC-DM achieves near-zero loss barriers for 160M parameter models (0.18), significantly outperforming Weight Matching (0.27) and single-sided Learned Matching (0.27). For 1B parameter models, the barrier remains low (0.23).
• Vision Domain: ViT-L maintains over 69% ImageNet top-1 accuracy across the entire interpolation path, a substantial improvement over naive interpolation (which drops significantly) and single-sided methods.
• Scaling Behavior: While Pythia models up to 6.9B parameters exhibit small loss barriers, OLMo-7B shows higher barriers, suggesting that LMC in large models is influenced by factors beyond just scale, such as training data and optimization dynamics.
• Ablation Studies: Removing QK or OV symmetries significantly increases loss barriers, particularly for ViTs. The dual-learning approach consistently outperforms single-sided variants, and Cayley-based parameterization for orthogonal symmetries proves more effective than unconstrained or matrix-exponential approaches.

Significance and Claims
The paper claims to provide the first demonstration of near-barrier-free linear connectivity for language models with medium-sized parameters (160M) and to show that billion-parameter LLMs can be connected with only small loss barriers when properly aligned.

The authors argue that their results suggest large pretrained Transformers are not isolated points in weight space but are "symmetry-related representations" that can be effectively aligned and merged through simple linear paths. By resolving parameter symmetries via dual optimization, the method enables reliable model merging and composition at scales previously thought to be intractable for linear mode connectivity. The work bridges the gap between theoretical loss landscape geometry and practical model reuse, offering a foundation for scalable composition of large pretrained models.