Optimal Transport Model Fusion

• The paper introduces OT-based fusion that aligns disparate model components via Wasserstein metrics to achieve superior integration compared to naive averaging.
• It details layer-wise alignment for neural network architectures, enabling robust fusion across modalities and handling neuron permutation challenges.
• Empirical results show that OT fusion improves accuracy and efficiency in multimodal tasks, sensor arrays, and model compression without extra inference cost.

Model fusion via optimal transport (OT) is a principled methodology for combining the parameters, intermediate representations, or predictions of multiple distinct machine learning models into a single, unified model. The approach leverages the mathematical machinery of optimal transport—particularly Wasserstein-type distances and entropic regularized variants—to softly align disparate components (neurons, features, or distributions) across models before fusing them. This formulation yields significantly better alignment than naive parameter averaging or stacking, unlocking robust fusion across architectures, modalities, and sensor types. The OT fusion paradigm has been instantiated for supervised neural networks, self-supervised learners, multi-modal systems, sensor arrays, and generative models.

1. Mathematical Foundations of Optimal Transport Fusion

At the core of model fusion via optimal transport is the alignment of two (or more) discrete probability measures representing model components. For two finite sets $\{x_i\}_{i=1}^n$ and $\{y_j\}_{j=1}^m$ with probability weights $a \in \Delta_n$, $b \in \Delta_m$ ($\Delta_k$ is the $k$-simplex), and ground costs $C_{ij}$, OT seeks a transport plan $T \in \mathbb{R}_+^{n\times m}$:

$$\min_{T\ge0} \sum_{i=1}^n \sum_{j=1}^m C_{ij} T_{ij} \quad \text{s.t.} \quad T\mathbf{1}_m = a,\, T^\top \mathbf{1}_n = b$$

The cost matrix $C_{ij}$ is typically the squared Euclidean or $\{y_j\}_{j=1}^m$0 norm between parameter vectors, channels, or feature representations. Entropic regularization adds a term $\{y_j\}_{j=1}^m$1 for $\{y_j\}_{j=1}^m$2, yielding a strictly convex problem solvable by Sinkhorn iterations in $\{y_j\}_{j=1}^m$3 per iteration (Singh et al., 2019, Imfeld et al., 2023).

A single-layer OT fusion between two neural nets aligns rows or columns using the optimal plan $\{y_j\}_{j=1}^m$4, transforms the weight matrix of one model, and computes a barycentric average:

$\{y_j\}_{j=1}^m$5

Layer-wise recurrence or indirect propagation through previous alignment plans (for Transformers, GCNs) enables global consistency (Imfeld et al., 2023, Ormaniec et al., 27 Mar 2025).

2. Layer-wise Alignment and Fusion in Deep Networks

Layer-wise OT-based fusion overcomes the neuron permutation symmetry problem inherent in deep nets. Naive weight averaging fails in the presence of different random initializations or data splits, since it does not account for potential neuron orderings. For convolutional, residual, and fully-connected nets, OT-based fusion is performed independently per layer by computing a ground cost between rows (neurons/filters) or activation vectors, solving for the alignment plan, and averaging aligned weights. The approach extends naturally to multi-layer perceptrons, CNNs, ResNets, and VGG variants (Singh et al., 2019, Ormaniec et al., 27 Mar 2025).

For graph neural networks, the cost functions include Euclidean feature distance, quadratic energy (graph-based diffusive distance on activations), and fused Gromov-Wasserstein; performance is best for naive Euclidean cost, since structure-aware costs introduce noise and computational burden (Ormaniec et al., 27 Mar 2025). In Transformers, fusion generalizes to self-attention weights, cross-head alignments, layer norm, and positional embeddings; alignment is propagated using a transportation-map flow graph that respects branching and residual connections (Imfeld et al., 2023). Fusing models with disparate widths is accomplished via rectangular transport plans, projecting larger models into the target model’s latent space.

3. Entropy-Regularized Barycenter and Sensor Fusion

The entropy-regularized Wasserstein barycenter problem fuses probabilistic spatial or spectral estimates under partial constraints. In the context of non-coherent sensor fusion and multi-array localization, the barycenter formulation seeks a "best" spectrum $\{y_j\}_{j=1}^m$6 that is simultaneously consistent with the partial (covariance-based) information from all $\{y_j\}_{j=1}^m$7 arrays:

$\{y_j\}_{j=1}^m$8

Here, $\{y_j\}_{j=1}^m$9 is the entropic Wasserstein cost, $a \in \Delta_n$0 the discretized forward operators, $a \in \Delta_n$1 the vectorized covariances, and $a \in \Delta_n$2 fit errors. The fusion objective admits globally-convergent block coordinate (Sinkhorn-Newton) ascent over dual variables. This framework is robust to sensor misalignments and spatial miscalibration, as the Wasserstein geometry penalizes mass movement smoothly across the domain, unlike $a \in \Delta_n$3 or $a \in \Delta_n$4 metrics (Elvander et al., 2018).

The barycenter formalism generalizes to arbitrary nonnegative measures—topic distributions, color histograms, or time-series spectra may be simultaneously fused via the same machinery.

4. Multimodal and Cross-Architecture Applications

Model fusion via OT applies to diverse modalities and architectural pairings:

