💬 NLP

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

InfiGFusion introduces a novel structure-aware model fusion framework that employs Graph-on-Logits Distillation with an efficient, closed-form Gromov-Wasserstein approximation to capture semantic dependencies between vocabulary channels, thereby significantly outperforming existing baselines in complex reasoning tasks.

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

Published 2026-05-25

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have three different expert chefs. One is a master of Italian pasta, another is a wizard with spicy Indian curries, and the third is a genius at baking French pastries. You want to create a single "Super Chef" who can cook all three cuisines perfectly without needing to hire three separate people.

This is the challenge of Model Fusion in Artificial Intelligence. The paper introduces a new method called InfiGFusion to solve this problem, specifically for Large Language Models (LLMs) like the ones that power chatbots.

Here is how the paper explains their solution, using simple analogies:

1. The Problem: The "Solo" Approach vs. The "Team" Approach

Most current methods for combining AI models are like asking the Super Chef to look at the ingredients for a pasta dish, then look at the ingredients for a curry, and decide on a flavor one ingredient at a time.

The Flaw: This ignores how ingredients interact. In a real recipe, garlic and lemon work together to create a specific taste. If you treat them as separate items, you miss the "flavor profile" or the semantic dependency.
The Paper's View: Existing methods look at the AI's "thoughts" (called logits) word-by-word. They say, "The AI thinks 'cat' is likely, and 'dog' is likely," but they don't understand that the AI sees a relationship between "cat" and "dog" as "animals." They miss the structure of the thinking.

2. The Solution: Drawing a Map of Thoughts (Graph-on-Logits)

InfiGFusion changes the game by not looking at words individually, but by looking at how they connect.

The Analogy: Imagine the AI's thoughts aren't just a list of words, but a social network.
- Nodes (People): Each possible word the AI could say is a person.
- Edges (Friendships): If the AI often thinks about "fire" and "smoke" together, there is a strong friendship line between them.
The Innovation: Instead of just comparing the list of words the AI predicts, InfiGFusion compares the entire map of friendships. It asks: "Does the Super Chef's map of how 'fire' relates to 'smoke' look like the map from the Italian Chef and the Indian Chef?"

3. The Secret Sauce: The "Gromov-Wasserstein" Distance (The Shape Matcher)

To compare these complex maps, the authors use a mathematical tool called Gromov-Wasserstein (GW) distance.

The Analogy: Imagine you have two different shapes made of clay. One is a cube, one is a sphere. You want to know how similar they are.
- Old Way: You measure the height, width, and depth of every single point. This is slow and messy.
- GW Way: You look at the shape of the object. Does the cube have corners? Does the sphere have curves? You compare the relationships between the points, not just the points themselves.
The Problem with GW: Usually, this "shape matching" is incredibly slow and computationally expensive (like trying to solve a puzzle with a billion pieces). It would take too long to train the AI.

4. The Breakthrough: The "Sorting" Shortcut

The paper's biggest technical achievement is a new way to do this shape matching fast.

The Analogy: Instead of trying to match every single person in the social network one by one, they realized they could just rank everyone by how popular they are (how many connections they have).
- If the "Super Chef" has a "Fire" person who is the 5th most popular, and the "Source Chef" also has a "Fire" person who is the 5th most popular, they match them!
The Result: This "sorting" trick turns a task that used to take forever (O(n⁴)) into a task that is almost instant (O(n log n)). It's like turning a 10-hour hike into a 10-minute bike ride.

5. The Results: Better at "Thinking"

The authors tested this new "Super Chef" on 11 different challenges, including math, coding, and logic puzzles.

The Outcome: InfiGFusion was significantly better at complex reasoning than previous methods.
- Example: On a "Multi-step Arithmetic" test (solving a math problem that requires several steps), it improved by 35.6 points.
- Example: On "Causal Judgement" (figuring out if A caused B), it improved by 37 points.
Why? Because it preserved the relationships between ideas. It didn't just memorize the answer; it learned the logic of how the source models reasoned.

Summary

InfiGFusion is a new way to combine AI models. Instead of just copying the answers word-for-word, it copies the structure of the thinking. It treats the AI's predictions like a map of connections and uses a clever "sorting" trick to align these maps quickly. The result is a fused AI that is much smarter at solving complex, multi-step problems than models that only look at words in isolation.

Technical Summary: InfiGFusion

