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

Abstract

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 InfiGFusion, the first structure-aware fusion framework with a novel 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.

Introduction

Recent advances in LLMs have sparked growing interest in aggregating diverse model capabilities to build stronger, more general-purpose systems. Existing collective approaches span a broad design space. Prompt-based techniques dynamically chain or compose pretrained LLMs with tools or APIs to expand functional scope. Multi-agent systems employ multiple LLMs to collaborate or debate, often outperforming single-agent models on complex reasoning tasks. Model ensembling aggregates outputs from different models to boost accuracy and robustness, while Mixture-of-Experts (MoE) architectures improve scalability by activating specialized subnetworks. Meanwhile, parameter-level merging models aim to consolidate multiple fine-tuned checkpoints into a single model. Each of these paradigms explores a different trade-off between performance, efficiency, and compatibility.

Among fusion strategies, model fusion via logit distillation has emerged as a flexible and efficient paradigm. It enables a single pivot model to absorb knowledge from multiple source models—often with diverse domains, sizes, or architectures—without increasing inference cost. A representative cross-tokenizer approach is ULD, which introduces a universal logit-distillation loss for aligning models with different tokenizers. Unlike traditional distillation, which transfers intermediate features or outputs, these methods align only the output logits using token-level objectives such as KL divergence or Wasserstein distance. However, token-level methods inherently treat each vocabulary dimension independently, overlooking semantic dependencies—the co-activation patterns and relational structures among tokens that underlie a model's reasoning process. This limits their ability to align models with divergent generation behaviors, particularly for tasks requiring multi-step reasoning or fine-grained relational inference.

The following figure illustrates this limitation. While token-level losses enforce dimension similarity, they fail to capture the structured dependencies among tokens. InfiGFusion addresses this by modeling logits as graphs, aligning not just distributions but also the interactions that reflect a model's internal reasoning.

Figure 1: Token-level vs. Structure-aware Fusion. Given pivot and source logits of shape [L,3] (sequence length L, vocab size 3), token-level methods (left) align dimensions independently, ignoring token interactions. GLD (right) aggregates outer products into [3,3] co-activation graphs, capturing semantic dependencies via structure-aware graph alignment.

To address this challenge, we model semantic dependencies in the logit space using graph structures. Rather than treating token logits independently, we represent model outputs as feature-level graphs, where nodes correspond to token dimensions and edges capture their co-activation patterns across sequence positions. This representation reflects how a model internally organizes semantic relationships, beyond surface-level token probabilities, to support coherent reasoning and inference. Constructing such graphs from high-dimensional logits is nontrivial: vocabulary sizes are large, and most logits are near-zero, making full pairwise modeling inefficient and noisy. To mitigate this, we apply top-$k$ sparsification to retain salient dimensions, followed by inner-product similarity to define edge weights. This yields interpretable graphs that faithfully reflect model reasoning behaviors.

Building on this, we propose InfiGFusion, a semantic structure-aware fusion framework that distills knowledge from multiple source models by aligning their logit graphs. At its core is the Graph-on-Logits Distillation (GLD) loss, which captures semantic dependencies by aligning structure-level information between source and pivot models. Specifically, we employ the Gromov-Wasserstein (GW) distance to match graph structures. However, standard GW incurs $\mathcal{O}(n^4)$ complexity, making it impractical for large vocabularies. To address this, we design a closed-form approximation based on node aggregation and sorting-based matching, reducing the cost to $\mathcal{O}(n \log n)$. We further provide theoretical guarantees on the approximation error, ensuring alignment fidelity.

We validate InfiGFusion across diverse fusion scenarios, demonstrating consistent improvements over state-of-the-art methods. By explicitly modeling semantic dependencies through logit graph alignment, InfiGFusion excels in tasks involving multi-step reasoning, relational inference, and fine-grained alignment—scenarios where token-level objectives fall short. Extensive experiments on 11 benchmarks, including reasoning, mathematics, and coding, confirm its effectiveness in producing more capable fusion models.