• Audio-visual fusion: LAVCap fuses audio and visual encoder outputs by solving an entropic OT alignment using a pairwise similarity matrix, enforcing semantic token-wise correspondence and driving a subsequent OT-attention fusion module. The fusion improves sample efficiency over vanilla concatenation and attention, and enables LLM-based captioning to outperform prior audio-visual methods (Rho et al., 16 Jan 2025).
• Multimodal traversable area segmentation: OT-Drive fuses RGB and surface-normal modalities in autonomous driving by transporting both distributions onto language-grounded anchor prototypes, using scene attributes as anchor indices. This yields robust generalization for out-of-distribution navigation (Zhao et al., 15 Jan 2026).
• Heterogeneous model fusion: OT-based alignment enables merging of supervised and self-supervised ASR models, parallel-branch fusion of distinct self-supervised PTMs (e.g. Mamba and attention-based models for speech emotion), and merging task-specific adapters into a backbone for continual multi-task learning (Pan et al., 24 Nov 2025, Fu et al., 2023, Phukan et al., 1 Jun 2025).
• Structured pruning and compression: Intra-Fusion replaces hard pruning by fusing all neurons in a layer into a lower-dimensional representation using OT, preserving function and maintaining accuracy without retraining (Theus et al., 2024).

5. Algorithmic and Computational Considerations

The computational bottleneck is the OT solver per layer or feature set. Entropic regularization permits scalable Sinkhorn iteration with complexity $a \in \Delta_n$5 per layer for $a \in \Delta_n$6 Sinkhorn steps. For very wide layers or high-dimensional features, block-sparse or low-rank acceleration is feasible. In high-dimensional localization or 3D fusion, greedy coordinate descent can amortize cost by exploiting sparsity and local quadratic structure (Jaouedi et al., 31 Mar 2026).

Hyperparameters include entropy weight $a \in \Delta_n$7 (determining softness of alignment), number of iterations, choice of cost metric, and for unbalanced or multimodal fusion, relaxation parameters weighting mass mismatch or scene factors.

No additional inference-time cost is accrued—fusion is performed offline, and the fused model matches the computational profile of a single parent. Post-fusion fine-tuning of the resulting network closes the residual performance gap to (or occasionally surpasses) the ensemble or best parent.

6. Empirical Performance and Robustness

OT-based fusion consistently outperforms vanilla parameter averaging, prediction ensembling (at lower resource cost), and most baseline fusion methods across supervised, self-supervised, multimodal, and pruning settings. For example:

• In VGG/ResNet fusion on CIFAR10, OT fusion (no fine-tuning) achieves 86–77% accuracy vs. 17–18% for vanilla averaging (Singh et al., 2019).
• In speech ASR, OT fusion reduces word and character error rates by 8–9% relative to single or ensemble models (Fu et al., 2023).
• In pruning and compression, OT fusion recovers 10–30 percentage points of performance with no retraining, and up to +60 pp in extreme sparsity regimes (Theus et al., 2024).
• For multi-modal AV captioning, OT-attention fusion delivers state-of-the-art metrics (CIDEr 84.9, SPIDEr 51.7), outperforming previous attention or joint-encoder techniques (Rho et al., 16 Jan 2025).
• In off-road segmentation, OT-Drive enhances OOD mIoU by +6.35 points over prior best, with robust cross-dataset transfer (Zhao et al., 15 Jan 2026).

OT approaches maintain robustness to input or sensor misalignments, as transport plans distribute mass smoothly rather than enforcing rigid, brittle correspondences. Empirical ablations confirm the gains arise from the alignment step, not the specific choice of scoring or importance metric.

7. Extensions, Limitations, and Open Challenges

Key extensions include:

• Heterogeneous-architecture fusion (cross-width, partially aligned layers) (Imfeld et al., 2023).
• Unbalanced/multimarginal transport for sensor fusion and partial mass alignment (Jaouedi et al., 31 Mar 2026).
• Continual and incremental fusion of task-specific adapters, using feature-space OT plans and learned masks to prevent catastrophic forgetting (Pan et al., 24 Nov 2025).
• Fusion of generative models, graph neural networks, and vision/language transformers.

Limitations comprise cubic scaling in layer width for exact OT, non-convex interaction across layers (layerwise greedy alignment may be suboptimal globally), need for exact or approximate correspondence in layer or feature shape, and sensitivity to the choice of ground cost or entropic parameter. Fusion of data-heterogeneous or structurally dissimilar models sometimes requires auxiliary projection layers or post-fusion fine-tuning. Approximate and block-sparse OT solvers, improved neuron importance metrics, and extension to variable-depth aggregation remain active areas.

OT-based model fusion constitutes a theoretically grounded, empirically robust, and computationally tractable strategy for synthesizing high-performing, resource-efficient, and robust models from heterogeneous, multimodal, and structurally diverse sources. It yields a general paradigm that subsumes classical parameter fusion, sensor fusion, multimodal alignment, and model compression within a unified geometric-transport framework (Elvander et al., 2018, Singh et al., 2019, Imfeld et al., 2023, Theus et al., 2024, Ormaniec et al., 27 Mar 2025, Pan et al., 24 Nov 2025, Rho et al., 16 Jan 2025, Zhao et al., 15 Jan 2026, Phukan et al., 1 Jun 2025, Fu et al., 2023, Jaouedi et al., 31 Mar 2026).