Research map / Large language models / Language, safety and interpretability
Probabilistic Token Alignment for Large Language Model Fusion
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.
Proposes PTA-LLM, which fuses pre-trained LLMs with different tokenizers by aligning their token distributions via dynamic token pairing and Sinkhorn-based optimal transport, then distilling into a Llama-2 7B target model.
Problem
Knowledge fusion of heterogeneous LLMs depends on manually predefined, hard token alignment between different tokenizers, which generalizes poorly and degrades performance on some tasks.
Approach
Reformulates token-level logit alignment as an optimal transport problem with an edit-distance-plus-logit cost, using dynamic token pairing followed by Sinkhorn soft alignment, so that the fused target distribution is a soft, distribution-aware mapping rather than a hard vocabulary mapping.
Key results
- PTA-LLM averages 29.00 across six benchmarks (78 tasks), a 1.72% relative gain over FuseLLM (28.51) and above Llama-2 CLM (27.90).
- On MultiPL-E, PTA-LLM reaches 15.88 vs. 15.56 for FuseLLM (+2.06%); on TyDi QA it reaches 34.07 vs. 32.99 (+3.27%).
- On the stability analysis of tasks where FuseLLM degrades, PTA-LLM reduces average degradation relative to Llama-2 from -13.04% to -4.07%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace FuseLLM's hard token mapping with probabilistic optimal-transport alignment | Llama-2 7B · Avg. 6 Benchmarks (78 tasks) | average score: 28.51 -> 29.00 | Helped |
| replace FuseLLM's hard token mapping with probabilistic optimal-transport alignment | Llama-2 7B · MultiPL-E | pass@1: 18.11 -> 15.88 | Hurt |
| replace FuseLLM's hard token mapping with probabilistic optimal-transport alignment | Llama-2 7B · BBH | score: 41.01 -> 41.08 | Helped |
| use threshold 1e-5 instead of 1e-4 (ablation) | PTA-LLM · BBH | score: 40.54 -> 41.08 | Helped |
| use threshold 1e-5 instead of 1e-4 (ablation) | PTA-LLM · ME | score: 15.88 -> 15.82 | Hurt |
| use threshold 1e-5 instead of 1e-4 (ablation) | PTA-LLM · MMLU | score: 48.99 -> 49.38 | Helped |
| use transport window size 5 instead of 10 (ablation) | PTA-LLM · BBH | score: 41.08 -> 40.68 | Hurt |
| use transport window size 5 instead of 10 (ablation) | PTA-LLM · ME | score: 15.88 -> 15.61 | Hurt |
| use transport window size 5 instead of 10 (ablation) | PTA-LLM · MMLU | score: 48.99 -> 49.38 | No clear effect |
| use combination weight 0.8 instead of 0.9 (ablation) | PTA-LLM · BBH | score: 40.39 -> 41.08 | Helped |
| use combination weight 0.8 instead of 0.9 (ablation) | PTA-LLM · ME | score: 15.72 -> 15.88 | Helped |
| use combination weight 0.8 instead of 0.9 (ablation) | PTA-LLM · MMLU | score: 48.93 -> 49.38 | Helped |
Surprising findings
- Fusing lower-performing source models (e.g., MPT on MMLU) yields the smallest gains, and the BBH gain over FuseLLM is marginal (+0.17%), attributed to weak source models.
Limitations
- Optimal transport step adds computational overhead; the authors suggest exploring faster alternatives such as Greenkhorn.
- Gains are modest on some benchmarks (e.g., BBH +0.17% over FuseLLM) and depend on quality of source models.
- Transport is performed in softmax-probability space, and the authors note that logit-space transport differs fundamentally and is left for future work.
Experimental setup
Datasets: MiniPile The Pile GitHub (HumanEval pilot)
Benchmarks: GSM8K BBH MultiPL-E MMLU ToxiGen TyDi QA HumanEval HellaSwag
Metrics: accuracy pass@1 exact match non-toxicity rate
Baselines: Llama-2 OpenLLaMA MPT Llama-2 CLM FuseLLM
Builds on: FuseLLM MinED token alignment Sinkhorn algorithm
Closest papers by meaning
- A Probabilistic Perspective on Unlearning and Alignment for Large Language Models (ICLR 2025 · similarity 0.90)
- Leveraging Importance Sampling to Detach Alignment Modules from Large Language Models (NeurIPS 2025 · similarity 0.89)
- Knowledge Fusion of Large Language Models (ICLR 2024 · similarity 0.89)
- Leveraging robust optimization for llm alignment under distribution shifts (NeurIPS 2025 · similarity 0.88)
- L3Ms — Lagrange Large Language Models (ICLR 2025 · similarity 0.88)
- Alignment of Large Language Models with Constrained Learning (NeurIPS 2025 · similarity 0.88)
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.