Overall, our contributions can be summarized as follows:

• InfiGFusion is the first structure-aware fusion framework that models semantic dependencies via feature-level logit graphs and aligns them via a novel Graph-on-Logits Distillation loss.
• We propose a novel approximation to Gromov-Wasserstein distance, reducing complexity from $\mathcal{O}(n^4)$ to $\mathcal{O}(n \log n)$, with provable error bounds and rigorous theoretical analysis.
• Experiments on 11 benchmarks show that InfiGFusion consistently outperforms baselines, with significant gains on complex reasoning tasks (+35.6 on Multistep Arithmetic, +37.06 on Causal Judgement), demonstrating superior multi-step and relational inference.

Related Work

Collective LLM: To aggregate the capabilities of specialized LLMs, several strategies have been proposed. Prompt-based integration methods, such as Toolformer, ReAct, and HuggingGPT, guide LLMs to invoke tools or auxiliary models through natural language prompts. Multi-agent collaboration frameworks like ChatArena and AgentBench treat LLMs as interacting agents that debate or coordinate to improve reasoning. Model ensembling aggregates outputs via voting or reranking, but suffers from high inference cost. Mixture-of-Experts (MoE) models activate sparse expert subsets per query, offering scalability but introducing routing complexity. Parameter merging methods combine weights from fine-tuned models, but require architectural alignment. While these strategies explore different trade-offs between performance, efficiency, and flexibility, they generally assume the underlying models remain fixed. In contrast, we focus on knowledge fusion, where a pivot model learns from multiple sources during training.

Model Fusion via Knowledge Distillation: Traditional knowledge distillation (KD) transfers knowledge from a large teacher to a smaller student via features, logits, or output distributions, primarily for compression or acceleration. In contrast, model fusion via KD aims to integrate multiple source models—often with varying architectures, scales, and vocabularies—into a single target of similar size. Recent methods fuse multiple LLMs by aligning token distributions using KL divergence or Wasserstein distance. These approaches typically require vocabulary alignment and assume semantic compatibility at the token level. However, they overlook semantic dependencies among logit dimensions—i.e., how token probabilities interact under internal reasoning, which are crucial for preserving logical structure. Our work addresses this gap by introducing a structure-aware distillation objective that explicitly aligns semantic dependencies via graph alignment. Unlike prior methods, it avoids explicit vocabulary alignment by modeling logits as graphs, enabling robust, vocabulary-agnostic fusion across heterogeneous models while preserving reasoning structures.

Gromov-Wasserstein Distance: The Gromov-Wasserstein (GW) distance is a relational optimal transport metric for aligning distributions based on internal structural similarity. It has been widely used in graph comparison and structure-aware representation learning, especially when node correspondences are unknown. However, its standard formulation incurs $\mathcal{O}(n^4)$ complexity due to fourth-order interactions, rendering it impractical for large-scale logit graphs in LLMs. To mitigate this, prior works apply entropic regularization, Sinkhorn-based solvers, or low-rank projections. While effective for small graphs, these methods lack closed-form solutions and suffer from convergence instability, making them less suited for distillation over vocabulary-sized graphs ($n > 10^5$). Our method takes a non-iterative, sorting-based approach: it compresses each similarity matrix into node-level features and aligns them via sorted-based matching, reducing complexity to $\mathcal{O}(n \log n)$. Unlike prior approximations, it yields a differentiable, training-friendly loss with a provable approximation bound. To our knowledge, this is the first closed-form, theoretically guaranteed GW approximation tailored for aligning logit graphs in LLM fusion.

Method

InfiGFusion fuses heterogeneous LLMs by distilling structured knowledge from multiple source models into a unified pivot model. Central to this process is the Graph-on-Logits Distillation (GLD) loss, which models logit vectors as feature-level graphs and aligns their structures across models.

