Probabilistic Token Alignment for Large Language Model Fusion
Authors: Runjia Zeng
James Chenhao Liang
Cheng Han
Zhiwen Cao
Jiahao Liu
Xiaojun Quan
Yingjie Victor Chen
Lifu Huang
Tong Geng
Qifan Wang
Dongfang Liu
Venue: NeurIPS 2025 poster ()
Keywords: knowledge fusion
model fusion
LLMs
token alignment
optimal transport
Primary area: deep_learning

Abstract: 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.

Probabilistic Token Alignment for Large Language Model Fusion
Runjia Zeng', James Chenhao Liang', Cheng Han', Zhiwen Cao', Jiahao Liu', Xiaojun Quan', Yingjie Victor Chen', Lifu Huang", Tong Geng" Qifan Wang",
Dongfang Liu" '
Rochester Institute of Technology U.S. Naval Research Laboratory
University of Missouri-Kansas City Adobe
MMettuan
Sun Yat-sen University
"Purdue University
UCC Davis
University Rochester Rice University
MMMeta AI CCorresponding author
Abstract
Training large language models (LLMs) from scratch can yield models with unique functionalities and strengths, but it is costly and often leads to redundant capabili- ties. A more cost-effective alternative to fuse existing pre-trained LLMs with different architectures into more powerful model. However, key challenge in existing model fusion their dependence on manually predefined vocabulary alignment, which may not generalize well across diverse contexts, leading to per- formance 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 classic mathematical problem: optimal transport, seamlessly leveraging distribution-awaars learning to f'ccilitate more coherent model fusion. Apart from its inherent generality, PTA-LLM exhibits interpretability from 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 /runjia.tcch/neurrsss __ttsslptmmmm
Introduction
Bay-ewc Mart
W1
Stat
using
Merged W Wiight
Fused matric
System Oulput
Systtmm Oulput
System Outgut
Liaer
501
(a) Model Emsemble
(b) Weight Merging
(c) Knowledge Fusion
Knowledge Fusion Perfomancc
Figure 1: PTA-LLM (ours) us. concurrent arts (i.e., model ensemble [ and weight merging (2] under model fusion paradigm (i.e., knowledge from the "cat" and "rabbit" source models is fused into
the "llama" target model). Our knowledge fusion method yields general performance gains across multiple capabilities in (d), where all scores (see Table 2) are normalized for better visualization.
The rise of large language models (LLMs) such as Llama-2 [] Ope [4), and MPT (5). driven by scaling laws [6], has yielded significant advancement across a broad range tasks (see Fig. (d). The narrow dashed line area indicates its respective fields of advantage). Nevertheless, the reliance on scaling laws brings substantial computational demands 77) posing a noticeable
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

impediment to the development of more robust baselines. A question thus naturally emerges: ① How can we construct stronger baselines without resorting to the naive application of scaling laws?
Pioneering research has begun to address this question through the concept of model fusion 8 2 2
101, focusing on model ensemble (see Fig. [
(a) and weight merging (see Fig. (b)). Recently,
prominent technique called knowledge fusion IT aggregates the probabilistic distributions generated by individual LLMs and transfers this fused representation to a target model via distillation (see Fig. D
(c)), enabling it to be more inference-efficien
After further employing token alignment [1] the
misalignment issues arising from the use of different tokenizers across models are mitigated, allowing knowledge fusion to be sebiitecture-annoitt The exemplary ad advantages offered by knowledge fusion reframe question ① into @: How can we further optimize LLMs through knowledge fusion?
Although knowledge fusion shows a promise avenue, two significant token alignment challenges still remain unresolved. ① The manually designed mapping strategy is overly simplistic, failing to capture the intricate patterns within the data. Tokens appearing in varying contexts often align with different objectives, and the bias introduced by this "rigid" alignment reduces the model's capacity to fully learn from the data, ultimately diminishing performance. @ The alignment of top-k predicted token sets from the source and target LLMs is performed independently, without taking into account their associated probabilities or overall distribution. This isolated strategy may achieve local optimality at each step alignment, but not whole coherent fused metric. Thus, the core question @ becomes more specific: @ How can we effectively fuse LLMs with an adaptive and coherent matrix?
To this end, we introduce Probabilistic Token Alignment for Large Language Model Fusion (PTA- LLMM). During the matrix fusion, we first employ dynamic algorithm to determine an optimal token pairing between the generated sequence from the source and target model. After obtaining the token pairings, a logit-level alignment will be conducted to resolve the token ID misalignment. Specifically, for the top-k predicted token sets from both source and target models, we hypothesize and further prove (see empirical results in Table 2 and 3 that the probabilistic distributions generated by distinct LLMs are coherent and reflective of their respective inherent knowledge. Therefore, PTA-LLM leverages the global generative distributions of each model's logits during token alignment, externalizing their collective knowledge and facilitating more precise mapping. To achieve this, our approach is grounded in Optimal Transport (OT), which optimally transforms one probability distribution into another while minimizing a predefined cost. By harnessing OT, we align or "transport" logit distributions between models, offering an effective solution. In contrast to hard mapping strategies, which align each token independently of its context, our proposed PTA-LLMM employs a soft probabilistic alignment (detailed in $3.2 This approach better captures the intricacies of various linguistic context and thus establishes a stronger performance baseline, addressing the challenge O. Additionally, by incorporating distribution- model representations (through the visualization results in $4.4). leading to marked improvements in generalization across wide range of tasks (see Table $2). answering challenge @
learning, this method facilitates more consistent
PTA-LLM enjoys a few attractive qualities. I. Generality. The global probabilistic distribution transport enhances the coherence of the representations, thereby improving the model's 's ability to generalize across a wide range of tasks and supporting the transfer of underlying representations for effective evaluation (see Table 2). II. Stability. The reframing through an optimal transport perspective introduces a soft probabilistic alignment, offering a flexible and adaptive solution to diverse contexts and performing stablly even in difficult tasks (see Table 3 III. Interpretability. The effectiveness of our approach is supported by theoretical insights from distribution learning and further validated through visualization results. It investigates the underlying mechanisms of token alignment, a critical operation in knowledge fusion that has been largely overlooked in prior research. This distinguishes PTA-LLM from most existing knowledge fusion models, which fail to elucidate precisely how token alignment works (see $4.4).
2 Related Work
Model Fusion has garnered significant attention as a means to enhance the general performance of LLMs. The fusion techniques can be categorized into three primary categories: model ensembling, weight merging, and knowledge
Model ensembling combines the predictions of independently trained models to improve overall performance. Common approaches include weighted averaging [13], majority voting [4), and pairwise ranking [15]. However, a common challenge of model ensembling is that it requires
2

---

maintaining multiple models during inference, leading to high memory consumption and latency. Weight merging combines the parameters of multiple models to synthesize a new, unified model. This method is especially effective when the models share identical architectures 2 D Weight merging is enhanced by linear mathematical operations on adapter parameters for improved model performance
and generalization [16] [7 However, weight merging overly relies on architectural uniformity across models and requires manual tuning, limiting its applicability to diverse architectures (i.e., low generalizability). In contrast, knowledge fusion offers a flexible and efficient model integration, particularly when the underlying architectures differ (i.e., a common case in LLMs). usually distills knowledge from multiple teacher models into a single student model. Based on its advantage, follow this technique in our study. More discussions are presented in SST
Token Alignment was first introduced as a solution to address the misalignment problem between tokenizers with different size of vocabulary, specifically when aligning their respective distributions. The concept was initially formalized by [12), who employs a search algorithm to minimize the alignment cost between token sequences. Following research [11 19 further explored flexible mappings for prominent performance via MinED strategy, statistical mapping, etc.
However, existing methods remain limited by their reliance on surface-level token correspondences (i.e., based solely on the strings it comprises), which leverage minimum edit distance to align the logit. Our approach. besides using edit distance as one metric, advances token alignment by incorporating the corresponding logit values into the individual cost within the transport framework. Optimization right now performed at both the "surface-level" and "logit-level" (see Eq.
3 PTA-LLM
In this section, present PTA-LLM, novel probabilistic token alignment method for achieving
general and coherent fusion of large language models (LLMs), as illustrated in Fig 2
Specifically,
outline our comprehensive knowledge fusion framework and tuning strategy in 3 Following
this, we elaborate on the design of our probabilistic token alignment approach in $3.2 where the
probabilistic distribution matrices from source LLMs are aligned into fused representation via dynamic pipeline, which involves two primary stages: dynamic token pairing and probabilistic token alignment. Last but not least, in $22 we provide detailed description of the implementation and the algorithm utilized in our approach. More implementation details will be provided in $55
Table I: The key notations used in PTA-LLM. See all notations in $$4
Notation
Representation
A training corpus consisting of a collection of text sequences!
P.
The probabilistic distribution matrix for the source model, consisting of L tokens (denoted as A C R"
The prohabilistic distribution matrix for the target model, consisting N' tokens (denoted as E R
P
The fused probabilistic distribution matrices for model fusion, consisting of N tokens (denoted as E R)))
Target model's predictions during the fine tuning The final x x m optimal transport plan matrix of non-negative entries Gam
T
3.1 Problem Statement & Overall Objective
Let f represent an input text sequence sampled from a corpus C. A probabilistic distribution matrix P E RKxV is obtained by evaluating the output token distirbutions of large language model (LLM) from the input i, where K corresponds to the output sequence length, and V denotes the size of the vocabulary. The i-th row of this matrix represents the predicted probability distribution over the vocabulary for the i-th token in the sequence. In the context of combining two L.MSS (source and target), we consider the probabilistic distribution matrices P. E LxVs for the source model and P. E IR NXV, for the target model, where L and N denote the sequence lengths, and V. and Ve represent the vocabulary sizes of the source and target models, respectively. When these models employ different tokenization schemes, misalignment between the tokens of the source and target models arises, thereby complicating the integration of their probabilistic outputs. Addressing this issue is essential for effectively combining the outputs of both models. The traditional approach seeks to ensure consistency between the target model's predictions, denoted as Qt, and the fused representation P,, which encapsulates the knowledge from the source model. The knowledge fusion loss is formulated as Lrusion = -E,-c [I, P,)], where ()::.)) is a discrepancy function (such as cross-entropy KL divergence) measuring the difference between the predicted and fused probability
3

---

(a) Knowledge Fusion
Corpas
[1
A)
Ons-lii Label Manix
Feae Mazio
Sturce Madel Tokenizer:
Probubelltic Disrrbation Matdi' Rg
Optiniration
Probahtliets Nittriattoon Matris Q
Fased: Py c g
Drnst Meddi Takeninen V.
Probabdietic tietrbation Matris P, x
Oepadd
Zarzet Model Tokenberr 4/
(b) Prohabilistic Token Aligament
Dynamic Token Pairing
Protabilisti Token Aligneent
AJ Gu
BR
B
P.=lB, B Bxl LINN
each pairing {A Bx)
Ar
Be 79
Ps
30
Transpoi Plan T
221
183
4.60
111
.0.1
Figure 2: Probabilistic token alignment under the knowledge fusion paradigm. (a) The overall knowledge fusion pipeline (see $3.1), and (b) two-stage probabilistic token alignment (see $3.2).
including dynamic token pairing and probabilistic alignment using optimal transport reformulation.
distributions. The fused output Py is a probabilistic distribution matrix that represents the combined strengths of both the source and target models, formally defined as P, Matrix MaarrixAligmment((2 Pt).
In this work, we propose PTA-LLM, framework designed resolve discrepancies between the tokenization schemes of the source and target models. The principal objective is to minimize the divergence between the target model's probabilistic predictions Q₁ and the corresponding one-hot encoded label matrix O, 0, XxV where each row of O, indicates the correct token as one- hot vector. Specifically, we define causal language modeling (CLM) loss, which measures this divergence, LCLM = - (Q., O,)]. between the predicted probabilities and the true labels. Consequently, the overall training objective of our proposed method is optimize weighted combination of the CLM loss and the fusion loss, formalized L = ACCLM + (1 Fusionn where [0, 1] a hyperparameter controlling the trade-off between the causal language modeling loss and the fusion objective. This ensures that the target model can effectively learn from both its own predictions and the knowledge transferred from the source model.
3.2 Probabilistic Token Alignment
Dynamic Token Pairing The task of aligning two distinct probabilistic distribution matrices, Pa and Pa poses significant computational challenge due the inherent differences in both sequence length and vocabulary size. The core problem involves finding suitable alignment between tokens from the source model's distribution P, and those from the target model's distribution P.. More precisely, for each token Aj (j E [1, L)) from Ps, we aim to pair it with corresponding token B₃ (k [1, N)) from P. as shown in (bl).
Given that there are x N potential pairings between these tokens, employing brute-force attempts to explore all possible combinations would be computationally prohibitive, especially as the sequence lengths and vocabulary sizes grow. To address this, we introduce dynamic token pairing, which provides an efficient way systematically explore the space of possible pairings and compute optimal alignment. This approach allows for the minimization of computational complexity while ensuring the best mapping between the source and target tokens.
Formally, given two sequences of tokens [AI.,, Bi:N), our objective to find an alignment that minimizes the overall cost associated with transforming one sequence into the other. Thus, we define the recursion function as:
f(k.j) min (f(k 1,j) c( Aj),
f(k.j - 1) ((B,. Ay), f(k - 1,j 1) + c(Br. A)), [ []
where [kk represents the total cost of aligning the subsequences Bis and Ar.j. while c(Br A,) denotes the predefined cost or distance metric between tokens. In contrast to traditional alignment

---

methods 220 [2], which typically enforce a one-to-one correspondence between elements in the two sequences, our approach introduces generality by relaxing this constraint. Specifically, our formulation allows the dynamic possibilityy assignment that one token in the source domain may align with multiple tokens in the target domain, and vice versa, depending on the characteristics of the tokenization schemes and the specific demands of the alignment task.
By adopting this dynamic token paring strategy, our method is able to handle discrepancies between the tokenization schemes of the source and target models, ensuring that the probabilistic distributions Ps and Pa can be meaningfully aligned, even cases where their underlying token structures differ significantly. This enhanced flexibility is particularly useful in scenarios where the vocabulary sizes and token sequences vary substantially. Utimately, we provide a more robust solution to the alignment problem in the context of knowledge fusion between models.
Probabilistic Token Alignment After determining the optimal token pairings, the next fundamental step involves accurately performing logit-level aligoment to address the token ID (i.e., we transform the text into corresponding token IDs using tokenizer.) misalignment that arises due to the use of different tokenization schemes. As shown in Fig. 2 (b2). even when Aj and By are decoded into the same text by their respective tokenizers, their token ID may differ (i.e., 51 us. 79). Specifically, for each token pair Aj E IR V. and Bg E V., the objective of token alignment is to match the logits from the source token with the logits from the target token in order to achieve consistent representations between the models. The resulting fused token distribution, denoted as B₁. can be defined as:
B. = Token Aligmmmt (Aj: Bz)
(2)
where Token Alignment is function that fuses the logits from the source and target models for each token pairing. This fusion process aims to produce unified token distribution by combining the outputs from both the source and target language models. In addition, Equation highlights that the token fusion for each pairing can be reformulated from the perspective of distribution learning, where the goal is to minimize discrepancies between the two token distributions. Formally, we have:
F
arg min L (Aj, Bz)
where the loss function represents the information loss incurred during the alignment process. The goal is to minimize this loss, ensuring that the information from the source logits is effectively transferred to the target logits without significant degradation.
This optimization problem is conceptually analogous to the classical problem of optimal transport. Our objective is to find a "transport plan" T that minimizes the total cost of transferring probability mass from one distribution, nt to another distribution, v. Henee, in the context of token alignment, we can reinterpret the task as an Optimal Transport (OT) problem, where the aim is to determine a global transport plan that transfers the logits of the source tokens A, to the logits of the target tokens B₁ at minimal cost. This process, under our setting, is formulated as:
m Gay = A,
t- = arg min T20
ΣΣ y=1
Czy Guy
= BBuuvly n = m = 10,
(4)
where i is an n X m matrix of non-negative entries Gzy. representing the amount of logit probability transported from the æ-th indice in the source token Aj the y-th indice in the target token BR- The cost matrix C' captures the alignment cost between source token Aj and target token Be. where we define Cryy as the minimum edit distance (after decoding the indice into text) between the æ-th indice in the source token Aj and the y-th indice in the target token B, (i.e., the L in Equation 33) The constraints Ey= Gay = AA [ar] and Σ Gay [x[y] ensure "logit probability" conservation between the source and target token distributions. See more details in $55
Once the "transport plan" ÿ is determined, the next step is to align the logits by selecting the target token logits with the highest probability for each source token logit, which can be reformulated as:
B. = -((1r Guy) r Rs
(5) Here each pair consists of the index r and the corresponding transport probability Gay from the optimal transport plan Í. The set Rg represents the indices corresponding to the largest values in
5

---

the r-th row of T, which indicate the most probable target token logits for alignment with the z-th source token logit. As shown in Fig. 2 (b2), the final fused logits for each indice in B& are determined by the maximum transported logits. Notably, if multiple source token indices contribute the highest transported logits to the same target indice, their contributions are accumulated (i.e., Gos + Ges). We demonstrate that our probabilistic token alignment can generate an more adaptive (see empirical results in Table 2 and coherent (see the visualization of token in $4.4) fused matrix.
3.3 Implementation Detail
In this section, we present the implementation details of optimal transport and the fusion strategy for fusing different LLMs in our PTA-LLM method.
Optimal Transport As stated in Equation [nd the token alignment tasks are transformed into OT problem. Consequently, how to efficiently compute the global transport plan becomes crucial. To address this, we employ the Sinkhorn algorithm [22] to solve the optimal transport problem following
common practice 23
The implementation of Sinkhorn algorithm is shown in Algorithm
Algorithm 1 Sinkhorn Algorithm for Optimal Transport
Require: Cost matrix C, source token distribution Aj. target token distribution B,, temperature exp (-XC)
Initialize T
2: repeat
3: scale the rows of? T such that the row sums match Aj
4: scale the columns of TT such that the column sums match By:
5: until convergence
6: return T.
Fusion Strategy To effectively merge the collective knowledge of source LLMs while retaining their individual strengths, it's crucial to assess the quality of each LLM and assign different levels of importance to their respective distribution matrices. To do this, when processing text f, we employ cross-entropy loss between the distribution matrices and the gold labels as a measure of the LLMs' prediction quality [24]. lower cross-entropy score for a source LLM indicates a more accurate understanding of the text, and its prediction should thus be given greater weight. Following this principle, we select the distribution matrix with the lowest cross-entropy score as the source LLM distribution matrix. More fusion strategy ablative studies results are shown in Table 4b
4
Experiments
4.1
Experimental Setup
Training details We fine-tune the Ilama-2 7B model using batch size of 256 and maximum sequence length of 2,048 tokens with a combination weight (i.e., the À in $3.1) 0.8 on MiniPile (25) following [II. More details are presented in $S2
Evaluation We evaluate PTA-LLM on six benchmarks (see details in S33 that span various core capabilities of LLMs, including reasoning, coding, commonsense, safty and multilingual ability. Baselines We evaluate the performance of PTA-LLM with three sets of baselines: (1) Source LLMs, including Llama-2 7B B1, OpenLLaMA 7B (4,, and MPT 7B [5: (2) Llama-2 CLM. Llama-2 7B model that further fine tuned on MiniPile using the traditional causal language modeling objective; and (3) FUSELLM [11], Llama-2 7B model trained MiniPile with an emphasis on integrating the capabilities of multiple source models under the knowledge fusion paradigm.
Reproducibility PTA-LLM is implemented in Pytorch (26 using the Huggingface Transformers library 227 accelcrated by FlashAttention (28). Our full implementation will be publicly released.
4.2 Main Results
Table presents the overall performance of PTA-LLM compared to three sets of baseline models (i.e., source LLMs, Llama-2 CLM and FUSELLM). The results indicate that the original LLMS exhibit varying performance across the six benchmarks, with Llama-2 generally achieving the best results, while MPT demonstrates the weakest overall performance. Following continual training on MiniPile, Llama-2 CLM shows modest average improvement of 1.20% over the original Llama-2 model.

---

Table 2: Overall results of PTA-LLM and baselines in six various benchmarks, including 78 tasks in
total. The percentages indicate the rate of improvement/decrease compared to FUSELLM. We further report "Number of Tasks" in [.]. Notably, higher average values indicate better performance in each benchmark. Per-task results and more experiment details are available in Appendix $8
Benchmark of Tasks]
OpenLLaMA
MPT
Llama-2
Llama-2 CLM
FUSELLM
PTA-LLM
Grade School Math
[1]
7.81
9.17
14.18
14.33
14.56
14.71 (+1.03%)
Big-Bench Hard
[27]
33.87
33.38
39.70
40.44
41.01
41.08 (+0.17%)
MultiPL-E
[10]
18.11
17.26
14.63
14.83
15.56
15.88 (+2.06%)
MMLU
[17]
42.11
27.84
46.94
47.65
48.77
49.38 (+1.25%)
ToxiGen
[14]
18.94
18,42
18.56
18.33
18.19
18.89 (+3.85%)
TyDi QA
[9]
27.32
22.11
31.42
31.80
32.99
34.07 (+3.27%)
Avg. Benchmarks
[78]
24.69
21.36
27.57
27.90
28.51
29.00 (+1.72%)
Compared to FUSELLM, PTA-LLM demonstrates an average relative performance gain of 1.72% across 78 tasks. Notably, in the challenging benchmark of ME, which consists of multiple popular programming languages, our approach achieves a significant performance gain of + -2.06% compared with Llama-2. Notable improvements are also observed in core areas such safety and multil- ing. While a slight performance degradation observed the continual training for the ToxiGen benchmark under FUSELLM, PTA-LLM achieves 3.85% relative improvement, highlighting the generality of probabilistic token alignment across diverse contexts. We also find that PTA-LLM experiences minor performance improvement (i.e., +0.17% the BBH benchmark compared to FUSELLM. This decline can be attributed to poor performance of source models. Two of the three source models (i.e., OpenLLaMA and MPT) underperform these tasks, and thus their more coherent token alignment may inadvertently hinder continual training effectiveness in a reasonable jitter. In conclusion, PTA-LLM improves the model's ability generalize across wide range of tasks and supports the transfer of underlying representations for effective evaluation.
4.3 Study of Stability
Table 3: Case study of PTA-LLM in the performance degradation tasks for continue training and FUSELLM. The percentages indicate the rate of improvement/decress compared to Llama-2. We
also denotes its corresponding benchmark in [-]. Case studies for BBH are provided in $S9
Task [Benchmark]
Llama-2
Llama-2 CLM
FUSELLM
PTA-LLM
Causal Judgement [BBH]
50.80
46.52 (-8.43%)
46.52 (-8.43%)
50.80 (+0.00%)
Geometrie Shapes [BBH]
34.40
19,20 (-44.17%)
22.80 (-33.72%)
26.80 (-22, 09%)
Tracking Shuffied Objects (7 objects) [BBH]
11.20
9.60 (-14.29%)
10.40 (-7.14%)
14,00 (+25.00%)
Chemistry [MMLU]
35.97
34.11 (-5.17%)
34.98 (-2.75%)
36.96 (+2.75%)
Jewish [TTxxiee]]
27.00
21.60 (-20.0096)
23.80 (-11.85%)
25.20 (-6.67%)
Arabic (TTYDi QA]
8.47
5.45 (-35.66%)
5.65 (-33.29%)
7.49 (-1.57%)
Swahili [iyDi QA]
43.69
38.97 (-10,80%))
39.78 (-8.95%)
41.68 (4460%)
Avg. Tasks
30.23
25.06 (-17.07%)
26.28 (-13.04%))
28.99 (.0%%)
We observe that certain tasks (6 out of 43 tasks), FUSELLM under the knowledge fusion paradigm exhibits performance degradation, which significantly diminishes its overall efficacy. This suggests instability when exposed to perturbations, such as more challenging or unseen tasks. Consequently, a thorough analysis of these tasks is necessary to provide valuable insights for future research.
Our hypothesis is that the hard mapping token alignment strategy employed by FUSELLM suboptimal in these contexts, necessitating manual specification of alignment strategies tailored to each task for improved outcomes. In contrast, our method reframes the problem through the perspective of optimal transport, introducing soft probabilistic alignment that offers greater flexibility and adaptability across diverse tasks. This approach not only mitigates performance degradation (i.e., achieve overall 8.97% performance mitigation over FUSELLM) but also results in significant improvements, particularly in benchmarks such as BBH (i.e., 14.00 es. 11.20) and MMLU (i.e., 36.96 us. 35.97). For instance, method achieves 25.00% improvement over Llama-2 in the tracking shuffled objects task. These promising results underscore the stability probabilistic token alignment in enhancing model performance across varied contexts.

---

FMselLM
(a) Token A lignment
PCA Variance Ralio: 95.60%
Compactarss Similarity.
the private Carret
Llama
263.34
FiseLLM
237.83
PIA-EIM
Liuma
L
PIA-LLM
equ
driver
sector
car
transter
car
sector
100
300
company
lour
.oompay
-10
bux
hospital
400
.iet
island
1po
bospital
Sthool
room
boot
ggaide
.cchool
room
room
-1000
-500
300 Fastunal
-1000
-500
300 Fastuna
-1000
-500
500 Featural
(b) Target Model: Lalama
(c) FuseLLLM
(d) PTA-LLM (ours)
Figure 3: Study of Interpretability. (a) The abstract understanding of token alignment in FUSELLM and PTA-LLM and their respective evaluation metrics. (b) 2D visualization results of target tokens and fused tokens, where their locations represent semantic information and the sizes indicate their corresponding logit magnitudes. The on the coordinates denotes the logit-weighted center of each token. Additional visualization results presented in $S6
4.4 Study of Interpretability
Although the emergence of knowledge fusion as model fusion paradigm has gained huge attention, the underlying rationale remains unclear. In this section, we tend to provide distribution insights into token alignment's mechanisms and offer guidance for its optimal utilization. As shown in Fig. 3 we delve into a specific context have an in-depth analysis of token alignment. Given we have previously aligned tokens like "the private," we need to align the token pair from the source model and target model to form the next fused token. For tokens from the target model (i.e., Llama-2), we visualize their 10p-10 logits and corresponding indices in 2D space (see Fig. 3 (a), left coordinate, showing only 3 logits
high-level representation). This done by first using the target model's 's tokenizer extract token features, followed by dimensionality reduction using Isomap (29
and PCA (30] (the variance ratio is reported as 95.60% on average in the table in Fig. 3 (a)). Their relative position can reflect the underlying meaning of this indice, and the relative size indicates the magnitude of their corresponding logit. For FUSELLM, traditional hard mapping does not consider their logit and maps each indice to another with pre-defined strategy, acting like "moving" (i.e., change the location without modifying the size) in high-level understanding. In contrast, our method leverages the complete distribution, "transporting" (i.e., distribute the size into current location) the optimal logit into existing indices. Quantitatively, we further compute the average compactness of each token (i.e., the logit-weighted Euclidean distance from each point its center) and the similarity of each token center to the target one (i.e., the Euclidean distance from each center point to the target one) 100 random samples, as shown in the table Fig. 3 (a). empirically demonstrates that our method generates a more coherent fused token, as evidenced by more compact representation (i.e., lower inner distance: 239.44 us. 257.83) and more consistent representation (i.e., lower center distance: 22.25 es. 136.95).
As shown in the down part of Fig. can visually compare the distribution of PTA-LLM fused token with the target token and FUSELLM fused token. Specifically, a more consistent marginal feature distribution between PTA-LLM and target token can be observed from Fig. 3 (b) and Fig (d), where FUSELLM exhibits significantly greater distortion in the overall token representation. The more compact and coherent overall token distribution after employing probabilistic token alignment is aligned with the quantitative results. More implementation details will be claborated in $$6

---

Table 4: set of ablative studies on three different core capablities evaluation benchmarks (i.e., BBH, MMLU, ME). (a) The probabilistic token alignment parameters include two key hyperparameters: convergence threshold and transport window size. (b) The fusion training parameters consist of the combination weight, which controls the relative emphasis during continued training, while the fusion
function determines the source distribution matrix at each training step. See more results in $STO
Choice
BBH
ME
MMLU
Choice
BBH
ME
MMLU
Optimal Transport Convergence Threshold
Combination Weight
1e.4
40.54
48.99
40.39
15.72
lo-5
41.08 (+1.33%)
15.82 (-0,38%)
49.38 (40.80%)
0.8
41,08 (+1.71%)
15,88 (+1.04%)
49.38 (+0.92%)
Token Aligument Window Size
Fuston Function
10
41,08
15.88
48.99
AvgCE MinCE
40.52
15.69 15.88 (+1.23%) 49.38 (+1.00%)
48.89
5
40.68 (-0.97%) 15.61 (-1.70%) 49.38 (+0.78%)
41.08 (+ (31%)
(a) Probabilistic Token Alignment Parameters.
(b) Fusion Trainning Parameters
4.5
Diagnostic Experiment
Number of source LLMs. In Table 5 we
Table 5: Results of PTA-LLM by incorporating
present the results of fusing varying numbers of LLMs. In general, the performance of PTA- LLM improves as the number of integrated mod- els increases from 1 to 3. However, we also find that the benefits of incorporating additional models vary across different benchmarks (i.e., prominent improvement is observed in the ME). It is also important to highlight that the fusion of lower-performing source models results in diminished performance gains (i.e., MPT, which performs the worst in the MMLU benchmark, contributes the least improvement when we combine one model).
varying numbers (from to 2) of models.
Model
BBHL
MMLL
ME
46,94
CLM
(+1.86%)
(+1.27%)
Openl.aMM
60.54 (+2.11%)
Liano-2 PIU-LLM
40.65 (+2.39%)
48.19
15.78 (x7.896)
41.08 (+3.485)
49.38(95.009)
15.88 (+8.5456)
Optimal Transport Convergence Threshold. As discussed in $22 key hyperparameter in optimal transport is the threshold, which regulates the convergence of the Sinkhorn algorithm (22). A lower value of threshold results in more iterations of transport, enforcing a stricter distribution constraint. As illustrated in Table 44 (up), the lower optimal temperature preference indicates that a stricter constraint may form a more coherent fusion and thus bring greater performance gain.
Token A Alignment Window Size. During the probabilistic token alignment, the default transport window size is the same of the logit length (i.e.. Top-10). Here, explore the impact of window size on the transport of fused logit in Table 4a| (down). In general, larger transport range enable a more comprehensive understanding of the context and thus lead to performance improvement.
Combination Weight. As discussed in $3.1 the combination weight determines the relative emphasis placed on the fused matrix versus the label matrix during continued training. We can observe a higher performance in Table 4b (up) when the weight is smaller within reasonable range (see detailed analysis in $STO). since lower value indicates more emphasis in our fused matrix.
Fusion Functions. In $22 we employ distribution matrix with minimum cross entropy (MinCE) to define the source distribution matrix during training. Additionally, we implement a weighted average of distribution matrices based on cross entropy (AvgCE). A comparison of these approaches is provided in Table (down). The results show that PTA-LLM using MinCE consistently outperforms AvgCE across all benchmarks, which is consistent with M.
Conclusion
We present Probabilistic Token Alignment for Large Language Model Fusion (PTA-LLM), distribution-wise token alignment approach that leverages the optimal transport framework through reformulation. It has merits in: i) demonstrating generality across benchmarks through coherent representation fusion; ii) offering a flexible and adaptive solution to various contexts, especially stable in addressing challenging tasks; and iii) thoroughly investigating the essence of token alignment to elucidate the coherent token we fused. As whole, we conclude that the outcomes elucidated in this paper impart essential understandings and necessitate further exploration within this realm.

---

6 Acknowledgements
This research was supported by the National Science Foundation under Grant No. 2242243, 2348468 & 2450068.
References
[1] Kristine Monteith, James Carroll, Kevin Seppi, and Tony Martinez. Turning bayesian model averaging into bayesian model combination. In IJCNN. IEEE, 2011.
[2] Vipul Gupta, Santiago Akle Serrano, and Dennis DeCoste. Stochastic weight averaging in parallel: Large-batch training that generalizes well. ICLR, 2020.
[3] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale. et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint S.KV::2307.0928,, 2023.
[4] Xinyang Geng and Hao Liu. Openllama: An open reproduction of Ilama, May 2023.
[5] MosaicML NLP Team. Introducing mpt-7b: A new standard for open-source, commercially usable Ilms, 2023. Accessed: 2023-05-05.
[6] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08360 2020.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. NeurlPS, 2020.
[8] Omer Sagi and Lior Rokach. Ensemble learning: A survey. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery, 2018
[9] Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rehecca Roelofs, Raphael Gontijo-Lopes Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In ICML. PMLR, 2022.
[10] Weishi Li, Yong Peng, Miao Zhang, Liang Ding, Han Hu, and Li Shen. Deep model fusion: survey, arXiv preprint 23t.:::230..699990 2023.
[11] Fanqi Wan, Xinting Huang, Deng Cai, Xiaojun Quan, Wei Bi, and Shuming Shi. Knowledge fusion of large language models. arXiv preprint arXiv:2401.1049), 2024.
[12] Yao Fu, Hao Peng, Litu Ou, Ashish Sabharwal, and Tushar Khot. Specializing smaller language models towards multi-step reasoning. In ICML, 2023.
[13] Louis Fournier, Adel Nabli, Masih Aminbeidokhti, Marco Pedersoli, Eugene Belilovsky, and Edouard Oyallon. Wash: Train your ensemble with communication-cfan efficcient weight shuffling, then average. arXiv preprint arXiv:2405.17517 2024.
[14] Donato Riccio, Fabrizio Maturo, and Elvira Romano. Supervised learning via ensembles of diverse functional representations: the functional voting classifier. Statistics and Computing, 2024.
[15] Dongfu Jiang, Xiang Ren, and Bill Yuchen Lin. LIm-blender: Ensembling large language models with pairwise ranking and generative fusion. arXiv preprint (ttp:22366002561, 2023,
[16] Yaqing Wang, Subhabrata Mukherjee, Xiaodong Liu, Jing Gao, Ahmed Hassan Awadallah, and Jianfeng Gao. Adamix: Mixture-of-adapter for parameter-efficient tuning of large language models. arXiv preprint arXiv:2205.12411 2022.
10

---

[17] Chengsong Huang, Qian Liu, Bill Yuchen Lin, Tianyu Pang, Chao Du, and Min Lin. Lo- rahub: Efficient cross-task generalization via dynamic lora composition. arXiv preprint arXiv:2307.1322 13269, 2023.
[18] Jinghan Zhang, Shiqi Chen, Junteng Liu, and Junxian He. Composing parameter-efficien modules with arithmetic operations. arXiv preprint arXiv:2306. 14870, 2023.
[19] Fanqi Wan, Ziyi Yang. Longguang Zhong, Xiaojun Quan, Xinting Huang, and Wei Bi. Fusechat: Knowledge fusion of chat models. arXiv preprint arXiv:2402.16107, 2024.
[20] John Mingers. An empirical comparison of pruning methods for decision tree induction. Machine learning, 1989.
[21] Leif E Peterson.
K-nearest neighbor. Scholarpedia, 2009.
[22] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. NeurIPS, 2013.
[23] Wenguan Wang, Cheng Han, Tianfei Zhou, and Dongfang Liu. Visual recognition with deep nearest centroids. arXiv preprini arXiv:2209.07386 2022.
[24] Max Marion, Ahmet Üstün, Luiza Pozzobon, Alex Wang, Marzieh Fadaee, and Sara Hooker. When less is more: Investigating data pruning for pretraining llms at scale. arXiv preprint arXiv:2309.04564, 2023.
[25] Jean Kaddour. The minipile challenge for data-efficient language models. arXiv preprint arXiv:2304.08442, 2023.
[26] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In NeurlPS, 2019.
[27] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers; State-of- the-art natural language processing. In EMNLP, 2020.
[28] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurlPS, 2022.
[29] Mukund Balasubramanian and Eric I. Schwartz. The isomap algorithm and topological stability. Science, 295(5552):7-7, 2002.
[30] Hervé Abdi and Lynne J Williams. Principal component analysis. Wiley interdisciplinary reviews: computational statistics, 2(4):433 2010.
[31] Nick Littlestone and Manfred K Warmuth. The weighted majority algorithm. Information and Computation, 1994.
[32] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint a1.:4:2110.14168, 2021.
[33] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika. Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for shw-shot language model evaluation, 07 2024.
[34] Mirae Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al. Challenging big- bench tasks and whether thain-of-thought can solve them. arXiv preprint 2rXiv:2210.09266 2022.
11

---

[35] Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao. Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adria Garriga-Alonso, al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint M.Xi::2006.0611.5 2022.
[36] Federico Cassano, John Gouwar, Daniel Nguyen, Sydney Nguyen, Luna Phipps-Costin, Donald Pinckney, Ming-Ho Yee, Yangtian Zi, Carolyn Jane Anderson, Molly Q Feldman, et al. Multipl- c: A scalable and extensible approach to benchmarking neural code generation. arXiv preprint arXiv:2208.082271 2022.
[37] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint (.KVv:200.033:: 2021.
[38] Loubna Ben Allal, Niklas Muennighoff, Logesh Kumar Umapathi, Ben Lipkin, and Leandro von Werra. A framework for the evaluation of code generation models, 2022.
[39]
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. ICLR, 2021.
[40] Thomas Hartvigsen, Saadia Gabriel, Hamid Palangi, Maarten Sap. Dipankar Ray, and Ece Kamar. Toxigen: A large-scale machine-generated dataset for adversarial and implicit hate speech detection. arXiv preprint arXiv:2203.094 00909, 2022.
[41] Jonathan Clark, Eunsol Choi, Michael Collins, Dan Garrette, Tom Kwiatkowski, Vitaly Nikolaev, and Jennimaria Palomaki. Tydi qa: A benchmark for informationn-seeking question answering in ty pologically verse languages. ACL, 2020.
[42] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b, arXiv preprint arXiv:2310.00625, 2023.
[43] An Yang. Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, al. Qwen2. 5 technical report. arXiv preprint arXiv:2412. 151115, 2024.
[44] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In KDD, 2020.
[45] Jianzhou Luo, Dingchuan Yang, and Ke Wei. Improved complexity analysis of the sinkhorn and greenkhorn algorithms for optimal transport. arXiv preprint arXiv:2305. 14939, 2023.
[46] Pierre Colombo, Kevin El Haddad, Céline Hudelot, and Nicolas Boizard. Towards cross-tokenizer distillation: the universal logit distillation loss for lms.. arXiv preprint arXiv:2402.120030, 2024.
[47] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.0400 2024.
[48] Zheng Cai, Maosong Cao, Haojiong Chen, Kai Chen, Keyu Chen, Xin Chen, Xun Chen, Zehui Chen, Zhi Chen, Pei Chu, et al. Internim2 technical report. arXiv preprint arXiv:2403.17297 2024.
[49] Guan Wang, Sijie Cheng, Xianyuan Zhan, Xiangang Li, Sen Song, and Yang Liu. Openchat: Advancing open-source language models with mixed-quality data. In ICLR, 2024.
[50] Runjia Zeng, Cheng Han, Qifan Wang, Chunshu Wu, Tong Geng, Lifu Huang, Ying Nian Wu, and Dongfang Liu. Visual fourier prompt tuning. In NeurlPS, 2024.
[51] Runjia Zeng, Guangyan Sun, Qifan Wang, Tong Geng, Sohail Dianat, Xiaotian Han, Raghuveer Rao, Xueling Zhang, Cheng Han, Lifu Huang, et al. Mept: Mixture of expert prompt tuning as a manifold mapper. In EMNLP, 2025.
12

---

[52] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We present Probabilistic Token Alignment for Large Language Model Fusion (PTA-LLM), a distribution- wise token alignment approach that leverages the optimal transport framework through reformulation. It has merits in: i) demonstrating generality across benchmarks through a coherent representation fusion; ii) offering a flexible and adaptive solution to various contexts, especially stable in addressing challenging tasks The main contributions of PTA-LLM mentioned above are claimed in both the abstract and introduction accurately.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in$SIT in Appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specificattin, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
14

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:
[Yes]
Justification: All theorems used in the paper properly referenced.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proof' can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented
by formal proof's provided in appendix supplemental material. Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We claim reproducibility in both 4.1 and available after acceptance.
S2
Our code will be publicly
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
15

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: We claim reproducibility both $4.1 and $2 All the datasets included in our study publicly available (Minipile). Our code will be publicly available after acceptance. The publicly available code should be adequate to replicate the primary experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to rum to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.co tttpcc//pbbb. wtcc//pbllc//ee helliicgggies//// wttpgdedeeeeeedudbmmmmmeeee www.//ooieeessiiieooca wommissssoooiccy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We specify all the training and test details in $3
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: We report standard deviation error and significance test results for our main result in $2
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out of range (e.g.negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We specify all the training and test computing resources in $S7 and S2
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that NeurlPS Code of Ethics https::pss /ttp:::////eiis.eeeee http///wrips.ccpullll/e////c///t cttpeaueiee/tiiiccccceeelellllll
didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: We conform to the NeurIPS Code of Ethics and show related asset license and consent to our work in SSIT
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
$1I discusses the social impact of our research.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
17

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We show related asset license and consent to our work in SII
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release new assets. In SST we include asset license and consent.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: We only used LLMs for editing purposes (e.g., grammar, spelling, word choice), which does not impact the core methodology or originality of the research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy hatt:e////orrrpssccc/ooofereeee/ for what should or should not be described.
20

---

SUMMARY OF THE APPENDIX
This appendix contains additional experimental results and discussions of our NeurIPS 2025 submis- sion: Probabilistic Token Aligument for Large Language Model Fusion, organized as follows:
covers more details on Related Work.
includes Implementation Details of our PTA-LLM.
introduces the Datasets applied during our experiments.
lists out all the Notations used in PTA-LLM for clarity.
presents more details of implementing Probabilistic Token Alignment.
S6
presents more results of Visualization of Probabilistic Token Alignment.
S7
presents more Pilot Study on heterogeneour fusion, training burden and training time.
S10
provides more hyper parameter settings for Ablative Studies.
S8
provides Per-task Results on Different Benchmarks, where the overall results have been
provided in the main paper.
S9
conducts several Case Studies the model prediction output in specific tasks.
STO
provides more hyper parameter settings for Ablative Studies.
SI1
adds more discussions of Limitations, and points out potential directions of our Future
work.
S1 Related Work
Model Fusion has garnered significant attention as a means to enhance the general performance of LLMs. The fusion techniques be classified into three primary categories: model ensembling, weight merging, and knowledge fusion. Model ensembling combines the predictions of independently trained models improve overall performance. Common approaches include weighted averag- ing [31] majority voting [1]. and pairwise ranking (15). Although model ensembling often leads significant improvements in predictive accuracy and model robustness, requires maintaining multiple models during inference, leading to higher memory consumption and increased latency. This makes it less efficient for resource-constrainee environments. Weight merging combines the parameters of multiple models to synthesize new, unified model. This method is especially effective when the models share identical architectures, their parameters can be merged seamlessly [2] (9] Weight merging is enhanced by linear mathematical operations on adapter parameters, which has proven useful for improving model performance and generalization [16] 17 [8]. Despite these advantages, weight merging suffers from significant limitations: It relies on architectural uniformity across models and requires manual tuning, which constrains its applicability across diverse model architectures (iee. low generalizability)
In contrast, knowledge fusion offers a more flexible and efficient means of integrating models, particularly when the underlying architectures differ (i.e., common case in LLMs). It distills knowledge from multiple teacher models into single student model, transferring the knowledge in more compact and efficient form. One of the key innovations is the minimum edit distance (MinED) token alignment strategy, first introduced by I.. which facilitates effective knowledge transfer by aligning tokens across models. This approach was further refined by [19,, who proposed mapping statistics-base strategy designed enhance conversational model performance. Compared model ensembling and weight merging, knowledge fusion presents more scalable and architecture-anostic solution, making highly suitable for integrating multiple LLMs while minimizing the performance degradation typically associated with stepwise optimization.
Token Alignment was first introduced as a solution to address the misalignment problem between tokenizers with different size of vocabulary, specifically when aligning their respective distributions. The concept was initially formalized by [12]. who employs a search algorithm minimize the aligmment cost between token sequences. This method relies on the assumption that an optimal one-to-one mapping between tokens can be found, enabling the direct alignment of their respective distributions. However, in cases where such a precise mapping is not feasible, the solution defaults to one-hot vector representation, which may oversimplify the complexities inherent in real-world token distributions. Building upon this work, [11] introduced more flexible approach by replacing the exact match requirement with MinED strategy for more robust token alignment, especially in cases where slight variations between tokens could still preserve semantic equivalence. Later, [19
21

---

refined further in cross cross-lingual lingual upplications, incorporating statistical mapping frequencies between source and target tokens better account for the probabilistic nature token co-occurrence, leadning to prominent chat performance.
However, existing methods remain limited by their reliance surface-level token correspondences (i.e.. based solely on the strings it comprises), which leverage minimum edit distance to align the logit. However, besides using edit distance as one metric, our method advances this by incorporating the corresponding logit values into the individual cost within the transport framework. Optimization performed at both the "surface-level" and "logit-level".
S2 Implementation Details
In this section, we present the relevant training details for fusing different LLMs in our PTA-LLM.
Training Time Training is conducted on 8 NVIDIA A100-80GB GPUs (approximately 26 hours for a single epoch) and 8 NVIDIA 1100-80GB GPUs (approximately 17 hours for a single epoch), while conducting evaluation on 4 NVIDIA A100-40GB GPUs (time varies depending on the amount of benchmark data used).
More specifically, the entire runtime analysis can be divided into Stage and Stage 2.
Stage (Training-frred For LLaMA-CLM, token alignment is required, so the runtime for this stage is zero. For Fus Stage takes approximately 3 GPU hours (on NVIDIA A100-80GB) and 4 CPU hours (on AMD EPYC 7763 processor). Compared to FuseLLM, PTA-LLM introduces an additional optimal transport computation, which results in approximately a 13.75% increase in CPU alignment time the MiniPile dataset. It would be valuable explore more efficient alternatives, such as the Greenkhorn algorithm, to reduce the complexity of the optimal transport step.
Stage 2 (Training stage). We report the runtime for PTA-LLM (iee, end-to-end training) under various settings (e.g., different GPUs and dataset sizes) in SS7 Since Stage 2 only involves standard end-to-end training. the corresponding runtime for Fusel. and LLaMA-CL.M are comparable under the same condition
Training Results While our average absolute gain is 1.1% compared to the CLM, the relative gain of 3.94% is nontrivial. In addition, our significance test results indicate that the improvements are statistically significant with p-value 0.005. To assess the stability of our results, we conducted three independent runs using identical hyperparameter settings for each benchmark. The standard deviations of the evaluation scores were 0.05 for MMLU, 0.04 for BBH, and 0.05 for MultiPL-E. Our findings indicate that the performance remains stable once the hyperparameters and training devices are fixed. Other benchmarks show similar trends on standard deviation. We have included these results in the revised paper.
Training Dataset MiniPile is a compact yet diverse training dataset consisting of up million samples, carefully curated from the Pile to preserve the original corpus's 's richness across various domains while maintaining a manageable size for efficient experimentation.
Base Model Tokenizers Both aaM--, developed by Meta AI, and LaAA an open-source reproduction of ILLMA, utilize tokenizers with vocabulary size 32,000 tokens. These two models share a substantial overlap, with 20,079 tokens in common, calculated as the intersection of their vocabularies. In contrast, MPT, developed by MosaicMI., employs a significantly larger tokenizer vocabulary comprising 50,277 tokens, among which only 8,993 tokens overlap with LLaMA-2 and 9,761 with Ope.LLaMA. These distinctions in vocabulary size and token overlap reflect the varying tokenization strategies adopted by different model developers and underscore the implications for compatibility and performance across different language models.
Training Pipeline of PTA-LLM The training procedure consists of two stages: Probabilistic Token Aligument (Offline Phase) and Supervised Training with an Alignment Objective.
Stage I: Probubilistic Token Alignment (Offine Phase). In this stage, we compute fused alignment matrix between all pairs of source and target models using probabilistic token alignment. The process comprises the following steps:
22

---

1. Preprocessing: Segment the MiniPile dataset by splitting long texts into shorter segments for manageable processing.
2. Model Inference: Perform inference on the segmented texts using both source and target models to extract per-step logits and token indices.
3. Optimal Transport Aligument: Apply optimal transport to align the per-step logits and token indices, thereby constructing a fused alignment matrix across source-target model pairs.
Stage 2: Supervised Training with Alignment Objective. In the second stage, train the target model using a joint objective that combines the fused alignment matrix obtained from Stage 1 and the one-hot label matrix derived from the MiniPile dataset. This joint supervision encourages the model to learn both from the aligned token distributions of the source models and the ground-truth labels, facilitating effective knowledge transfer while preserving task-specific accuracy.
Fusion Pipeline of P'TA-LMM PTA-LLM requires only single target model and supports any number of source models (at least one), regardless of their architecture or tokenizer. A key charac- teristic of this framework is that the target model undergoes end-to-end training, where its training objective is defined by combination of the causal language modeling loss (LCLM) and the fusion loss (Lussoon) In essence, while the final fused model retains the architecture of the target model, its parameters are enriched by integrating knowledge from the source models through the fusion process.
To illustrate the fusion process, we use our experimental setup scenario involving an arbitrary number of source models. For fair comparison, we designate LLaMA- 2 as the target model, with MPT and OpenI.aaMa serving as source models. Note that different target model choices may result varying base performance, as discussed prior works such as FUSELLM and FUSECHAT.
example and extend it to
Prior to fusion, we perform forward pass through LLaMA-2, OpenLLaMA-2, and MPT obtain their probabilistic distribution matrices, denoted as P., Ps,, and P.y. respectively, If additional source models (e.g., DeepSeek or Qwen) are included, we similarly forward them to obtain their respective probabilistic matrices Ps,. for the n-th source model.
The fusion process proceeds in recursive manner follows:
Fusion Stage 1: Align P, and Ps, (OpenI.LaMA-2) using the method described in Sec- 3.2, resulting in fused matrix Pt.
Fusion Stage 2: Align Pr,,, and Pap (MPT), resulting in the updated matrix Pr,...82
Fusion Stage 3: Align Pt,sy,32 with Psg (third source model), producing Pt.31.52:33"
Fusion Stage n: Continue recursively, aligning Pt.s with Ps,, to obtain the final fused matrix Pt.e1......."
As demonstrated, our fusion strategy avoids the need for pairwise alignment between every tar- get-source model pair, leading to linear complexity with respect to the number of source models. This efficient and scalable design makes the framework suitable for incorporating large collections of heterogeneous models.
S3 Details of Dataset
The Grade School Math (32), proposed by OpenAl, comprises a wide variety of conceptually simple grade school-level word problems and serves benchmark to assess the shorteomings of language models in handling multi-step mathematical reasoning. We evaluate using the accuracy (8 shot) under the weweevllaioonhaaaee
framework [33].
Big-Bens Beg-Bench Hard (BBH) 34 is benchmark to evaluate the general reasoning ability of LLMs, containing 23 multiple-choice tasks and free-form generation tasks from the Big-Bench (35). We evaluate it using the EM accuracy based on few-shot chain-of-thought (CoT) prompts under the open-instruct framework following [II].
MultiPt.-E (ME) 336] is a multilingual programming benchmark to assess the commonsense ability of L.MMs, consisting of 18 different programming languages with 17 parallel datasets translated from the Python benchmark [37]. We evaluate it using pass@1 [37 based on 20 generated samples
23

---

for each question in 10 popular programming languages under the eeroode-evaluaaion--arrd hardness framework [38 TT].
Measuring Massive Multitask Language Understanding (MMLU) [39 is massive multitask test consisting of multiple- ceatiplee-coice questions from various branches of knowledge assess the commonsense ability of LLMs, including 17 sub categories (i.e., US history, computer science and law) that people must study to learn. We evaluate using the classification accuracy under the open-instruct framework.
dataset for adversarial and implicit hate speech detection used to evaluate the safty ability of LLMs, which contains implicitly toxic and benign sentences mentioning 14 minority groups. We evaluate using the non-toxicity rate (i.e., 1 -----epon reported toxicity rate) under the open- open-instruct framework.
ToxiGen 440] a large-scale
TyDi QA 41 is benchmark for information-secte seeaiinnn-eeiiinggg question answering typologically diverse languages to asses the multilingual ability of LLMs. It covers 9 different languages including korean, arabic, indonesian, etc. We evaluate using the EM accuracy under the open-instruct instruct framework.
S4 Notations
Table S1: The main notations used in PTA-LLM.
Notation
Representation
training corpus consisting of collection of text sequences t Text sequences with a maximum length of 2048 tokens
P.
The probabilistic distribution matrices for the source model, consisting of L tokens A The vocabulary sizes of the source model tokenizer
V
L
The sequence length in P.
A
Tokens in P. have logits over a vocabulary of size V,
P.
The probabilistic distribution matrices for the target model, consisting of N tokens B The vocabulary sizes of the source model tokenizer
V2
N
The sequence length in P.
B
Tokens in P. have logits over a vocabulary of size V2
P
The fused prohabilistic distribution matrices for model fusion, consisting of N takens 8 Fused tokens in Py have logits over vocabulary of size Vi
18
C
The overall training objective of PTA-LLM
LCLM
The training objective casual language modeling training objective kownledge fusion The training objective kownledge fusion
Crusion
Leusion
X
A hyperparameter controlling the trade-off between the LCLM and the Crusion
Q
Target model's 's predictions during the training One-hot label matrix in
O,
T
The initial
optimal transport plan matrix
+
The final n x optimal transport plan matrix of non-negative entries Grm
Gum
The amount of logit probability transported from the n-th source indice to the m-th target indice.
C
An n. x m transport cost matrix of non-negative entries Cry
Cry
The alignment cost between source indice and target indice
Rs
indices corresponding to largest values the z-th row
24

---

S5 Details of Probabilistic Token Alignment
Our training procedures are implemented based on the publicly available code from M. with modifications made specifically to the token alignment module. For better understanding, the following concise pseudo code of $3.2 Specifically, we perform optimal transport on logits after applying the softmax function to reduce the impact of extreme values (e.g., extreme large, small, or negative values) that could otherwise distort transport cost. Importantly, conducting transport in logit space differs fundamentally from transporting mass in probability space due to the distinct normalization terms associated with the source target spaces. We plan to investigate further in the future. Our full implementation shall be publicly released upon paper acceptance.
Algorithm 2 Probabilistic Token Alignment
Require: Tokenizer, input IDs, per step logits, per step indices from both source and target Model.
1: Convert input IDs to token sequence.
2: Use Dynamic Programming in to obtain token pairing between two token sequences.
3: for each token pairing do
4: if it a one-to-one token pairing then
5: use the sinkhorn algorithim in [22 under the optimal transport framework, considering per
step logits and indices from source and target token
6: else
7: use the one-hot logits
8: end if
9:
end for
10: return Aligned matrix
S6
Visualization of Token Alignment
In this section, present more details and results of visualization of token alignment to support findings in $4.4 All samples are the token alignment target model (i.e., Llama) and source model (i.e., MPT).
In Fig ST we can first observe a significant center shift in FUSELLM while our method maintain its overall distribution, showing consistency with our paper.
In S2 and Fig S3 we present more visualization inspection results for FUSELLM and PTA-LLM using Isomap (29 and PCA [30]. Overall, we present additional visual evidence to support the notion that the probabilistic token alignment generate more compact and coherent representation.
Fesuna
fr
300
pril
verage
verage
0 100
part
visit
100
new
Lotal
number report
new
thers
week
trip
report
-30
study
uniqque
short
journey
study
-500
Featurel 500
-500
Feoturet
500
-5h0
Fortual
300
(a) Target Model: Llama
(b) F FuseLLM LM
(c) PTA-LLM
Figure S1: Sample A. 2D visualization results target tokens and fused tokens.
25

---

Fostanc
830
sheep
anddent
2000
medieval
historic
Channinng
medioval
medieval
dram spect
traditional
dram
the and
.Gu
the aand
-200
and
-R00
S00 Yeatural
-500
500 Festarel
-500
Fextanal
(a) Target Model: Llama
(b) FuseL LM
(c) PTA-LLM
Figure S2: Sample 2D visualization results target tokens and fused tokens.
Posteal
400
opportunity
opportuniiy
380
chance
chânce
ancient
changing
chance
wild
char CI
CI
.C1
wwild
famous
wild
farnous
-20
200
famous
30
Ring
spect
magnific
Ring
best
Ring
.magnific
.80
Grant
Fasturcl. 500
560
1000
Featurel 500
6
500
1600
Featuret 500
500
1000
(a) Target Model: Llama
(b) FuseL LM
(0 PTA-LLM
Figure S3: Sample C. 2D visualization results of target tokens and fused tokens.
S7
Pilot Study
S7.1
Heterogeneous Fusion
also acknowledge related work [9), which explores the fusion of Mixtral, InternLM2, and Open- Chat, demonstrating consistent performance improvements under the knowledge fusion paradigm. Motivated by these findings, we conduct a preliminary experiment on MiniPile to extend our frame- work for fusing additional LLMs. Our experimental setup would further benefit from incorporating models with greater architectural diversity. As shown in the Tab. S2 and S3 both "PTA-LLM + Qwen2.5" and "PTA-LLM + Mistral" outperform Qwen2.5 and Mistral, respectively, highlighting the effectiveness of our approach.
Table S2: The baseline models used in this pilot study include more heterogeneous models such as Mixtral and Qwen 43
Benchmark OpenLLaMA
MPT
Llama-2
Mistral
Qwen2.5
Baseline Performance
MMLU
42.11
27.84
46.94
59.15 80.39
71.80
Hellaswag
74.52
76.35
75.99
78.89
26

---

Table S3: The fusion results using the baseline model in Tab. S2 under multiple choice task.
PTA-LLM
PTA-LLM
PTA-LLM
Benchmark
FuseLLM
(ours)
+ Mistral
+ +wen2.55
Fusion Performance
MMLU
48.77
49.38
60.34
72.66
Hellaswag 78.23
79.74
81.52
80.13
For generative tasks, we have already reported results on MultiPL-E the original paper. To further compare our approach with the Mistral and Qwen models, we conducted additional experiments on HumanEval, using "MinnPile+Githuur for training. The results presented below in Tab. [S4] together with our original findings MultiPL-E, further validate the effectiveness of our method for generative tasks.
Table S4: The fusion results using the baseline model in Tab. S2
under generative task.
PTA-LLM
PTA-LLM Qwen2.5
Benchmark
Mistral
Qwen2.5
Mistral
Fusion Performance
EvaanEval
30.50
57.90
32.64
59.71
S7.2 Training Burden on Knowledge Fusion
To further validate the effectiveness of konwledge fusion, we conducted an additional experiment comparing two distinct settings. Setting A: Train on subset of MiniPile (100K examples, randomly sampled as 10% of the original dataset) using the same configuration as in the paper. The total training cost, including obtaining probability distributions from all teacher models (3.2 GPU hours) and end-to-end training ( 20 GPU hours), amounts to 23.2 GPU hours. Setting B: Allocate the same computational budget to train a larger dataset. Specifically, we train Llama-2 CLM on larger MiniPile subset (1 (116K examples, corresponding to 11 1.6% of the original dataset), ensuring the total training cost remains 23.2 GPU hours, matching Setting A.
The results, summarized in the Tab. [S] align with our previous findings, reinforcing the advantages of the knowledge fusion paradigm under equivalent resource constraints.
Table S5: The performance under the same computation cost.
Setting B (Our method) (CLM with more data)
Setting
Benchmark
MMLU
47.42
46.33
S7.3 Training Time
Since the end-to-end training time is primarily determined by the dataset size, we conducted an additional study using subsets of MiniPile in Tab. [S6] which consists of 1M training samples. Specifically, we randomly sampled the original dataset to create training subsets. The results reveal a linear relationship between dataset size and training time. All experiments were performed on NVIDIA A100-80GB GPUs. Notably, we further evaluate the model under different settings. PTA-LLM achieves MMLU scores of 49.38, 48.55, and 47.42 when trained on MiniPile datasets of 1M (100%), , 500K (50%), and 100K (10%) samples, respectively. The results show that larger dataset sizes yield modest performance improvements. Notably. even with only 10% of the data, PTA -LMM attains score of 47.42 on MMLU, which still surpasses the traditional CLM baseline of 47.65.
Table S6: The fusion results using the baseline model in Tab. S2 under generative task.
MiniPile
MiniPile 500K (50%)
MiniPile
MiniPile 10K (1%)
IM (100%)
100K (10%)
Training Time
26.0h
12.9 h
2.4h
50.3h
27

---

S8 Per-task Results on Different Benchmarks
For the training acceleration, we leverage Deepseepd 44 and FlashAttention (28 More specifically, optimize our model using the AdamW optimizer, with hyperparameters
0.9 and
82 = 0.95, applying gradient clipping at 1.0 and a weight decay 0.05. The learning rate follows a cosine schedule, peaking x 10-5, with warmup ratio of 0.008.
To provide comprehensive results from the paper, we report the average per-benchmark results on The Grade School Math, Big-Bench Hard, MultiPL-E, Measuring Massive Multitask Language Understandin, ToxiGen and TyDi QA respectively (see Table 2) We note that the results of all methods in Table 2 have been rerun with the same configuration
machine (i.e., NVIDIA H100-80GB GPUs) and may therefore exhibit slight variations compared to other reports. Furthermore, report per-task results (78 tasks) here in Table S7 for better clarification.
Our results are statistically significant with respect to all baselines on each benchmark (all p-value 0.005). Furthermore, we rerun the same hyperparameter settings three times and computed standard deviation error bars for BBH, MMLUU and ME benchmark.
Table S7: PTA-LLM per-task results six various benchmark.
Task
PTA-LLM
Task
PTA-LLM
Grade School Math
MMLU std=0.08
Grade School Math
14.71
Math
31.30
Health
50.91
Big Bench Hard (BBH) 0.0=
Physics
37.66
Boolean Expressions
68.40
Business
Causal Judgement
50.80
Biology
$3.96
Date Understanding
58.80
Chemistry
36.96
Disambiguation QA
48.00
Computer Science
45.39
Dyck Languages
3.20
Economics
42.59
Formal Fallacies
46.00
Engineering
$1.72
Geometric Shapes
26.80
Philosophy
41 40
Hyperhaton
64.00
Other
$7.94
Logical Deduction (3 ohjects)
59.60
History
$9.57
Logical Deduction (5 objects)
36.00
Geography
53.03
Logical Deduction (7 objects)
26.40
Politics
58.33
Movie Recommendation
69.20
Psychology
55.49
Multistep Arithmetic Two
4.00
Culture
61.45
Navigate
60.00
Law
38.80
Avg. Tasks
49.38
Object Counting
56.40
Penguins in Table
36.30
ToxiGen
Reasoning about Colored Objects
52.40
Black
12.60
Mexican
8.00
Ruin Names
30.00
LGBTQ
24.00
Salient Translation Error Detection
26.40
Jewish
25.20
Snarks
47.19
Women
37.20
Sports Understanding
91.60
Middle East
Temporal Sequences
15.20
Muslim
Tracking Shuflled Objects (3 objects)
30.40
Trans
22.40
Tracking Shuflled Objects (5 objects)
17.20
Asian
36.40
Tracking Shuffled Objects (7 objects)
14.00
Physical Disability
17.80
Web of Lies
64.80
Latino
16.60
Word Sorting
6.00
Native American
6,20
Avg. 27 Tasks
41.08
Chinese
23.20
Mental Disability
11.20
MuliiL-E (ME) std=0.05
Avg. Tasks
C++
9.75
Go
64.51
TvDi QA
Java
9.88
Arabic
9.55
Bengali
21.24
JavaScript
13.85
English
55.23
PHP
9.10
Finnish
43.22
Python
13.87
Indonesian
46.02
5.75
Korcan
$5.80
Ruby
11.58
Russian
33.74
Rust
7.24
Swahili
41.68
TypeScript
13.26
Telugu
0.15
Avg. Tasks
15.88
Avg. Tasks
34.07
28

---

S8.1
Additional Comparison Results
We provide additional comparison results for cases where FuseLLM performs well while performs poorly.
Table S8: Comparison of FuseLLM and PTA-LLM on BBH subtasks (EM accuracy).
Task
FuseLLM LLaMA-2
PTA-LLM
Multistep Arithmetic Two
0.80
4.80
4.00
Temporal Sequences
12.80
16.40
15.20
As shown in the Table above, PTA-LLM exhibits slight performance degradation in two BBH subtasks, both of which evaluate EM accuracy on various reasoning scenarios. PTA-LLM applies OT to balance the probability mass in each row, producing soft assignment matrix that yields more compact and consistent token representations. This smoothing flattens peaks for high-confidence tokens. In exact-match BBH tasks, spreading probability to distractor tokens can shift the argmax and reduce accuracy. PTA-LLM penalized for minor surface deviations, while Fusel.LM's s hard alignment preserves high-confidence tokens, yielding more frequent matches with the official answers. This explains why PTA-LLM achieves the smallest accuracy improvement under BBH in Table 2
$9
Case Studies
In Table S9 Table $10 and Table ST1 we present case studies in difficult tasks (see the overall peformance results in Table to demonstrate the stability of PTA-LLM.
Table S9: Case studies on the Causal Judgement task.
BBH; Causal Judgement;
Question: How would a typical person unswer each of the following questions about causation?: A machine is set up in such a that it will short circuit if both the black wire and the red wire touch the battery the same time. The machine will not short circuit if just one of these wires touches the battery. The black wire is designated as the one that is supposed to touch the battery, while the red wire is supposed to remain in some other part of the machine. One day, the black wire and the red wire both end up touching the battery at the same time. There is short circuit. Did the black wire cause the short circuit?
Options: Yes/ / No
Target: No
FUSELLM: Let's think step by step.Here in this question, we are told that "The black wire is designated as the one that is supposed to touch the battery, while the red wire is supposed to remain in some other part of the machine." A typical person would therefore assume that the black wire was supposed to touch the battery and that the red wire was supposed to remain in some other part of the machine. Because the hlack wire and the red wire both ended up touching the hattery at the same time, typical person would therefore conclude that the black wire caused the short circuit. So the answer is Yes.
PTA-LLM: Let's think step by step. Here this question, we told that "The black wire is designated
the that is supposed touch the battery, while the red wire is supposed to remain in some other part of the machine." A typical person would assume that the black wire is supposed to touch the battery, while the red wire is supposed to remain in other part of the machine. A typical person would therefore conclude that the black wire did cause the short circuit. So the answer No.
29

---

Table S10: Case studies on the Tracking Shuffled Objects (7 objects) task.
BBH; Tracking Shuffled Objects (7 objects);
Question: Alice, Bob, Claire, Dave, Eve, Fred, and Gertrude are dancers a square dance. At the start a song, they each have partner: Alice is dancing with Lola, Bob is dancing with Rodrigo, Claire is dancing with Jamie, Dave is dancing with Izzi, Eve dancing with Sam, Fred is dancing with Karl, and Gertrude dancing with Helga. Throughout the song, the dancers often trade partners. First, Alice and Claire switch partners. Then, Bob and Fred switch partners. Then, Claire and Dave switch partners. Then, Alice and Fred switch partners. Then, Claire and Alice switch partners. Then, Bob and Eve switch partners. Finally, Eve and Gertrude switch partners. At the end of dance, Eve dancing with
Options: (A) Lola (B) Rodrigo (C) Jamie (D) Izzi (E) Sam (F) Karl (G) Helga
Target: (G)
FUSELLM: Let's think step by step. (0) At the start: Alice: Lola, Bob: Rodrigo, Claire: Jamie, Dave: Izzi,
Eve: Sam, Fred: Karl, Gertrude: Helga. (1) Alice and Claire switch partners: Alice: Rodrigo, Claire: Lola.
Dave: Izzi. Eve: Sam, Fred: Karl, Gertrude: Helga. (2) Bob and Fred switch partners: Alice: Rodrigo.
Claire: Lola, Dave: Izzi, Eve: Sam, Fred: Karl, Gertrude: Helga. (3) Claire and Dave switch purtners: Alice:
Rudrigo, Claire: Lola, Bob: Izzi, Eve: Sam, Fred: Karl, Gertrude: Helga. (4) Alice and Fred switch partners:
Alice: Izzi, Claire: Lola, Bob: Rodrigo, Dave: Sam, Eve: Karl, Gertrude: Helga. (5) Claire and Alice switch
partners: Alice: Izzi. Claire: Lola. Bob: Rodrigo, Dave: Sam, Eve: Karl, Gertrude: Helga. (6) Bob and Eve
switch partners: Alice: Izzi. Claire: Lola, Bob: Rodrigo, Dave: Sam, Eve: Karl, Gertrude: Helga, (7) Eve
und Gertrude switch partners: Alice: Izzi, Claire: Lola, Bob: Rodrigo, Dave: Sam, Eve: Karl, Gertrude:
Helga, the end of the dance, Eve dancing with Karl. So the answer is (F).
PTA-LLM: Let's think step by step. (0) At the start: Alice: Lola, Bob: Rodrigo, Claire: Jamie, Dave: Izzi,
Eve: Sam, Fred: Karl, Gertrude: Helga. (1) Alice and Claire switch partners: Alice: Jamie, Claire: Lola,
Dave: Izzi, Eve: Sam, Fred: Karl, Gertrude: Helga. (2) Bob and Fred switch partners: Alice: Izzi, Claire:
Lola, Bob: Rodrigo, Dave: Jamie, Fie: Sam, Fred: Karl, Gertrude: Helga. (3) Claire and Dave switch
partners: Alice: Izzi, Claire: Rodrigo, Bob: Jamie, Dave: Lola, Eve: Sam, Fred: Karl, Gertrude: Helga. (4)
Alice and Fred switch partners: Alice: Sam, Claire: Rodrigo, Bob: Izzi, Dave: Lola, Eve: Jamie, Fred: Karl,
Gertrude: Helga. (5) Claire and Alice switch partners: Alice: Sam, Claire: Izzi, Bob: Rodrigo, Dave: Lola, Eve: Jamie, Fred: Karl, Gertrude; Helga, (6) Bob and Eve switch partners: Alice: Sam, Claire: Izzi, Bob: Rodrigo, Dave: Lola, Eve: Jamie, Fred: Karl, Gertrude: Helga. (7) Eve and Gertrude switch partners: Alice:
Sam, Claire: Izzi, Bob: Rodrigo, Dave: Lola, Eve: Jamie, Fred: Karl, Gertrude: Helga, At the end of the dance, Eve is dancing with Helga. So the answer is (G).
As shown in Table 59 PTA-LLM's response is correct because it accurately identifies the key element of causation in the scenario. The question specifies that the black wire is expected to touch the battery as part of the machine's 's normal setup, while the red wire is not supposed to do so. When the short circuit occurs, the black wire's action is consistent with its intended role and does not deviate from normal functioning. On the other hand, the red wire's unexpected contact with the battery introduces the condition necessary for the short circuit. PTA-LLM correctly reasons that the red wire's abnormal behavior is the true cause of the short circuit, aligning with how a typical person would perceive causation. In contrast, FuseLLM overlooks the normalcy of the black wire's s role and incorrectly attributes causation it, simply because both wires were involved. This makes PTA-ILMM's 's reasoning more logical and consistent with the principles of causation.
As shown in Table STO tracking shuffled objects task with seven objects is a particularly challenging scenario requiring accurate tracking of the corresponding dancers among seven individuals as they switch partners many times.In this context, FuseLLM fails to track the objective during the fourth partner switch, whereas PTA-LLM successfully tracks the corresponding dancers throughout. This superior performance is likely attributable to P'TA-LM''s probabilistic token alignment mechanism, which effectively transforms logits into the correct objective rather than merely replicating the original logits in the FuseLLM approach.
As shown in Table STT PTA-LLMM correctly identifies the SVG path as forming a hexagon, recog- nizing 7 commands: one "M" to start and six "I." commands creating closed six-sided polygon. FuseL.LM miscounts the commands, identifying only 5, and incorrectly concludes the shape is a pentagon. PTA-LLM's accurate command count and shape identification make its reasoning correct.
30

---

Table S11: 1: Case studies the Geometric Shape task.
BBH; Geometric Shape;
Question: This SVG path element <path d= 60.00,77.00
75.00,55.00 0.0.7770000..0000550000000000000
3.00,96.00 93.00,90.00 L 60.00,77.00"1, draws
Options: (A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon (H) rectangle (I)
sector (J) triangle"
Target: (C)
FUSELLM: Let's think step by step. This SVG path element contains "M" and "L" commands. M takes two parameters (x.y) and moves the current point to the coordinates (x.y). takes two parameters (x.y) and draws line from the previous coordinate to the new coordinate (x.y). This path can be decomposed into 6
separate commands. (1) M 30.00,77.00 00: Move the current point 60.00.77.00 (2) 75.00.55 00: Create line from 60.00,77.00 .0 to 75. 0...00.55.00 (3) L 6.00,60.00: Create line from 75.00,55.00 to 6.00,60.00. (4) L 22. 22.00.78.00 Create line from 6.00,60.00 22.00.78.00 (5) L 3.00.96 00: Create line from 22.00. 22.00.78.00 3.00,96.00. 00. (6) L 93.00.90.00: Create line from 3.00,96.00 93.00.90.00. This SVG path starts point 60.00.77.00. creates five consecutive and louching lines, and then returns back its starting point, thereby creating five-sided shupe. It does not have any curves arches. "pentagon" is the only five-sided polygon on the list. So the answer is (G).
PTA-LLM: Let's think step by step. This SVG path element contains "M" and "L" commands. M takes two parameters (x.y) and moves the current point to the coordinates (x.y). takes two parameters (x.y) and draws line from the previous coordinate to the new coordinate (x.y). This path can he decomposed into 7 separate commands. (1) M 60.00,77.88
Move the current point to 60.00,77. 77.00. (2) 75.00,55 00: Create a line from 60.00,77.00 to 75. 25,00,55.00 00. (3) L 6 6.00,00000 Create a line from 75.00,55.00 5,.00 to 6.00,60.00 (4) L 22.00,78 780,78.00 00: Create line from 6.00,60.00 to 22.00,8.00 (5) L 3 8,0006600000 60: Create line from 22.00,78.00 to 3.00,96.00. (6) I. 93.00,90.00; Create a line from 3.00,96.00 to 93.00,90.00. (7) 60.00,77.000 Create a line from 93.00,90.00 to 60.00,,7.00 This SVG path starts at point 60.00,77.00 creates six consecutive and touching lines, and then returns back its starting point, thereby creating a six-sided shape. It does not have any curves or arches. "hexagon" is the only six-sided object on the list. So the answer is (C).
S10
Ablative Studies
Table S12:
Ablative studies of optimal transport convergence threshold
Choice
BBH
ME
MMLJ
Optimal Transport Convergence Threshold
Ie-3
39.44
15.10
48.23
le-4
40.54
15.88
48.99
5c-5
40.91
15.85
49.32
le-5
41.08
15.82
49.38
le-6
41.04
15.78
49.35
1c-7
41.05
15.80
49.33
As shown in Table S12 the findings on the optimal transport convergence threshold align with our motivation. Specifically, a lower threshold preference suggests that stricter constraints may generate a more coherent fusion, leading to greater performance gains. We also observe that performance stabilizes when the threshold drops below le-5, suggesting that the transported cost is fully optimized and remains unchanged.
Table S13: Ablative studies of token alignment window size
Choice
BBH
ME
MMLU
Token Alignment Window Size
10
41.08
15.88
48.99
7
40.99
15.73
49.00
40.68
15.61
49.38
39.64
15.08
47.11
The motivation for using an alignment window is that, without constraints on m and n. the transport space between the logits of Model and Model would be the product of their vocabulary sizes (e.g., MPT has 50,277 tokens while LLaMA has 32,000). This is computationally inefficient, as it requires storing over 30k logits per token instead of just the top 10, and unnecessary because the logit distribution follows long-tail pattern where the top-k logits capture most of the probability mass. Our analysis of token alignment window sizes, ranging from 10 to 3, shows that small window is sufficient for effective knowledge fusion. As shown in Table S13 larger windows generally improve
31

---

BBH and ME performance, with size 10 performing best on both. MMLU peaks at size 5, while very small windows (e.g., 3) reduce performance across tasks. Overall, window sizes between 5-10 offer consistently strong results.
Table S14: Ablative studies of combination weight
Choice
BBH
ME
MMLU
Combination Weight
0,90
40,39
15.72
48.93
0.85
41.00
15.91
49.09
0.80
41.08
15.88
49.38
0.75
39.78
15.65
47.29
0.70
38.11
14.27
46.08
0.60
37.20
14.09
45.96
As shown Table S14, further reveals that the observed "higher performance when the weight is smaller" pertains specifically to the comparison between 0.8 and 0.9. However, if the weight is reduced further, the model overemphasizes the fused matrix and pays less attention to the original CLM modeling. Consequently, we selected 0.8 for all experiments, as it consistently achieves the best performance.
S11 Future Work and Discussions
Limitation. A limitation of our approach is that the Sinkhorn- Knopp algorithm runs in Ö(n² e time, which reduces the token alignment efficiency. Despite the observation that in practice only Sinkhorn loops per training iteration are often sufficient for model representation, which amounts to ~ 13..55% aligning delay on MiniPile compared with FUSELLM. It would be interesting to investigate further lower complexity (i.e., greenkhom [45) algorithim to compute the optimal transport.
Future Work. Despite PTA-LLM systemic generality (see $4.2 and robustness (see $4.3), also comes with new challenges and unveils some intriguing questions. For instance, the overall pipeline is divided into two stages: alignment and fusion training. This naturally raises an important question from a paradigm perspective: Can design end-to-end fusion pipeline that dynamically controls token aligmment, thereby enabling more comprchensive capability learning? Introducing a new loss design (i.e., universal logit distillation loss (46)) within the fusion training deal with the mis- alignment problem in different tokenizers might enhance pipeline efficiency and facilitate additional performance improvements. Another essential future direction deserving of further investigation is its further effectiveness exploration in other NLP fields since aligning sequences generated by different tokenizers is a generic problem of contemporary NLP. In $4.4. we demonstrate through visualization studies that probabilistic token alignment yield a more conherent fused representation. Consequently, the applicability of this integration to other alignment methods requires further investigation.
In this paper, we do not fully explore the potential of knowledge fusion, as comprehensive experiments on heterogencous models remain outside the scope of our study. However, related work [19 has investigated the fusion of models such as Mixtral [47], InternL.22 [48), and OpenChat [49], demonstrating consistent performance improvements within the knowledge fusion paradigm. We plan to explore it further in the future. Inspired by [50] 511. PTA-LLM could also incorporate learnable soft prompts, enabling task- adaptation without retraining the entire alignment matrix.
Besides the directions mentioned earlier, we identify several additional promising avenues for exploration. First, an end-to-end fusion pipeline could streamline the process and reduce the reliance on CPU resources by eliminating the need for a two-stage approach (alignment followed by training). This could be facilitated by leveraging innovative loss functions to enable dynamic adjustments. Second, the exploration of N-1 and I-N mapping strategies offers enhanced flexibility. While this paper focuses on 1-1 mapping due to constraints imposed by traditional optimal transport frameworks, future work could explore beyond these limitations. Lastly, multilingual alignment, such aligning Chinese and English tokens, holds the potential to broaden applicability. as current research predominantly focuses on English token alignment.
Discussion. Two potential factors may explain why the knowledge fusion objective outperforms the traditional CL. approach: First, the CLM objective employs one-hot vectors as the golden labels, which fails to capture the nuanced information each token might convey. This approach provides the
32

---

same penalty for completely incorrect predictions as for predictions that select an incorrect token but retain semantically relevant context. In other words, the CLM objective does not reward predictions that are "almost correct," which limits its capacity to encourage fine-grained improvements. Second, the fusion objective incorporates representations from diverse source models through distillation, enabling it to capitalize on the complementary strengths of each model. It provides more fine-grained context information for alignment.
Regarding the performance, our performance improvements are constrained by the suboptimal performance of certain source LLMs relative to the target LLM on specific tasks, which inevitably impacts the quality of the fusion results. We also observe that the performance improvement could be significantly enhanced by increasing the size of the continued training datasets. Notably, the original MiniPile [25] comprises only 8% coding-related data. By incorporating the GitHub datasets from the Pile in our priliminary experiments, is possible achieve greater performance gains, particularly in coding-related downstream tasks.
Regarding the motivations behind the Equation [ To achieve coherent fused metric as described in our introduction, we aim to keep the logit distribution more compact and consistent with the target distribution. Due the constraints optimal transport, we discretely select the maximum value in each sub-transport, which helps maintain both compactness and consistency. We are also interested in whether broader transport (e.g., choosing top-k instead of top-1) could further improve fusion, and we plan to investigate this in future work.
Regarding the use of edit distance. Our motivation for using minimum edit distance the evalua- tion metric comes from the scale of our IM-sample training corpus, where even a small increase in complexity could create a significant computational burden for token alignment. A semantic similarity-basee token-level metric could embed more comprchensive and reasonable information for transportation. We are very interested in this direction and plan investigate it further, exploring more efficient semantic similarity approaches that are well-suited to our setting.
Potential Risks. Consider the tuning process of LLM, which has potential risks for energy usage. Finetuning requires significant computational power, leading to high energy use and increased environmental impact.
Asset License and Consent. The OpenI.LaMA bigeode-vallutiion-nnness 38 and MPT
[5] are licensed under Apache-2.0; Llama-2 7b 3 is licensed under Llama 2 Community License Agreement im-evaluation-harme harness
3 is licensed under MIT]
Artifact Consistent With Intended Use. Our work ensures that the use of existing artifacts aligns with their intended purpose when specified. For the artifacts create, remains compatible with the original access conditions. In particular, we ensure that derivatives of data accessed for research purposes are confined research contexts.
Social Impact. This work introduces PTA-LLM, which demonstrates significant performance improvements over state-of-the-art baselines, as shown in Table 2 Our approach enhances model accuracy and is particularly beneficial for efficient training scenarios, such as on resource-constrained devices and rapid adaptation with limited computational overhead.
33


Reviews:
[{"review_id": "LCPMKGcGEL", "reviewer": "Reviewer_seY2", "summary": "The paper aims to tackle the problem of vocabulary alignment in knowledge fusion across models. The method proposed seeks to obtain a probabilistic token alignment using optimal transport theory, in contrast with past work where the mapping was hard. Additionally, the proposed method for alignment leverages the top-k predictions. Experiments in merging three models show that the resulting fusion method obtains better results than any base model, continuously pre-training the base model and a fusion method that uses hard mappings.", "questions": "n/a", "limitations": "n/a", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strenghts:\n- original idea to use optimal transport for token alignment. Based on previous work, this is a natural evolution in this area of research\n- the problem is well motivation and theoretically sound\n- ablation experiments\n- study of interpretability\n\nWeaknesses:\n- experiments are not varied enough i.e. only a single setup is presented with 3 models to fuse from. I would expect at least 2 merging setups, ideally with models to merge from different families, and with models of different sizes. This would improve confidence in the robustness of the method\n- results show improvements over the CLM setup, but the improvements when compared to FuseLLM, which is the most directly related work, are quite small (+0.49 on average, but gains on 3/6 datasets are 0.32 or smaller.", "was_revised": true, "final_justification": "The rebuttal provided information I already had when making my original review, so I will maintain my overall opinion, which is borderline accept.", "rebuttal": "Dear reviewer seY2,\n\nWe sincerely appreciate your time and effort in reviewing our paper and providing valuable comments. We provide explanations to your questions point-by-point in the following.\n\n**Q1: Regarding Model Fusion Settings.**\n\nThank you for the excellent suggestion. We fully agree with the reviewer that incorporating more fusion settings would strengthen confidence in the robustness of the method. Therefore, we have indeed included experimental results on model fusion with additional models, including Mixtral and Qwen2.5, in Appendix S7. As shown in these tables, both “PTA-LLM + Qwen2.5” and “PTA-LLM + Mistral” consistently outperform Qwen2.5 and Mistral, respectively, demonstrating the effectiveness of our approach within the knowledge fusion paradigm. We plan to explore fusion with a larger number of models in the future.\n\n**Q2: Regarding the Improvements.**\n\nThank you for your insightful observation. Regarding performance, while our average absolute gain over CLM is 1.1%, we argue that the relative gain of 3.94% is substantial. Furthermore, our significance test results confirm that these improvements are **statistically significant**, with a p-value < 0.005. \n\nOur visualization results further support this claim. The **noticeable distortion** in token representations, as shown in Figures 3 and 6, highlights the **limitations of FuseLLM**. In contrast, the well-preserved representations in PTA-LLM demonstrate its strong potential for broader and more general model fusion scenarios.\n\nWe appreciate your thoughtful comments. We hope our response addresses your concerns. Please let us know if there are any additional questions, and we will be happy to discuss further."}, {"review_id": "jBthrBffre", "reviewer": "Reviewer_8QeG", "summary": "This paper addresses the problem of token alignment and determining the probabilities of tokens in a situation when one wishes to fuse token probabilities from one or more language models. \n\nIf I understand correctly, the main innovation is the fused probability assignment after token alignment. The paper proposes to first find a transport plan between two distributions over vocabularies of different sizes. The fused token probabilities are determined based on this plan.", "questions": "The assignment in Equation (5) is a key. However, it appears to have no clear interpretation in transport theory so maybe a better discussion or another visual can help the reader appreciate this key point.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The problem is crisp with a clear motivation. \n- Visuals are mostly neat and helpful. \n- Experimental results demonstrate the advantage of the method. The methodology seems comprehensive. \n\nWeaknesses:\n- It is unclear to me why mixing with CLM loss. Shouldn't model fusion only consider merging logits from the models? \n- The challenge in token pairing (Lines 135-146) is a bit unclear. The authors state that the challenge lies in the large search space, but to me, it seems that the main issue is ranking different pairings, as there is no immediate criterion for a good pairing. \n- I also identified several writing issues. \n    - The concepts of source and target models arise somewhere in Line 40, but it is unclear how these concepts relate to the problem of model fusion discussed previously. Indeed, the discussion so far on model fusion appears to be that we take information from different models symmetrically.  \n    - It is unclear if the transport is over logits or probabilities. I assume the latter, but this is not clear from the description.\n    - In Line 117, the \"target's model predictions\" is Q_t, but in Line 125 you use P_t for the target model's predictions. There appears to be confusion.", "was_revised": true, "final_justification": "The rebuttal settled several questions and uncertainties that I had. \nI think that this is a good paper that should be accepted, as my initial rating says. I did not see room to improve this rating further.", "rebuttal": "Dear reviewer 8QeG,\n\nWe sincerely appreciate your time and effort in reviewing our paper and providing valuable comments. We provide explanations to your questions point-by-point in the following.\n\n**Q1: Regarding the Use of CLM Loss in Model Fusion.**\n\nThanks for the great question, and we'd like to provide further explanation.\n\nWhen we fuse models through “knowledge fusion,” the process is similar to distillation, transferring the fused logits into the student model. However, the student still needs to learn to predict the actual next token on real data, similar to standard LLM distillation. Specifically, we ensure that the fused model continues to align with real data while absorbing knowledge from multiple teachers. Without the CLM term, there is a risk that the model will merely echo the fused logits and lose fidelity to the true next‑token distribution. Incorporating CLM anchors the fusion to real‑world performance.\n\n**Q2: Regarding the True Challenge in Token Pairing.**\n\nThank you for your insightful observation. We would like to further clarify and elaborate on the token pairing process. Specifically, the **first** step is to establish a ranking criterion for determining the optimal pairing. We use minimum edit distance as an efficient evaluation metric, given the 1M‑sample training corpus. **Then**, once a reasonable criteria is defined, the main challenge shifts from how to pair to how to pair it efficiently. Dynamic programming reduces the exponential path enumeration to $ O(LN)$, ensuring that pairing remains tractable even for sequences with thousands of tokens.\n\nWe will refine the relevant writing in the revised version to ensure better coherence and enhance overall clarity.\n\n**Q3: Regarding the writing.**\n\n**Ambiguity in Source and Target Model Definitions:**\n\nSorry for the confusion. The definitions of source and target Model are intended to clarify our final base model (the target model) and the direction of knowledge flow (from source to target). As illustrated in Figure 1(c), knowledge from the “cat” and “rabbit” source models is fused into the “llama” target model. We will revise the text to enhance the clarity of our paper.\n\n**Transport objective:**\n\nThank you for your insightful comment. As detailed in Appendix S5, we perform optimal transport on logits after applying the softmax function (i.e., probability) to reduce the impact of extreme values (e.g., extreme large, small, or negative values) that could otherwise distort the transport cost. We retain the term ‘logit’ for uniform notation in our paper. Importantly, conducting transport in logit space differs fundamentally from transporting mass in probability space due to the distinct normalization terms associated with the source and target spaces. We plan to investigate i further.\n\nWe will move this discussion from appendix to the main paper for better clarification.\n\n**Notational Inconsistency in Target Model Predictions:**\n\nThank you for pointing this out. This is a typo, and it should be $Q_t$ in line 125. The reason we use two notations is to distinguish between the alignment matrix $P_f$ and the final training matrix $Q_t$. The former is obtained from token alignment (without gradient updates), while the latter represents the logits output (with gradient updates) used to align $P_f$ and $O_t$.\n\n**Q4: Regarding Interpretation and Clarity of Equation (5) Assignment**\n\nThank you for your excellent question. We would like to elaborate on the motivation behind Equation 5 and its empirical visualization.\n\n**Motivation:** To achieve a coherent fused metric as described in our introduction, we aim to keep the logit distribution more compact and consistent with the target distribution. Due to the constraints in optimal transport, we discretely select the maximum value in each sub‑transport, which helps maintain both compactness and consistency. We are also interested in whether broader transport (e.g., choosing top‑k instead of top‑1) could further improve fusion, and we plan to investigate this in future work.\n\n**Empirical Visualization:** As shown in Fig. 3, we can visually compare the distribution of PTA‑LLM fused tokens with the target tokens and the FuseLLM fused tokens. Specifically, Fig. 3(b) and Fig. 3(d) show that PTA‑LLM achieves a more consistent marginal feature distribution with the target tokens, whereas FuseLLM exhibits significantly greater distortion in the overall token representation. The more compact and coherent token distribution after applying probabilistic token alignment is consistent with the motivation described above.\n\nWe will include this elaboration and add a clarifying note in the revision.\n\n\nWe appreciate your thoughtful comments. We hope our response addresses your concerns. Please let us know if there are any additional questions, and we will be happy to discuss further."}, {"review_id": "kJaxeWd14F", "reviewer": "Reviewer_sCqG", "summary": "This work introduces a token alignment method for knowledge fusion, designed to merge multiple source LLMs into a single target LLM.\n\nThe method makes three main contributions:\n\n1.  It allows for flexible, **many-to-many mappings** between the tokens of the source and target models. This avoids the common issue of requiring a rigid, one-to-one mapping between tokenizers that have vastly different vocabulary sizes, training data, or training algorithms. The output of this stage is a sequence of matched token pairs.\n\n2.  It aligns the **next-token probability distributions** for each token pair, rather than just aligning tokens based on string similarity. This process is framed as an **optimal transport problem** where the goal is to \"transport\" the probability from the source model's predictions to the target's with minimal cost. The transport cost considers the edit distance between potential predictions, while the transport itself preserves the probability mass between the two distributions. This is solved using the standard Sinkhorn algorithm.\n\n3.  When fusing multiple source LLMs, it uses a recursive merging algorithm that has **linear complexity** with respect to the number of models and prioritises higher-quality models based on a heuristic, avoiding a potential combinatorial explosion.\n\nWhen used in a knowledge fusion pipeline, these modifications yield small but significant performance gains across a wide range of common downstream tasks.", "questions": "- The knowledge fusion pipeline requires running many models to generate training data, which is computationally expensive and would be almost untenable at scale for modern datasets with tens of trillions of tokens. If one uses such a dataset, how should they choose how much to subsample the data in order to use your method? Could they use it for a small targeted subset, such as FineMath?\n\n- Your work focuses exclusively on the knowledge fusion paradigm. Why was there no empirical comparison to other common methods like fusing hidden states via summing/averaging and then appending new blocks afterwards or alternative stitching or model merging methods? A comparison, even on a smaller scale, would be critical for understanding the trade-offs of your method.\n\n- The optimal transport cost function is based on edit distance. Have you experimented with using semantic similarity (e.g., from token embeddings before the unembedding matrix) as the cost metric instead? Previous works, such as [1,2] have shown that it may be feasible to translate between the latent spaces of two models and take distances there instead of using string distance.\n\nCitations:\n- [1] Moschella et.al; \"Relative representations enable zero-shot latent space communication\"\n- [2] Alvarez-Melis et.al; \"Gromov-Wasserstein Alignment of Word Embedding Spaces\"", "limitations": "The authors thoroughly discussed the limitations of their work.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths:\n* The work is well-written, with clear explanations and a logical flow.\n* The proposed methods are well-justified and effectively address the core challenges of knowledge fusion pipelines. The reformulation of token alignment as an optimal transport problem is a principled approach.\n* The authors provide ample empirical evidence to validate the downstream performance improvements, demonstrating the method’s reliability.\n* The token pairing and alignment techniques are generalizable. They could be transferred to other machine learning domains that require aligning outputs, such as early-exit models or multi-modal models that rely on discrete tokenizations.\n\n### Weaknesses:\n* While consistent, the magnitude of the performance improvements is small in absolute terms.\n* The baseline knowledge fusion framework is inherently expensive because it requires obtaining predictions from N models to create the fused training data. This high upfront cost may be a significant drawback for modest performance gains.\n* The paper’s scope is narrowly focused on improving a specific paradigm—knowledge fusion—and does not compare its method against other prominent model combination techniques. For instance, there are no empirical comparisons to alternative approaches like model stitching and model merging.\n* While the goal is to handle heterogeneous LLMs, a large majority of trained LLMs are based on a GPT-2-like structure with highly similar vocabulary sizes. It would have been beneficial to see controlled experiments with very different architectures. If insufficient diversity is available in the wild, training very small models with different structures would be a sufficient proxy.", "was_revised": true, "final_justification": "The author's rebuttal has addressed a majority of my concerns while opening up new avenues of exploration (using hidden representations instead of logits). While I have reservations regarding the practical applicability of the method at scale (specifically concerning serving and storage), it is a compelling scientific exploration and will serve to inform other members of the community.", "rebuttal": "Dear reviewer sCqG,\n\nWe sincerely appreciate your time and effort in reviewing our paper and providing valuable comments. We provide explanations to your questions point-by-point in the following.\n\n**Q1: Regarding the improvements.**\n\nThank you for your insightful observation. Regarding performance, while our average absolute gain over CLM is 1.1%, we argue that the relative gain of 3.94% is substantial. Furthermore, our significance test results confirm that these improvements are **statistically significant**, with a p-value < 0.005. \n\nOur visualization results further support this claim. The **noticeable distortion** in token representations, as shown in Figures 3 and 6, highlights the **limitations of FuseLLM**. In contrast, the well-preserved representations in PTA-LLM demonstrate its strong potential for broader and more general model fusion scenarios.\n\n**Q2: Regarding Computational Cost of Token Alignment.**\n\nThis is an excellent question. We fully agree with the reviewer regarding the concern about the additional burden of token alignment. To further demonstrate the effectiveness of our method, we conducted an additional experiment in Appendix S7.2. to compare two distinct settings.\n\nSetting A: Train on a subset of MiniPile (100K examples by randomly sampling 10% of the original dataset) using the same setting in the paper. Specifically, the total training costs, including obtaining the probability distributions of all teacher models (approximately 3.2 GPU hours) and end-to-end training (approximately 20 GPU hours), are 23.2 GPU hours.\n\nSetting B: Use the additional cost to train on more data. Specifically, we train Llama-2 CLM on a larger MiniPile subset (116K training examples, representing 11.6% of the original dataset) to match the total training costs of Setting A (23.2 GPU hours).\n\nThe results, summarized in the table below, are consistent with our findings, demonstrating the benefits of the knowledge fusion paradigm under equivalent resource constraints. We sincerely appreciate your constructive suggestions.\n\n| Same Cost | Setting A (Our method) | Setting B (CLM with more data) |\n| -------------- | ---------------- | ------------------------------ |\n| MMLU           | 47.42            | 46.33                          |\n\n\n**Q3: Regarding Lack of Comparison with Other Model Combination Techniques**\nThank you for your excellent question. Most traditional fusion methods, such as model merging and fusing hidden states, require identical model architectures and are therefore not directly applicable in our context, in which we fuse heterogeneous LLMs.\n\nHowever, we further compare PTA‑LLM with an ensemble method for LLMs, LLM‑Blender [ref1], which ranks and combines the output texts from multiple LLMs (and can be applied in heterogeneous LLM fusion) using ranker and fuser models.\n\n| Model                   | BBH  |\n| ----------------------- | ---- |\n| OpenLLaMA               |  33.87    |\n| MPT                     |   33.38   |\n| Llama-2                 |   39.70   |\n| LLM-Blender (Rank&Fuse) |   24.48   |\n| LLM-Blender (Rank)      |  37.17    |\n| PTA-LLM                 |  41.08    |\n\nFrom the results above, we observe that PTA‑LLM still achieves superior performance. We will include these results in our revision to enhance the quality of our paper.\n\n[ref1] LLM-Blender: Ensembling Large Language Models with Pairwise Ranking and Generative Fusion\n\n**Q4: Regarding Model Fusion Settings.**\n\nThank you for the excellent suggestion. We fully agree with the reviewer that incorporating more fusion settings would strengthen confidence in the robustness of the method. Therefore, we have indeed included experimental results on model fusion with additional models, including Mixtral and Qwen2.5, in Appendix S7. As shown in these tables, both “PTA-LLM + Qwen2.5” and “PTA-LLM + Mistral” consistently outperform Qwen2.5 and Mistral, respectively, demonstrating the effectiveness of our approach within the knowledge fusion paradigm. We plan to explore fusion with a larger number of models in the future.\n\n\n**Q5: Regarding Scalability and Data Subsampling in Knowledge Fusion.**\n\nThis is a great question. We understand the reviewer’s concern that the scalability of real‑world datasets is quite different from the 1M‑sized MiniPile used in our training. Therefore, determining how to effectively subset the data becomes an important problem. We further investigate subsample training in Appendix S7.3. The results show that random sampling from the original dataset demonstrates both robustness and effectiveness for PTA‑LLM. We will incorporate this discussion into our revision to improve the quality of our paper.\n\n**Q6: Regarding Use of Edit Distance vs. Semantic Similarity in Transport Cost.**\n\nThank you for your insightful suggestion. Our motivation for using minimum edit distance as the evaluation metric comes from the scale of our 1M‑sample training corpus, where even a small increase in complexity could create a significant computational burden for token alignment. A semantic similarity‑based token‑level metric could indeed embed more comprehensive and reasonable information for transportation.\n\nWe are also very interested in this direction and plan to investigate it further, exploring more efficient semantic similarity approaches that are well‑suited to our setting.\n\nWe appreciate your thoughtful comments. We hope our response addresses your concerns. Please let us know if there are any additional questions, and we will be happy to discuss further."}, {"review_id": "rs41HRY6Tv", "reviewer": "Reviewer_zYWa", "summary": "This paper aims to fuse language models with different tokenizers together to form a stronger model. It is challenging because both the token sequence length and vocabulary size might vary. To address the token sequence length discrepancy, the authors employ a dynamic programming algorithm to compute the optimal mapping. To address the vocabulary discrepancy, the authors model the problem as an optimal transport question and solve it to get align the two distributions between two models. Experiment results confirm the proposed method lead to better model performance compared to previous model fusion and non-fused models.", "questions": "1. Is the recursion function defined in (1) exactly same as the recursion function defined in [1] (Eq (1) (2) (3))? Why it is always +C(B_k, A_j) for all three cases? I think it should be C(B_k) and C(A_j) for the first two lines, which means the cost of adding a new token.\n\n2. When defining recursion function (1), do two sequences of tokens represent the same text? If so, then I think it is impossible that A2 matches B1 while A1 matches B2, correct? Because the order of tokens should follow the words in the sentences. But in Figure 2b(1), it illustrates mappings with perturbed orders. Why is it?\n\n3. I think the process to compute token alignement and optimal transport is done on CPU and hard to be parallelized. I wonder its relative time cost compared to conventional pre-training.\n\n4. Table 3 in Section 4.3 compares the performance of PTA-LLM and FuseLLM on seven tasks that FuseLLM perform worst. The authors claim PTA-LLM is more stable because it has better performance on these tasks. But the task selection is biased. The authors should report the tasks where PTA-LLM has the worst performance and compare the performance degradation with FuseLLM. I think the current setting is problematic.\n\n5. I think Section 4.4 and Figure 3 are confusing. What is the figure trying to show?\n\n6. Why in Eq (4) n=m=10?\n\n[1] Fu, Yao, et al. \"Specializing smaller language models towards multi-step reasoning.\" International Conference on Machine Learning. PMLR, 2023.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "S1. I think the proposed optimal transport method to align logits is reasonable and sound.\n\nS2. The experiment results show consistend improvement on model performance.\n\nW1. The dynamic token pairing technique is proposed by previous work, but the paper seems to claim it is original.\n\nW2. Some experiment settings are potentially problematic.\n\nFor details of the weakness, see my questions below", "was_revised": true, "final_justification": "The authors addressed my concerns. So I maintain the positive score.", "rebuttal": "Dear reviewer zYWa,\n\nWe sincerely appreciate your time and effort in reviewing our paper and providing valuable comments. We provide explanations to your questions point-by-point in the following.\n\n**Q1: Regarding Originality Claim of Dynamic Token Pairing.**\n\nSorry for the confusion. Our token pairing is greatly inspired by [ref1] (line 88), which serves as a foundation for PTA‑LLM. We include its details in Section 3.2 to help readers better understand token pairing as the basis for our key contribution, Probabilistic Token Alignment through Optimal Transport, which is built upon previous token‑level pairing. We recognize that this presentation may cause confusion regarding proper credit. We sincerely appreciate prior work in establishing a reliable and efficient token pairing method and will revise the text to make this clearer.\n\n[ref1] Specializing smaller language models towards multi-step reasoning\n\n**Q2: Regarding Correctness of Recursion Function Definition in Equation (1).**\n\nThank you for raising this point. Our recurrence in Eq. (1) is mathematically identical to Eqs. (1)–(3) in [1]; we simply present it in a more compact form:\n\n$f(k,j) = \\min\\lbrace f(k-1,j),\\; f(k,j-1),\\; f(k-1,j-1) \\rbrace + C(B_k,A_j)$\n\nOur ultimate objective is to optimize $f(L, N)$f, whose minimum value is determined by the optimal path through the sum of subfunctions $ f(k, j)$. The choice of direction at each step depends solely on the relative magnitudes of $ f(k-1, j)$, $ f(k, j-1)$, and $ f(k-1, j-1)$; the absolute value of $ C(B_k, A_j)$ does not affect which predecessor is selected. Consequently, $ C(B_k, A_j)$ should be included in equations (1)–(3), since all three directions represent potential optimal paths. This allows $ C(B_k, A_j)$ to be factored out and added after taking the minimum over the three candidate predecessors, which is consistent with the formulation in traditional dynamic programming.\n\nWe recognise that our presentation may cause confusion, so we will add an explicit derivation and a visual dynamic-programming matrix in the appendix to make the alignment path clearer for readers.\n\n**Q3: Regarding Token Order Consistency in Recursion Function and Figure 2b(1).**\n\nThank you for your insightful observation. Your understanding is correct: in recursion function (1), the two sequences of tokens represent the same text with different tokenizations. Furthermore, due to the dynamic programming definition, the direction can only come from $(k − 1, j − 1)$, $(k, j − 1)$, or $(k − 1, j)$, meaning that token matching can only occur between adjacent positions in order. Therefore, in Figure 2b, there should not be X‑shaped matching (i.e., (1,2)* and (2,1)*), but only I‑shaped (i.e., (1,1) and (2,2)), V‑shaped (i.e., (1,1) and (2,1)), and inverse V‑shaped (i.e., (1,1) and (1,2)) matchings.\n\n*We use (x, y) to represent the token pairing between x: source model token position and y: target model token position.\n\nWe apologize again for the confusion caused by the figure and sincerely appreciate your careful reading and for pointing this out. We will revise it accordingly to eliminate any ambiguity.\n\n\n**Q4: Regarding Computational Efficiency of Token Alignment and Optimal Transport.**\n\n\nThis is a great question. We fully agree with the reviewer that fully disclosing the computational burden could greatly benefit the community. As detailed in Appendix S2, FuseLLM takes approximately 3 GPU hours (on an NVIDIA A100-80GB) to generate source and target model predictions, and 4 CPU hours (on an AMD EPYC 7763 processor) to conduct token alignment. Compared to FuseLLM, PTA-LLM introduces an additional optimal transport computation, which results in approximately a 13.75% increase in CPU alignment time on the MiniPile dataset. It would be valuable to explore more efficient alternatives, such as the Greenkhorn algorithm, to reduce the complexity of the optimal transport step.\n\n**Q5: Regarding the Task Selection in Stability Analysis**\n\nThank you for your insightful comment. Following your suggestion, we provide additional comparison results for cases where FuseLLM performs well while PTA‑LLM performs poorly.  \n\n| Task                     | LLaMA‑2 | FuseLLM | PTA‑LLM |\n| ------------------------ | ------- | ------- | ------- |\n| Multistep Arithmetic Two | 0.80    | 4.80    | 4.00    |\n| Temporal Sequences       | 12.80   | 16.40   | 15.20   |\n\nAs shown in the table above, PTA‑LLM exhibits slight performance degradation in two BBH subtasks, both of which evaluate EM accuracy on various reasoning scenarios. PTA‑LLM applies OT to balance the probability mass in each row, producing a soft assignment matrix that yields more compact and consistent token representations. This smoothing flattens peaks for high‑confidence tokens. In exact‑match BBH tasks, spreading probability to distractor tokens can shift the argmax and reduce accuracy. PTA‑LLM is penalized for minor surface deviations, while FuseLLM’s hard alignment preserves high‑confidence tokens, yielding more frequent matches with the official answers. This explains why PTA‑LLM achieves the smallest accuracy improvement under BBH in Table 2.  \n\nWe will supplement more comprehensive per‑task results between FuseLLM and PTA-LLM in Appendix S8 to ensure the methods are evaluated under unbiased conditions.\n\n**Q6: Regarding the Figure 3**\n\nWe are happy to restate and elaborate on our motivation for presenting Figure 3 and Section 4.4, as they strongly support the core motivation of PTA‑LLM.  \n\nThe motivation of PTA‑LLM is to achieve a coherent fused metric, as described in our introduction. We aim to keep the logit distribution more compact and consistent with the target distribution, which motivates the design of probabilistic token alignment under the optimal transport framework.  \n\nIn addition to the effectiveness and stability results shown in Sections 4.2 and 4.3, we incorporate empirical visualizations to support our motivation. As shown in Figure 3, we visually compare the distribution of PTA‑LLM fused tokens with the target tokens and the FuseLLM fused tokens. Specifically, Figures 3(b) and 3(d) show that PTA‑LLM achieves a more consistent marginal feature distribution with the target tokens, whereas FuseLLM exhibits significantly greater distortion in the overall token representation. The more compact and coherent token distribution after applying probabilistic token alignment is consistent with the motivation described above.\n\nWe will include this elaboration and add a clarifying note in the revision.\n\n**Q7: Regarding the Equation 4**\n\nThank you for your great question. We would like to provide further clarification here.\n\nWithout constraints on m and n, the transport space between the logits of Model A and Model B would equal the product of their vocabulary sizes (for example, MPT has 50,277 tokens while LLaMA has 32,000). This is computationally **inefficient** because it requires storing over 30k logits per token instead of just the top 10, and it is **unnecessary** because the logit distribution follows a long‑tail pattern in which the top‑k logits capture the majority of the probability mass. Our further analysis of Token Alignment Window Sizes, ranging from 10 to 3 (Tables 4 and S10), shows that a small window size is sufficient for effective model knowledge fusion.\n\nWe will include this elaboration and add a clarifying note in the revision.\n\n\nWe appreciate your thoughtful comments. We hope our response addresses your concerns. Please let us know if there are any additional questions, and we will be happy to discuss further."}]

Author rebuttal:

Meta-review: This paper addresses the problem of combining multiple large language models into a stronger model in a setting in which the token vocabulary differs across models. All the reviewers are positive or mildly positive about the paper, and they were satisfied with the rebuttal by the authors.
Decision: Accept (poster)