InfiGFusion: Graph-on-Logits Distillation via Efficient Gromov-Wasserstein for Model Fusion

📄 arXiv: 2505.13893v1 📥 PDF

作者: Yuanyi Wang, Zhaoyi Yan, Yiming Zhang, Qi Zhou, Yanggan Gu, Fei Wu, Hongxia Yang

分类: cs.CL

发布日期: 2025-05-20

💡 一句话要点

InfiGFusion提出一种基于logits图蒸馏的高效Gromov-Wasserstein模型融合方法，提升模型融合质量。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 模型融合 知识蒸馏 图神经网络 Gromov-Wasserstein距离 大型语言模型

📋 核心要点

  1. 现有基于logits的模型融合方法忽略了词汇维度间的语义依赖，限制了融合模型的性能。
  2. InfiGFusion通过logits图蒸馏（GLD）损失显式建模词汇间的语义依赖，提升模型融合质量。
  3. 实验表明，InfiGFusion在推理、编码和数学任务上优于SOTA模型，尤其在复杂推理任务上提升显著。

📝 摘要（中文）

大型语言模型（LLMs）的最新进展推动了将异构开源模型融合为统一系统的努力，以继承它们互补的优势。现有的基于logits的融合方法保持了推理效率，但独立处理词汇维度，忽略了跨维度交互编码的语义依赖性。这些依赖性反映了token类型在模型内部推理下的交互方式，对于对齐具有不同生成行为的模型至关重要。为了显式地建模这些依赖性，我们提出了InfiGFusion，这是第一个具有新颖的logits图蒸馏（GLD）损失的结构感知融合框架。具体来说，我们保留每个输出的前k个logits，并聚合它们在序列位置上的外积，以形成全局共激活图，其中节点表示词汇通道，边量化它们的联合激活。为了确保可扩展性和效率，我们设计了一种基于排序的闭式近似，将Gromov-Wasserstein距离的原始O(n^4)成本降低到O(n log n)，并具有可证明的近似保证。跨多个融合设置的实验表明，GLD始终提高融合质量和稳定性。InfiGFusion在涵盖推理、编码和数学的11个基准测试中优于SOTA模型和融合基线。它在复杂的推理任务中表现出特别的优势，在多步算术上比SFT提高了+35.6，在因果判断上提高了+37.06，证明了卓越的多步和关系推理能力。

🔬 方法详解

问题定义：现有基于logits的模型融合方法通常独立处理词汇表中的每个token，忽略了它们之间的语义关系。这种忽略导致融合后的模型无法充分利用不同模型的优势，尤其是在需要复杂推理的任务中，模型性能提升有限。现有方法计算效率高，但融合质量有待提高。

核心思路：InfiGFusion的核心思路是通过构建logits图来显式地建模词汇之间的语义依赖关系。该方法认为，词汇之间的共激活模式蕴含了模型内部的推理过程，通过学习这些模式，可以更好地对齐不同模型的生成行为，从而提升融合模型的性能。

技术框架：InfiGFusion框架主要包括以下几个步骤：1) 从每个模型的logits输出中选择top-k个logits；2) 计算这些logits在序列位置上的外积，构建全局共激活图，其中节点代表词汇，边代表它们之间的联合激活强度；3) 使用Gromov-Wasserstein距离衡量不同模型共激活图之间的差异；4) 通过logits图蒸馏（GLD）损失函数，最小化这些差异，从而实现模型融合。

关键创新：InfiGFusion的关键创新在于提出了Graph-on-Logits Distillation (GLD)损失，它通过Gromov-Wasserstein距离来衡量不同模型logits共激活图之间的差异，从而显式地建模了词汇之间的语义依赖关系。此外，为了提高计算效率，论文设计了一种基于排序的闭式近似方法，将Gromov-Wasserstein距离的计算复杂度从O(n^4)降低到O(n log n)。

关键设计：关键设计包括：1) 使用top-k logits来构建共激活图，以减少计算量并保留最重要的信息；2) 设计了基于排序的闭式近似方法来高效计算Gromov-Wasserstein距离，保证了算法的可扩展性；3) GLD损失函数的设计，它鼓励融合后的模型学习不同模型的共激活模式，从而提升融合质量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

InfiGFusion在11个基准测试中优于SOTA模型和融合基线，尤其在复杂推理任务中表现出色。在多步算术任务上，InfiGFusion比SFT提高了+35.6，在因果判断任务上提高了+37.06。这些结果表明，InfiGFusion能够有效提升模型的多步和关系推理能力。

🎯 应用场景

InfiGFusion可应用于各种需要融合多个大型语言模型的场景，例如构建具有特定领域知识的专家系统、提升通用模型的推理能力、以及在资源受限的环境中部署高性能模型。该方法能够有效整合不同模型的优势，提高模型性能，降低部署成本，具有广泛的应用前景。

📄 摘要（原文）

Recent advances in large language models (LLMs) have intensified efforts to fuse heterogeneous open-source models into a unified system that inherits their complementary strengths. Existing logit-based fusion methods maintain inference efficiency but treat vocabulary dimensions independently, overlooking semantic dependencies encoded by cross-dimension interactions. These dependencies reflect how token types interact under a model's internal reasoning and are essential for aligning models with diverse generation behaviors. To explicitly model these dependencies, we propose \textbf{InfiGFusion}, the first structure-aware fusion framework with a novel \textit{Graph-on-Logits Distillation} (GLD) loss. Specifically, we retain the top-$k$ logits per output and aggregate their outer products across sequence positions to form a global co-activation graph, where nodes represent vocabulary channels and edges quantify their joint activations. To ensure scalability and efficiency, we design a sorting-based closed-form approximation that reduces the original $O(n^4)$ cost of Gromov-Wasserstein distance to $O(n \log n)$, with provable approximation guarantees. Experiments across multiple fusion settings show that GLD consistently improves fusion quality and stability. InfiGFusion outperforms SOTA models and fusion baselines across 11 benchmarks spanning reasoning, coding, and mathematics. It shows particular strength in complex reasoning tasks, with +35.6 improvement on Multistep Arithmetic and +37.06 on Causal Judgement over SFT, demonstrating superior multi-step and relational inference.