To capture semantic dependencies, we adopt Gromov-Wasserstein distance for structure-level alignment between logit graphs. However, the standard GW incurs $\mathcal{O}(n^4)$ complexity, prohibitive for LLM-scale vocabularies ($n > 10^5$). We propose a sorting-based approximation that reduces this to $\mathcal{O}(n \log n)$ by compressing graph structures into node-level features. Furthermore, we provide a provable error bound of $\frac{n - 1}{n^2} + \frac{m - 1}{m^2}$, ensuring faithful and efficient graph alignment.

Figure 2: InfiGFusion framework. Given instruction-response pairs, source and pivot models produce logits, sparsified into feature-level graphs capturing semantic dependencies. We align graphs via an efficient Gromov-Wasserstein approximation (GLD), reducing complexity from O(n^4) to O(n log n). The overall objective combines structure-aware distillation (GLD) with token-level distillation (ULD) and supervised signals (SFT) for robust fusion.

Overview of InfiGFusion

InfiGFusion integrates knowledge from diverse source models $\{M_1, \dots, M_S\}$ into a unified pivot model $M_0$, enhancing reasoning capabilities without increasing inference cost.

Given an instruction-response dataset $\mathcal{D}$, $M_0$ is trained to align with both human supervision and source model behaviors via two complementary objectives: (i) Universal Logit Distillation (ULD): a coarse token-level alignment combining supervised fine-tuning (SFT) and a sorting-based Wasserstein-1 approximation; (ii) Graph-on-Logits Distillation (GLD): a structure-aware objective aligning relational dependencies in logit space. The overall loss is:

where ULD ensures distributional alignment and GLD captures fine-grained reasoning structures. For token-level ULD, we adopt the linear-time Wasserstein-1 approximation, which compare the logit distributions between pivot model $t$ and source models $s$:

where $\mathbf{f}^{\downarrow}$ denotes sorting. This provides an efficient coarse alignment signal, complementing GLD's structural matching.

Dynamic Sparse Graph Construction

To capture structural dependencies among logit dimensions while avoiding the cost of full $d \times d$ modeling ($d$ is the vocabulary size for each model), we adopt a two-stage strategy: top-$k$ sparsification followed by dynamic graph construction.