Problem Statement
Recent efforts to fuse heterogeneous Large Language Models (LLMs) into unified systems have largely relied on logit-based distillation methods. While these approaches maintain inference efficiency, they typically treat vocabulary dimensions independently, aligning token distributions via objectives like KL divergence or Wasserstein distance. This token-wise alignment overlooks semantic dependencies—the cross-dimension interactions and co-activation patterns that encode a model's internal reasoning processes. Consequently, existing methods struggle to align models with divergent generation behaviors, particularly in tasks requiring multi-step reasoning, relational inference, and fine-grained semantic alignment.

Methodology: InfiGFusion
The authors propose InfiGFusion, a structure-aware fusion framework that explicitly models these semantic dependencies. The core innovation is the Graph-on-Logits Distillation (GLD) loss, which operates as follows:

Logit Graph Construction: Instead of treating logits as independent vectors, the method constructs a global co-activation graph. For each output, the top- $k$ logits are retained to ensure sparsity. The outer products of these logits across sequence positions are aggregated to form a graph where nodes represent vocabulary channels and edges quantify their joint activations (similarity).
Structure-Aware Alignment: To align the logit graphs of source models with a pivot model, the framework employs the Gromov-Wasserstein (GW) distance. GW is chosen for its ability to match distributions based on internal structural similarity rather than direct node correspondence.
Efficient Approximation: Standard GW distance computation has a prohibitive $O(n^4)$ $O (n^{4})$ complexity, making it infeasible for large vocabularies ( $n > 10^5$ $n > 1 0^{5}$ ). The authors introduce a novel sorting-based closed-form approximation that reduces complexity to $O(n \log n)$ .
- Graph Feature Compression: The similarity matrix is compressed into scalar node-level features (specifically, row means or degree-style features).
- Sorting-Based Matching: The alignment is performed by sorting these features and matching the $i$ -th largest entry of the source to the $i$ -th largest of the target.
- Theoretical Guarantees: The paper provides a provable approximation error bound of $\frac{n-1}{n^2} + \frac{m-1}{m^2}$ , which decays as $O(1/n)$ . Additionally, the authors demonstrate that the GW-guided loss has a significantly lower Lipschitz constant compared to KL and standard Wasserstein losses, suggesting tighter generalization bounds and more stable training.
Unified Objective: The final training objective combines the structure-aware GLD loss with a coarse-grained Universal Logit Distillation (ULD) loss (for token-level distribution alignment) and standard Supervised Fine-Tuning (SFT) signals.

Key Contributions

First Structure-Aware Fusion Framework: InfiGFusion is the first method to model semantic dependencies via feature-level logit graphs and align them using a novel GLD loss, moving beyond independent token alignment.
Scalable GW Approximation: The authors propose a closed-form approximation for Gromov-Wasserstein distance that reduces computational complexity from $O(n^4)$ to $O(n \log n)$ with provable error bounds, making graph-based alignment feasible for LLM-scale vocabularies.
Empirical Superiority in Reasoning: Extensive experiments across 11 benchmarks demonstrate that InfiGFusion consistently outperforms state-of-the-art fusion baselines (e.g., MiniLogit, FuseLLM, FuseChat) and SFT baselines.

Results
The method was evaluated on 11 benchmarks spanning reasoning, coding, mathematics, and instruction following, fusing multiple source models (Qwen2.5, Mistral, Qwen2.5-Coder) into a Phi-4 pivot model.

Complex Reasoning Gains: InfiGFusion showed significant improvements in tasks requiring multi-step and relational inference. Notably, it achieved a +35.6 point improvement on Multistep Arithmetic and a +37.06 point improvement on Causal Judgement over SFT baselines.
Overall Performance: The model outperformed the pivot and source models across diverse domains, achieving an average score of 83.85 compared to the pivot's 79.54.
Efficiency: Despite the structural alignment, the method maintained competitive efficiency, requiring only 195 GPU hours for training, significantly less than multi-step fusion pipelines like FuseChat (650 GPU hours).
Ablation Studies: Experiments confirmed that both the ULD (coarse alignment) and GLD (structural alignment) components are necessary, with GLD specifically driving gains in reasoning tasks.

Significance
The paper claims that InfiGFusion addresses a critical gap in model fusion: the inability of token-level objectives to capture the relational structures underlying model reasoning. By explicitly aligning the "shape" of the logit space rather than just the marginal distributions, the framework enables the integration of heterogeneous LLMs without increasing inference costs. The authors posit that this structure-preserving alignment is essential for advancing collective LLM capabilities, particularly in complex reasoning scenarios where logical consistency and multi-hop inference are paramount. The work suggests that future progress in model fusion depends on moving beyond surface-level token matching to preserve the internal semantic dependencies that drive coherent reasoning.