Research map / Architectures and efficiency / Attention and transformers
Transformer Fusion with Optimal Transport
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends OTFusion to Transformers (ViT, BERT) via a transport-map flow graph, handling residuals, multi-head attention, LayerNorm and embeddings, and shows soft Sinkhorn alignment beats hard EMD alignment for one-shot fusion and fine-tuned fusion beats parents.
Problem
Model fusion (merging independently trained networks) has been developed for MLPs, CNNs, ResNets and RNNs, but layerwise OTFusion does not handle Transformer components such as multi-head self-attention, layer normalization, embeddings and residual streams.
Approach
A Transportation Map Flow Graph that describes how transport maps propagate through residual connections, attention, embeddings and normalization, combined with soft (Sinkhorn) alignment and cross-head alignment, enabling one-shot fusion and heterogeneous-width fusion of Transformers.
Key results
- On CIFAR10 two-model one-shot fusion, activations-based soft alignment reaches 60.87% vs 7.59% for vanilla fusion (parents ~92.3%).
- After finetuning on CIFAR100, weights-based soft-alignment fusion of 6 ViTs reaches 65.98% vs 63.19% for vanilla fusion and 64.94% best parent; ensemble of 6 reaches 70.71%.
- On ImageNet-1k, fused ViT after finetuning reaches 75.80% vs best parent 75.33% and vanilla fusion 67.83%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use soft Sinkhorn activation alignment instead of hard EMD alignment (ablation) | ViT OTFusion · CIFAR10 one-shot, 2 models | accuracy: 24.50 ± 5.66 -> 60.87 ± 0.44 | Helped |
| Use non-zero Sinkhorn regularization rather than hard alignment (ablation) | ViT OTFusion · CIFAR10 one-shot | accuracy | Helped |
| Use 6x6 center-window token filtering for activation alignment (ablation) | ViT activation-based OTFusion · CIFAR10 one-shot | accuracy | Helped |
| Use averaging for residual transportation maps instead of weighted residual policies (ablation) | ViT OTFusion · CIFAR10 one-shot | accuracy | Helped |
| Remove the equality constraint between query and key transport maps (ablation) | ViT soft OTFusion · One-shot accuracy | accuracy | Helped |
| Use soft Sinkhorn weight alignment instead of hard EMD alignment (ablation) | ViT OTFusion · CIFAR100 finetuned, 2 models | accuracy: 64.72 -> 65.80 | Helped |
| Use soft Sinkhorn weight alignment instead of hard EMD alignment (ablation) | ViT OTFusion · CIFAR100 finetuned, 6 models | accuracy: 65.42 -> 65.98 | Helped |
| Use activation-based alignment rather than weight-based alignment (ablation) | ViT OTFusion · CIFAR100 finetuned | accuracy: 65.80 vs 65.35; 65.98 vs 65.25 | Mixed |
| Increase the number of fused parent models from 2 to 6 (ablation) | ViT OT-wts followed by finetuning · CIFAR100 | accuracy: 65.80 -> 65.98 | Helped |
| Fuse independently trained models with soft OTFusion rather than vanilla parameter averaging | ViT · CIFAR10 one-shot, 2 models | accuracy: 7.59 -> 60.87 ± 0.44 | Helped |
| Fuse independently trained models with soft OTFusion rather than vanilla parameter averaging | ViT · CIFAR100 finetuned, 2 models | accuracy: 64.91 -> 65.80 | Helped |
| Fuse independently trained models with soft weight-based OTFusion rather than vanilla parameter averaging | ViT · ImageNet-1k finetuned, 2 models | accuracy: 67.83 -> 75.80 | Helped |
Surprising findings
- Unlike MLPs, CNNs and ResNets, hard (EMD) alignment underperforms soft Sinkhorn alignment for Transformers in one-shot fusion.
- Fusing a random-init model with a converged model still yields a 10% one-shot accuracy at fully soft alignment, while fusing two converged models gives 47%.
Limitations
- Heterogeneous fusion across different depths is not supported, inherited from layerwise OTFusion.
- Vanilla fusion and one-shot performance remain well below parents; gains rely on finetuning.
- Fused models underperform the ensemble in accuracy after finetuning.
Experimental setup
Datasets: CIFAR10 CIFAR100 Tiny ImageNet ImageNet-1k Wikipedia (simple subset) GLUE
Benchmarks: CIFAR10 CIFAR100 Tiny ImageNet ImageNet-1k GLUE
Metrics: accuracy MLM loss one-shot accuracy finetuned accuracy
Baselines: Vanilla Fusion Ensembling OTFusion with EMD (hard alignment) Individual parent models
Builds on: OTFusion Sinkhorn-Knopp Earth-Mover's Distance
Closest papers by meaning
- Model Fusion via Optimal Transport (NeurIPS 2020 · similarity 0.87)
- The Need for Speed: Pruning Transformers with One Recipe (ICLR 2024 · similarity 0.87)
- HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization (NeurIPS 2025 · similarity 0.86)
- Towards Meta-Pruning via Optimal Transport (ICLR 2024 · similarity 0.86)
- Slicing Vision Transformer for Flexible Inference (NeurIPS 2024 · similarity 0.85)
- An Image is Worth More Than 16x16 Patches: Exploring Transformers on Individual Pixels (ICLR 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.