Transport and Merge: Cross-Architecture Merging for Large Language Models

1. Problem Statement

Large Language Models (LLMs) achieve superior performance through massive scaling of parameters and data. However, real-world deployments often require smaller, resource-constrained models (e.g., for low-resource languages like Malay, Cantonese, or specialized domains like finance/medicine).

• The Gap: There is a significant performance gap between high-resource, large-scale models and smaller, low-resource targets.
• The Limitation of Existing Methods:
  • Gradient-based adaptation (SFT/Distillation): Requires massive training data and computational resources, which are often unavailable for low-resource scenarios.
  • Standard Model Merging: Existing merging techniques (e.g., Model Soups, TIES-Merging) assume the source and target models share the same architecture. They cannot directly merge weights between heterogeneous models (e.g., merging a 7B Qwen model into a 1B LLaMA model) because the weight dimensions and neuron structures do not align.
• Core Question: Can we enable high-resource to low-resource knowledge transfer across heterogeneous architectures through direct model merging without extensive retraining?

2. Methodology: Transport and Merge

The authors propose a framework based on Optimal Transport (OT) to align heterogeneous models in activation space and translate these alignments into weight-space fusion.

A. Core Insight

Despite architectural differences, language models processing the same inputs often exhibit correlated internal activations. The method leverages this by inferring cross-neuron correspondences between a source model ($M_B$) and a target model ($M_A$) using a small set of calibration inputs.

B. Step-by-Step Pipeline

1. Activation Extraction & Feature Alignment (Feature-Level OT)

• Input: A small dataset $D$ is passed through both the source and target models.
• Process: Intermediate activations are extracted from linear sublayers (Attention projections $Q, K, V$ and MLP projections).
• Cost Matrix: A correlation-based cost matrix $C_{\ell m}$ is constructed between target activation channels and source activation channels using Pearson correlation distance ($1 - \rho$).
• Optimization: An entropically regularized Optimal Transport problem is solved to find a transport plan $Q_{\ell m}$. This matrix represents a soft correspondence, indicating how target features can be reconstructed as weighted combinations of source features.
  $$Q_{\ell m} = \arg\min_{Q} \langle C_{\ell m}, Q \rangle - \epsilon H(Q)$$

2. Layer Alignment (Layer-Level OT)

• Process: The feature-level transport costs are aggregated to compute a layer-to-layer cost matrix $C_{layer}$.
• Optimization: A second OT problem is solved to find a global layer correspondence matrix $P$, determining which source layers best align with which target layers.
• Effective Correspondence: Pre-activation and post-activation transport plans are combined to form an effective layer correspondence $P_{eff}$.

3. Weight Fusion via Neuron Replacement

• Neuron Selection: To ensure robustness and avoid interference, the method selects the top-$k$ most activated neurons in the target model on the calibration set. Fusion is restricted only to these neurons.
• Weight Transport: The source weights are "transported" into the target's neuron basis using the learned transport maps ($Q_{in}, Q_{out}$) and layer correspondence ($P_{eff}$).
• Fusion Formula: The fused weights for a target layer $\ell$ are computed as a residual addition:
  $$W^{\ell, fused}_A = W^{\ell}_A + \alpha \cdot M_\ell \odot \left( \sum_{m} P_{eff}[\ell, m] Q^{out}_{\ell m} W^m_B (Q^{in}_{\ell m})^\top - W^{\ell}_A \right)$$
  Where $M_\ell$ is a mask for the top-$k$ neurons, and $\alpha$ controls the fusion strength.

4. Residual-Frozen Adaptation (Optional)

• After fusion, the transferred residual parameters ($\Delta W$) are frozen.
• Only the original base parameters of the target model are fine-tuned on the downstream task.
• Finally, the residual is folded back into the base weights for inference, resulting in a model with the original target architecture but enriched parameters.

3. Key Contributions

1. Cross-Architecture Merging Framework: The first method to enable direct weight-space fusion between LLMs of different architectures and capacities without distillation or full fine-tuning.
2. Optimal Transport for Neuron Alignment: A principled approach using OT to infer structured, neuron-level correspondences between heterogeneous models based on activation statistics.
3. Theoretical Interpretation: The paper provides a mathematical proof (Theorem 4.1) showing that weight transport is equivalent to mapping target features to source coordinates, applying the source linear map, and mapping back.
4. Efficiency: The method requires only a small calibration set (no massive training data) and performs fusion in a single pass, making it highly efficient for low-resource scenarios.

4. Experimental Results

The method was evaluated across low-resource languages (Malay, Indonesian, Thai, Cantonese) and specialized domains (Finance, Medical).

• Low-Resource Languages:
  • Malay (MalayMMLU): The fused model with adaptation improved accuracy by +6.51% in Humanities and +5.99% in Social Science over the base model.
  • Cantonese (CMMLU): Achieved a +2.18% overall improvement.
  • Thai: Significant gains in MGSM (math reasoning) from 0.50 to 0.72 with adaptation.
• Domain Transfer (Finance & Medical):
  • Consistent improvements on MMLU subsets (Business Ethics, Microeconomics, Anatomy, Medical Genetics).
  • Fusion served as an effective initialization, often outperforming the base model even without post-fusion adaptation.
• Comparison with SFT and Distillation:
  • The proposed method consistently outperformed Supervised Fine-Tuning (SFT) and Knowledge Distillation across all domains, particularly in scenarios where data is scarce.
• Robustness:
  • The method is robust to different source model backbones (e.g., LLaMA-3 vs. Qwen2.5).
  • Performance scales with source model size (merging a 32B source yields better results than an 8B source).
  • General Capabilities: The method preserves the target model's general reasoning abilities, showing no significant degradation in general benchmarks (ARC, CommonsenseQA, etc.).

5. Significance

• Democratizing LLMs: This approach allows organizations to leverage the knowledge of massive, high-resource models to boost smaller, specialized models without the prohibitive cost of training or distillation.
• Solving the Heterogeneity Bottleneck: It breaks the constraint that merging requires identical architectures, opening the door to mixing and matching models from different families (e.g., Qwen, LLaMA, Mistral) for specific tasks.
• Practical Deployment: By requiring only a small calibration set and offering a "plug-and-play" fusion mechanism, it provides a practical solution for deploying high-performance LLMs in resource-constrained environments (e.g., edge devices, low-resource languages).

In summary, Transport and Merge establishes a new paradigm for knowledge transfer, moving from gradient-based adaptation to activation-aligned weight fusion, effectively bridging the gap between large-scale foundation models and specialized, low-resource applications.