Top-$k$ Logit Sparsification. Given a logit tensor $\mathbf{Z} \in \mathbb{R}^{B \times L \times d}$, we extract the top-$k$ logits for each position to obtain $\mathbf{Z}^{\text{top}} \in \mathbb{R}^{B \times L \times d'}$, with $d' \ll d$. We then unify the top-$k$ dimensions across tokens to obtain a consistent $d'$-dimensional representation per sample.

Dynamic Graph Construction. To model the interdependencies among logit dimensions, graph construction can generally follow two paradigms: (i) token-wise similarity across the sequence, or (ii) feature-wise similarity among logit dimensions. The former focuses on token interactions, while the latter captures intrinsic correlations between semantic dimensions. In this work, we adopt the feature-wise approach, as it empirically yields better alignment performance and more effectively preserves structural information. Specifically, for each sample $b$, we construct a graph $G_h = (V, E)$ over the reduced logits $Z_b \in \mathbb{R}^{L \times d'}$, where each node corresponds to a selected logit dimension. The adjacency matrix $C_b \in \mathbb{R}^{d' \times d'}$ is defined via dot-product similarity across the sequence:

This construction approximates cosine similarity after normalization and is applied to both pivot and source models at every training step, resulting in dynamic, data-dependent graphs that evolve with model predictions.

Efficient Approximation Algorithm

Directly applying Gromov-Wasserstein (GW) distance for structure-level alignment incurs $\mathcal{O}(n^4)$ complexity, which is infeasible for LLM-scale graphs. We propose an efficient approximation by summarizing graph structures into node-level features and aligning them via sorting-based matching, reducing complexity to $\mathcal{O}(n \log n)$. This method maintains alignment fidelity with a provable error bound of $\frac{n - 1}{n^2} + \frac{m - 1}{m^2}$.

Given intra-graph similarity matrices $C_s, C_t \in \mathbb{R}^{d' \times d'}$ from source and pivot models, the squared GW distance is defined as:

Despite its expressiveness, the GW objective requires $\mathcal{O}(n^2 m^2)$ complexity due to its fourth-order tensor form, making it impractical for large logits ($n,m$ are the vocabulary sizes). To address this, we derive a closed-form approximation by compressing the graph structure into scalar features.

Graph Feature Compression. Our key insight is that the structure of each similarity matrix $C \in \mathbb{R}^{n \times n}$ and $D \in \mathbb{R}^{m \times m}$ can be summarized using scalar node-level features, where $n, m$ are vocabulary size. Specifically, we compute degree-style features:

which captures the overall importance or connectivity of each logit dimension in the graph. This reduces the original $d' \times d'$ structure to two vectors in $\mathbb{R}^{d'}$.

Separable Relaxation. Substituting these features into the original objective yields a simplified, separable form:

where symmetry is assumed across $\gamma_{ik}$ and $\gamma_{jl}$. This relaxation significantly lowers computational cost. The error is theoretically bounded:

Proposition 1 (Approximation Error Bound). Let $C \in \mathbb{R}^{n \times n}$ and $D \in \mathbb{R}^{m \times m}$ be the similarity matrices derived from logit self-inner products and row-normalized such that $\sum_{j} C_{ij} = 1$ and $\sum_{l} D_{kl} = 1$. Then the absolute error between the exact and approximated GW distances satisfies:

In practice, the size of $n,m$ is more than $10^5$. This bound guarantees the error decays as $\mathcal{O}(1/n)$, enabling scalable use in large-vocabulary settings.

Sorting-Based Matching. To avoid solving for the optimal transport plan $\gamma$ explicitly, we take a direct strategy like ULD that adopts a deterministic transport strategy based on sorting. Let $f_C^{\downarrow}, f_D^{\downarrow}$ be the descending-sorted versions of $f_C$ and $f_D$. We construct a matching that aligns the $i$-th largest entry of $f_C$ to the $i$-th largest entry of $f_D$:

This effectively reduces the original optimization to an $\mathcal{O}(n \log n)$ sorting-based assignment.

Stability Analysis. Beyond efficiency, our approximation improves training stability. We compare the Lipschitz constants of common alignment losses:

Proposition 2 (Lipschitz Constants Comparison). We have the following relationship between the Lipschitz constants ($L$) for the Gromov-Wasserstein, Wasserstein, and KL losses:

for vocabulary size $D$ and logit range $R$, which means the GW-guided training models have tighter generalization bounds.

The logit range $R$ is usually smaller than 120 for 14B LLMs. For example, with vocabulary size $D = 150,000$, $R$ is observed in the range of $[40, 120]$. This suggests GW-based alignment yields the lowest sensitivity to perturbations, contributing to improved generalization and stable training.

Unified GLD Loss

To support multi-source fusion, we extend the GLD loss into a unified formulation. For each source model $M_s$, we compute a structure-aware alignment loss with respect to the pivot model $M_0$ and aggregate all source contributions:

where $C_s$ and $C_0$ denote the similarity matrices of the source and pivot graphs, respectively.

This unified objective enables efficient and scalable distillation from multiple source models, while preserving both semantic and relational structures across models.

Experiments

We evaluate InfiGFusion on large-scale LLM fusion, focusing on integrating models with diverse architectures and reasoning styles while maintaining inference efficiency. Our experiments verify that GLD enables faithful and effective fusion beyond token-level objectives. InfiGFusion outperforms state-of-the-art fusion methods and strong open-source baselines, with significant improvements in multi-step reasoning and relational inference.

Experimental Setup

Datasets. We construct a novel multi-task training dataset comprising 130k examples across general reasoning, mathematics, and coding. (i) The 52K general data are sourced from Infinity-Instruct. (ii) Our mathematical dataset comprises 39k pairs of math questions and corresponding answers. The questions are sourced from the NuminaMath_1.5 dataset, while the answers are distilled from the DeepSeek-R1-671B model. NuminaMath_1.5 represents the second iteration of the widely acclaimed NuminaMath dataset. It offers high-quality data for competition-level mathematical problems across various domains, including Algebra, Geometry, Combinatorics, Calculus, Inequalities, Logic and Puzzles, and Number Theory. (iii) In the code generation domain, we utilized the KodCode-V1-SFT-R1 dataset, which contains approximately 268k samples. Each sample was fed into our pivot model, which was prompted to generate 5 random responses per input. The generated outputs were then passed through a sandbox-based evaluation system. For each sample, if at least one of the five responses failed, the sample was flagged for further consideration. From these flagged samples, we filtered and selected 39k high-quality examples for further distillation.

Models and Baselines. We fuse three source models: Qwen2.5-14B, Qwen2.5-Coder-14B, and Mistral-24B, into a Phi-4 pivot model. Fusion baselines include MiniLogit, InfiFusion, FuseLLM, and FuseChat.

Evaluation. We assess fusion performance on 11 benchmarks spanning reasoning (BBH, ARC-C, MMLU), mathematics (GSM8K, Math, TheoremQA), coding (MBPP, HumanEval), text reasoning (DROP, HellaSwag), and instruction following (IFEval).

Main Results

The table above summarizes the performance of InfiGFusion across eleven benchmarks, compared with the pivot model, source models, and representative model fusion methods.

Comparison with Pivot and Source Models. InfiGFusion consistently outperforms both the pivot model and source models in overall performance. Compared to the pivot model, it achieves a significant improvement (+2.53 avg), particularly on complex reasoning benchmarks such as BBH, ARC, and MMLU, where semantic richness and multi-step inference are crucial. For instance, InfiGFusion surpasses the pivot by +16.7 on BBH, demonstrating its ability to aggregate diverse reasoning styles beyond single-source capabilities. Compared to source models, InfiGFusion effectively integrates complementary strengths. While individual source models excel in isolated tasks (e.g., Mistral-Small on BBH, Qwen2.5-Coder on HEval), they struggle with generalization across domains. InfiGFusion surpasses these models with a more balanced performance, avoiding their trade-offs and consolidating knowledge into a unified pivot.

Comparison with Model Fusion Baselines. Against existing fusion methods (MiniLogit, FuseLLM, FuseChat), InfiGFusion exhibits superior fusion quality. Notably, while other distillation methods provide moderate improvements over the pivot, they often suffer from degraded performance in high-dependency tasks (e.g., logical reasoning, multi-condition alignment). InfiGFusion consistently delivers better results on structurally challenging tasks, indicating its advantage in preserving inter-token dependencies during fusion. Furthermore, compared to simpler fusion objectives, the structure-aware GLD loss of InfiGFusion yields better alignment without sacrificing robustness.

Efficiency and Resource Usage. In terms of computational cost, InfiGFusion maintains competitive efficiency. Despite introducing structure-aware alignment, its GPU hours remain modest (195 GPU hours), substantially lower than multi-step fusion pipelines like FuseChat (650 GPU hours). Model size remains identical to the pivot (14B), ensuring no additional inference overhead.

Complex Reasoning Performance

We evaluate InfiGFusion on 18 reasoning datasets spanning BBH and MMLU, covering multi-step logic, structural dependencies, and fine-grained semantic alignment, as shown in the following table. These tasks are known to challenge token-level fusion methods due to their complex reasoning path.

Improved Reasoning over SFT. Compared to Phi4-14B after supervised fine-tuning (SFT), InfiGFusion delivers consistent gains, boosting the average accuracy from 77.94% to 83.79%. Notable improvements are observed in tasks demanding complex relational reasoning: Multistep Arithmetic (+34.8%) and Causal Judgement (+34.39%). This demonstrates InfiGFusion's ability to enhance multi-hop inference and causal understanding by explicitly modeling semantic dependencies.

Structural Alignment Benefits. InfiGFusion excels in Tracking Shuffled Objects and Logical Deduction tasks, where aligning token relations is critical. For instance, it surpasses Phi4-14B SFT on Logical Deduction Five Objects (+1.6%) and Seven Objects (+0.4%), showcasing robust structure-level fusion beyond token alignment.

Preserving Pivot Strength. InfiGFusion maintains competitive performance in knowledge-centric tasks like International Law and Ruin Names, matching or slightly improving upon SFT baselines. This indicates that InfiGFusion not only integrates external model knowledge but also preserves the pivot model's original strengths. These results confirm that InfiGFusion's structure-aware fusion consistently enhances reasoning capabilities compared to token-level methods.

Ablation Study

Effect of Source Model Diversity. The following table evaluates the effect of fusing different source models. Individually, Qwen-Coder (Qc), Qwen-Instruct (Qi), and Mistral (M) each provide moderate gains, reflecting their domain-specific strengths. Notably, combinations such as Qi-M further amplify reasoning (+3.57), math (+3.06), and coding (+8.22) performance, demonstrating that complementary inductive biases from different models enhance the pivot model. InfiGFusion achieves the best results (+2.53 avg), confirming that our framework effectively consolidates diverse reasoning behaviors into a unified model.

Effect of Loss Components. The following table analyzes our loss design. Removing ULD (w/o ULD) leads to a significant drop (-1.52 avg), particularly in reasoning (-1.8), indicating that coarse-grained token-level alignment is crucial for providing a reliable starting point. ULD narrows large distribution mismatches, effectively preparing the model for the finer-grained structural alignment performed by GLD. Excluding GLD (w/o GLD), on the other hand, yields a smaller overall decline (-0.69 avg), but reasoning performance still drops notably (-1.36). This highlights that while ULD handles first-order token matching, it lacks the capacity to capture higher-order semantic dependencies. When combined, ULD and GLD form a synergistic alignment mechanism: ULD aligns distributions coarsely and stabilizes optimization, while GLD refines alignment by enforcing consistency in the relational structure of logits. The empirical results validate this design, with the full loss achieving the best performance. All ablation studies are tested in the Top-$k=10$ setting.

Effect of Top-$k$

InfiGFusion sparsifies logits by retaining top-$k$ token dimensions before graph construction, selecting the most salient indices per sequence position. This inductive bias suppresses noisy activations and emphasizes meaningful token dependencies, serving as the foundation for graph-based semantic alignment. We evaluate Top-$k \in \{5, 10, 15, 20, 25, 30\}$ and report the results below. Increasing $k$ from 5 to 10 brings notable gains across all tasks (Avg +0.54), as larger $k$ captures richer token interactions essential for semantic graph construction. However, further increasing $k$ beyond 10 yields diminishing returns and even slight degradations, particularly in reasoning (-0.49 from Top10 to Top30). This suggests that excessive low-confidence tokens introduce spurious edges, diluting graph discriminability. The observed trend aligns with the heavy-tailed nature of LLM logits, where only a small subset of token dimensions are informative. Overall, Top-10 achieves the best trade-off, balancing informative graph structure and computational efficiency.

Case Study

To illustrate InfiGFusion's strengths beyond token-level alignment, we analyze representative cases, comparing outputs from a standard SFT-tuned model (Phi4) and our InfiGFusion model. The selected examples highlight InfiGFusion's superior reasoning in multi-step causality and relational inference.

Figure 3: Case study.

Case 1: Frank T. Shooting Incident. While both models predict "No," InfiGFusion performs a deeper step-by-step causality analysis. It explicitly identifies the causal chain's disruption—distinguishing between "intent," "misfire," and "accidental result"—showcasing robust multi-step causality disambiguation capabilities.

Case 2: Wallace's Dual Cause of Death. Unlike Phi4's surface-level judgment, InfiGFusion reasons through the latent causal structure, identifying "organized crime" as a distal cause. This reflects its strength in relational and indirect causality inference, where it effectively models complex event chains and upstream dependencies.

These cases demonstrate that InfiGFusion excels at: (i) Multi-step inference: Decomposing complex causal chains into interpretable reasoning steps. (ii) Relational causality reasoning: Capturing indirect and upstream causal factors often missed by SFT models. (iii) Fine-grained disambiguation: Distinguishing intent, action, and outcome with structured alignment of reasoning behaviors.

Analysis of the GW Approximation

In addition to the main results, we also conducted some experiments to analyze our proposed Gromov-Wasserstein approximation. It compares the quality of GW estimation in a simple simulated scenario. For example, by obtaining samples from two Gaussian distributions, and comparing the value of exact GW, sinkhorn-based approximation, and our proposed approximation. While our experiments demonstrate that the approximation is suitable as a training objective, it would be interesting to also understand its behavior as GW estimator.

Experiment setup: For each $n\in\{50,100,200,500,1000,1500,2000,2500,3000\}$, we sample $n$ points from $\mathcal{N}(0,I)$ and $\mathcal{N}(1,I)$, build their $n\times n$ Euclidean distance matrices, and compute: (i) Exact GW and Sinkhorn GW via ot.gromov.gromov_wasserstein2(..., loss='square_loss'); (ii) Approx GW via our sorting-based closed-form. Each configuration is repeated $3\times$ to report mean$\pm$std. Approx RE quantifies the approximation accuracy relative to the exact GW: $\mathrm{RE}=\frac{|\widehat{GW}_{\text{approx}}-GW_{\text{exact}}|}{GW_{\text{exact}}}$ (mean$\pm$std over 5 seeds). In this symmetric Gaussian setting, Exact and Sinkhorn GW coincide to our reported precision (even with $\varepsilon=10^{-1}$), evidencing the high accuracy of the Sinkhorn estimator here.

Results and analysis: Each setting is repeated 5 times with different seeds; we report mean±std. Across all sizes, Sinkhorn GW matches Exact GW to our reported precision (RE ≈ 0), hence we merge them into a single column in the table above. Our Approx GW exhibits decreasing relative error as $n$ grows (from ≈ 0.59 at $n=50$ down to ≈ 0.39 at $n=3000$), consistent with the $O(1/n)$ behavior. In terms of runtime, Exact/Sinkhorn scale poorly and rapidly become impractical beyond a few hundred points (already seconds at $n=1000$ and tens of seconds by $n\ge1500$ on our machine), whereas Approx GW remains below ~0.3 s even at $n=3000$, demonstrating its $O(n\log n)$ scalability.

Conclusion

In this work, we propose InfiGFusion, a structure-aware fusion framework that explicitly models semantic dependencies among logits through a novel Graph-on-Logits Distillation (GLD) loss. By leveraging logit graph representations and an efficient Gromov-Wasserstein approximation, InfiGFusion enables the integration of heterogeneous LLMs without increasing inference cost. Extensive experiments on 11 benchmarks demonstrate that InfiGFusion consistently outperforms state-of-the-art fusion methods and baselines. In particular, it achieves notable gains in multi-step and relational reasoning tasks, such as +35.6 on Multistep Arithmetic and +37.06 on Causal Judgement over SFT models. It highlights the importance of structure-preserving alignment for effective model fusion in complex reasoning tasks, which is a promising direction for advancing collective LLM.

Limitations

While InfiGFusion demonstrates clear advantages in tasks involving multi-step reasoning, causal inference, and relational alignment, its performance shows limitations in scenarios dominated by literal token matching or factual recall. For tasks where the correct answer is encoded by a few dominant logit dimensions (e.g., factoid QA), the added structural alignment of GLD brings marginal benefits and may even introduce noise from irrelevant dependencies. Moreover, when source models present conflicting factual knowledge, InfiGFusion focuses on preserving relational consistency but lacks explicit mechanisms for factual conflict resolution, sometimes amplifying semantic ambiguities. These limitations reflect inherent trade-offs in structure-aware fusion and motivate future work on adaptive dependency modeling and source reliability estimation.

Our model is available at https://huggingface.co/InfiX-ai/InfiGFusion-14B.

BibTeX

@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}
}