InfiGFusion: Graph-on-Logits Distillation for Scalable Model Fusion

🚀 Executive Snapshot

Notation. Parenthetical deltas (e.g., (+6.0↑) / (-0.8↓)) show absolute change versus the Phi-4 (14B) baseline. Bold indicates the best score in the row.

InfiGFusion (14B) achieves the highest average across 18 tasks, improving +6.73 points over the strongest SFT baseline (Phi-4 SFT).

Interpretation.

• Against a strong SFT baseline (Phi-4 SFT, 14B), InfiGFusion (14B) boosts average accuracy from 77.94 → 84.16 (+6.73).
• Gains are especially pronounced on tasks emphasizing multi-step reasoning and structure-sensitive signals (e.g., Multistep Arithmetic, Dyck Languages, Causal Judgement).
• On a few tasks (e.g., Formal Logic), specialized systems (e.g., DeepSeek-R1-Distill-Qwen) may still lead, illustrating that teacher diversity/selection remains a complementary lever—InfiGFusion cleanly plugs into such strategies.

📌 At-a-Glance

InfiGFusion is a unified framework for multi-source model fusion built on two complementary loss:

• Universal Logit Distillation (ULD): token-level distribution alignment between teacher(s) and student.
• Graph-on-Logits Distillation (GLD): structural alignment by constructing graphs on logits and aligning them with approximated Gromov–Wasserstein (GW). We introduce a sorting-based closed-form approximation that reduces complexity from O(n⁴) to O(n log n), together with stability and error analyses.

Result: fuse knowledge from multiple experts without teacher weights and without changing the student architecture, delivering consistent improvements on cross-domain, muti-step reasoning, and other complex reasonging tasks with limited training cost.

🧩 Motivation: Distribution Alone Isn’t Enough

Token-level objectives (e.g., KL/CE) align marginals but ignore the geometry between classes or semantic dimensions. In multi-expert fusion, complementarity and conflicts abound:

• Which classes (or semantic axes) are close or distant?
• How should we preserve relative structure across teachers?

Key idea: logits are not just probabilities—they encode relational structure. Aligning the internal geometry (who is similar to whom, how strongly) is crucial for robust fusion.

🛠️ Method

InfiGFusion framework overview.

ULD (Universal Logit Distillation) — Token-Level Alignment

• A unified, multi-teacher friendly objective that aligns p(y|x) distributions.
• Provides a strong, stable training signal.

GLD (Graph-on-Logits Distillation) — Structural Alignment on Logits

• Build sparse graphs over logits (e.g., top-k neighbors or thresholded edges).
• Use GW to align pairwise distance structures between teacher and student graphs.
• Sorting-based closed-form approximation: replaces intractable O(n⁴) GW with O(n log n) while maintaining training usefulness; we also provide error bounds and stability guarantees.

Intuition: ULD makes the student look like the source(s); GLD makes the student’s relationships look like the source(s). The two are complementary.

🧪 Why Structure Helps (GLD vs. KL)

Comparison between GLD and KL divergence on logits.

• KL/CE/ULD only: may over-concentrate on a few dominant classes and miss inter-class geometry.
• Add GLD: preserves relative structure, amplifying cross-teacher complementarity and improving long-tail robustness and reasoning-heavy tasks.

⚖️ Case Study: Cross-Domain Complementarity & Conflicts

Cross-domain relations heatmap showing complementarity and conflicts.

• Domains such as math, physics, multi-step reasoning show structural proximity on sub-tasks.
• GLD helps amplify complementarities and mitigate conflicts, yielding more consistent fused models.

📊 Simulation Study for the GW Approximation

We conducted a controlled simulation: sample n ∈ {50, 100, 200, 500, 1000, 1500, 2000, 2500, 3000} points from N(0, I) and N(1, I), build the two n×n Euclidean distance matrices, and compare:

• Exact GW
• Sinkhorn GW (widely-used, high-quality approximation)
• Approx GW (ours) (sorting-based closed-form estimator)

Key takeaways:

• Estimator quality: Sinkhorn matches Exact (RE ≈ 0). Approx GW’s relative error decreases with n (≈ O(1/n)): from ~0.59 @ n=50 to ~0.33 @ n=3000.
• Runtime: Exact/Sinkhorn quickly become impractical beyond ~n=50 (seconds → tens of seconds); Approx stays sub-second even at n=3000.
• Conclusion: Our Approx GW is designed as a training objective that balances scalability and fidelity—precisely what KD/fusion scenarios require at scale.

@article{wang2025infigfusion,
  title={InfiGFusion: Graph-on-Logits Distillation via Efficient Gromov-Wasserstein for Model Fusion},
  author={Wang, Yuanyi and Yan, Zhaoyi and Zhang, Yiming and Zhou, Qi and Gu, Yanggan and Wu, Fei and Yang, Hongxia},
  journal={arXiv preprint arXiv:2505.13893},
  year={2025}
}