Probabilistic Token Alignment for Large Language Model Fusion

📄 arXiv: 2509.17276v1 📥 PDF

作者: Runjia Zeng, James Chenhao Liang, Cheng Han, Zhiwen Cao, Jiahao Liu, Xiaojun Quan, Yingjie Victor Chen, Lifu Huang, Tong Geng, Qifan Wang, Dongfang Liu

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-09-21

备注: NeurIPS 2025

💡 一句话要点

提出基于概率Token对齐的PTA-LLM，用于解决大语言模型融合中的词汇对齐问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型融合 Token对齐 最优传输 概率模型 知识迁移

📋 核心要点

  1. 现有大语言模型融合方法依赖于手动预定义的词汇对齐，泛化性差，导致性能下降。
  2. 论文提出概率Token对齐（PTA-LLM），将Token对齐转化为最优传输问题，实现分布感知的模型融合。
  3. 实验结果表明，PTA-LLM能够提升目标模型在多种能力上的性能，并提供Token对齐的分布视角解释。

📝 摘要（中文）

从头训练大型语言模型（LLM）可以产生具有独特功能和优势的模型，但成本高昂且通常导致能力冗余。一种更具成本效益的替代方案是将现有的、具有不同架构的预训练LLM融合为更强大的模型。然而，现有模型融合的一个关键挑战是它们依赖于手动预定义的词汇对齐，这可能无法很好地推广到不同的上下文中，从而导致多个评估中的性能下降。为了解决这个问题，我们从分布学习中汲取灵感，并提出概率token对齐方法作为一种通用且软的对齐映射，命名为PTA-LLM。我们的方法创新性地将token对齐重新定义为一个经典的数学问题：最优传输，无缝地利用分布感知学习来促进更连贯的模型融合。除了其固有的通用性之外，PTA-LLM还从分布的角度展示了解释性，从而提供了对token对齐本质的见解。经验结果表明，概率token对齐提高了目标模型在多种能力上的性能。我们的代码可在https://runjia.tech/neurips_pta-llm/上找到。

🔬 方法详解

问题定义：现有的大语言模型融合方法依赖于手动定义的词汇对齐，这种对齐方式缺乏灵活性，难以适应不同架构和训练数据的大语言模型。当模型词汇表存在显著差异时，手动对齐可能导致次优的融合效果，最终影响融合后模型的性能。因此，如何自动且有效地进行词汇对齐是大语言模型融合的关键问题。

核心思路：论文的核心思路是将Token对齐问题转化为一个最优传输问题。最优传输是一种在概率分布之间寻找最佳匹配的数学框架。通过将源模型和目标模型的Token表示视为概率分布，可以利用最优传输理论找到一个软对齐方案，该方案能够捕捉Token之间的语义相似性，而不仅仅是简单的词汇匹配。这种方法允许模型在融合过程中更好地利用源模型的知识。

技术框架：PTA-LLM的整体框架包括以下几个主要步骤：1) Token嵌入提取：从源模型和目标模型中提取Token的嵌入表示。2) 概率分布构建：将每个Token的嵌入表示视为一个概率分布。3) 最优传输计算：利用最优传输算法计算源模型和目标模型Token之间的传输矩阵，该矩阵表示Token之间的对齐概率。4) 模型融合：利用计算得到的传输矩阵，将源模型的知识迁移到目标模型中，例如通过调整目标模型的参数或logits。

关键创新：该论文的关键创新在于将Token对齐问题重新定义为最优传输问题，并提出了一种概率Token对齐方法（PTA-LLM）。与传统的手动对齐方法相比，PTA-LLM能够自动学习Token之间的软对齐关系，从而更好地适应不同架构和训练数据的大语言模型。此外，PTA-LLM还提供了一种从分布角度解释Token对齐的视角，有助于理解模型融合的本质。

关键设计：在最优传输计算中，可以使用不同的距离度量来衡量Token嵌入之间的相似性，例如余弦相似度或欧氏距离。传输矩阵可以通过Sinkhorn算法等方法进行求解。在模型融合阶段，可以使用不同的融合策略，例如线性插值或知识蒸馏。损失函数的设计需要考虑对齐的准确性和融合后模型的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PTA-LLM在多个任务上优于基线方法，例如在知识问答、文本生成等任务上取得了显著的性能提升。与手动对齐方法相比，PTA-LLM能够更好地捕捉Token之间的语义关系，从而实现更有效的模型融合。具体性能提升幅度取决于任务类型和模型架构，但总体趋势表明PTA-LLM具有较强的泛化能力。

🎯 应用场景

该研究成果可应用于各种需要融合不同大语言模型的场景，例如：将擅长不同任务的模型融合以构建更通用的模型；利用小模型的能力增强大模型的特定功能；以及在资源受限的环境下，通过融合多个小模型来提升整体性能。该方法有望降低大语言模型训练和部署的成本，并促进更高效的模型开发。

📄 摘要（原文）

Training large language models (LLMs) from scratch can yield models with unique functionalities and strengths, but it is costly and often leads to redundant capabilities. A more cost-effective alternative is to fuse existing pre-trained LLMs with different architectures into a more powerful model. However, a key challenge in existing model fusion is their dependence on manually predefined vocabulary alignment, which may not generalize well across diverse contexts, leading to performance degradation in several evaluation. To solve this, we draw inspiration from distribution learning and propose the probabilistic token alignment method as a general and soft mapping for alignment, named as PTA-LLM. Our approach innovatively reformulates token alignment into a classic mathematical problem: optimal transport, seamlessly leveraging distribution-aware learning to facilitate more coherent model fusion. Apart from its inherent generality, PTA-LLM exhibits interpretability from a distributional perspective, offering insights into the essence of the token alignment. Empirical results demonstrate that probabilistic token alignment enhances the target model's performance across multiple capabilities. Our code is avaliable at https://runjia.tech/neurips_pta-llm/.