Probabilistic Token Alignment for Large Language Model Fusion

Dear reviewer zYWa,

We 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.

Q1: Regarding Originality Claim of Dynamic Token Pairing.

Sorry 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.

[ref1] Specializing smaller language models towards multi-step reasoning

Q2: Regarding Correctness of Recursion Function Definition in Equation (1).

Thank 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:

$f(k,j) = \min\lbrace f(k-1,j),\; f(k,j-1),\; f(k-1,j-1) \rbrace + C(B_k,A_j)$

Our 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.

We 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.

Q3: Regarding Token Order Consistency in Recursion Function and Figure 2b(1).

Thank 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.

*We use (x, y) to represent the token pairing between x: source model token position and y: target model token position.

We 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.

Q4: Regarding Computational Efficiency of Token Alignment and Optimal Transport.

This 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.

Q5: Regarding the Task Selection in Stability Analysis

Thank you for your insightful comment. Following your suggestion, we provide additional comparison results for cases where FuseLLM performs well while PTA‑LLM performs poorly.

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

We will supplement more comprehensive per‑task results between FuseLLM and PTA-LLM in Appendix S8 to ensure the methods are evaluated under unbiased conditions.

Q6: Regarding the Figure 3

We 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.

The 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.

In 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.

We will include this elaboration and add a clarifying note in the revision.

Q7: Regarding the Equation 4

Thank you for your great question. We would like to provide further clarification here.

Without 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.

We will include this elaboration and add a clarifying note in the revision.

We 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.

Dear reviewer sCqG,

We 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.

Q1: Regarding the improvements.

Thank 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.

Our 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.

Q2: Regarding Computational Cost of Token Alignment.

This 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.

Setting 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.

Setting 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).

The 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.

Q3: Regarding Lack of Comparison with Other Model Combination Techniques Thank 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.

However, 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.

From 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.

[ref1] LLM-Blender: Ensembling Large Language Models with Pairwise Ranking and Generative Fusion

Q4: Regarding Model Fusion Settings.

Thank 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.

Q5: Regarding Scalability and Data Subsampling in Knowledge Fusion.

This 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.

Q6: Regarding Use of Edit Distance vs. Semantic Similarity in Transport Cost.

Thank 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.

We 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.

We 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.

Dear reviewer 8QeG,

We 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.

Q1: Regarding the Use of CLM Loss in Model Fusion.

Thanks for the great question, and we'd like to provide further explanation.

When 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.

Q2: Regarding the True Challenge in Token Pairing.

Thank 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.

We will refine the relevant writing in the revised version to ensure better coherence and enhance overall clarity.

Q3: Regarding the writing.

Ambiguity in Source and Target Model Definitions:

Sorry 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.

Transport objective:

Thank 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.

We will move this discussion from appendix to the main paper for better clarification.

Notational Inconsistency in Target Model Predictions:

Thank 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$.

Q4: Regarding Interpretation and Clarity of Equation (5) Assignment

Thank you for your excellent question. We would like to elaborate on the motivation behind Equation 5 and its empirical visualization.

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.

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.

We will include this elaboration and add a clarifying note in the revision.

We 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.

Dear reviewer seY2,

We 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.

Q1: Regarding Model Fusion Settings.

Thank 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.

Q2: Regarding the Improvements.

Thank 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.

Our 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.

We 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.