Probabilistic Token Alignment for Large Language Model Fusion

Dear reviewer qifk,

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 potential application to other source models.

Thank you for the excellent suggestion. We also note related work [ref1] that explores the fusion of Mixtral, InternLM2, and OpenChat, reporting consistent performance improvements under the knowledge fusion paradigm. We thus conduct a preliminary experiment on a MiniPile subset to extend our framework to fuse additional LLMs. Specifically, in this additional experiment, our approach is trained on a subset of MiniPile (100K training examples in total) created by randomly sampling 10% of the original dataset.

The results presented in the table below demonstrate that our approach exhibits robustness and effectively integrates knowledge from diverse source models. Due to the constrained computational resources and time limit, these experiments are ongoing. We will provide updated results once they have been completed. We are excited to explore the training on full-size MiniPile further.

[ref1] FuseChat: Knowledge Fusion of Chat Models

Q2.1: Regarding the improvement over FuseLLM.

Thank you for the great observation. While our average absolute gain is 0.49% compared to the FuseLLM, the relative gain of 1.72% is nontrivial. Moreover, our significance test results show that the improvements are statistically significant with p-value < 0.005. We also reported the standard deviation over three runs in the table below.

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.

Q2.2: Regarding the evaluation of GSM.

Thank you for the insightful suggestion. Following your suggestions, we reevaluated our method using an 8-shot setting under a widely adopted evaluation framework, lm-evaluation-harness. The gsm8k-cot evaluation results for Llama-2 7B (i.e., 14.18%) align closely with those reported in the original paper [ref1] (i.e., 14.60%).

From the results, we observe a consistent performance improvement when utilizing probabilistic token alignment. We have revised the reported results to ensure a fairer and more general comparison.

[ref1] Llama 2: Open Foundation and Fine-Tuned Chat Models

(Due to character limitations, we will continue to answer in another comment.)

(Continue from where rebuttal ends.)

Q3: Regarding the study of stability.

Thank you for the great question. In addition to the performance comparisons detailed in Table 2, we provide supplementary case studies in Appendix D to offer deeper insights into the effectiveness of our proposed method. For instance, consider the Tracking Shuffled Objects task with seven objects, 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 PTA-LLM's probabilistic token alignment mechanism, which effectively transforms logits into the correct objective rather than merely replicating the original logits from FuseLLM. We have provided an expanded discussion of this analysis in the revised paper.

Q4: Regarding the ablation experiments.

Following your suggestion, we provide additional settings for each hyperparameter to reveal the impact trends of each factor. As observed from the table, within a reasonable range, our method is robust. We have supplemented these experiments with further discussion in the revision. Thank you again for your suggestion.

Q5: Regarding the typos.

Thank you for pointing them out. We have revised accordingly and reformulated all the equations to ensure clarity and precision.

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.

We sincerely appreciate the reviewer's engagement in the discussion and their valuable comments, which are crucial in improving the quality of our work. We’d like to take the opportunity to provide further clarifications regarding Q1 and Q4.

Regarding Q1

First, we completely agree that our experiment would benefit from incorporating models with greater architectural diversity. Our initial choice to include three models was based on the FuseLLM setting to ensure a fair comparison.

Second, please note that the results reported above were obtained using only 10% of the data to quickly generate preliminary findings. The primary aim was to demonstrate that fusing additional models can indeed improve performance compared to PTA-LLM. For instance, PTA-LLM + Qwen2.5 improves MMLU performance from 47.42 to 51.28, achieving an absolute ~4% improvement with just 10% training data. We have since completed training on 100% of the data and present fair comparisons with Mistral and Qwen2.5 below. The results show that both PTA-LLM + Qwen2.5 and PTA-LLM + Mistral outperform Qwen2.5 and Mistral, respectively, demonstrating the effectiveness of our approach.

Third, for generative tasks, we have already included results on MultiPL-E in the original paper. To provide further comparisons with the Mistral and Qwen models, we conducted additional experiments on HumanEval using MiniPile+Github for training. The results below, combined with our original findings on MultiPL-E, demonstrate that our method is effective for generative tasks as well.

Regarding Q4

Thank you for these excellent questions. We’d like to clarify that the additional experiments are consistent with the findings presented in Section 4.5. Specifically:

1. transport convergence threshold: 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.

2. token alignment window size: The findings of larger token alignment window sizes lead to better performance is consistent with previous results. A larger transport range allows for a more comprehensive understanding of the context, which in turn improves performance.

3. combination weight: Sorry for the confusion on this part. The paper mentions "higher performance when the weight is smaller" is referring to a comparison between 0.8 and 0.9. However, if the weight is reduced further (e.g., to 0.6), 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. Thank you for highlighting this point; we have revised the paper accordingly.

We greatly value each comment and suggestion from your review, and are hoping that our additional clarifications and experimental results address your concerns. We are eager to address any remaining issues during the discussion phase. Thank you very much.

(Continue from where rebuttal ends.)

Q3: Why does knowledge fusion outperform the CLM objective?

Thank you for the insightful question. Two potential factors may explain why the knowledge fusion objective outperforms the traditional CLM 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 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. We have expanded on these points further in the revised version of the manuscript, offering a more comprehensive discussion.

Q4: Regarding the combination weight.

Thank you for your question. In addition to the ablation study presented in Section 4.6, we conducted supplementary experiments to further analyze the impact of the combination weight hyperparameter.

We observe consistent performance improvements when λ is set within the range of 0.8 to 0.9. Furthermore, we conclude that different benchmarks exhibit slight variations in hyperparameter preferences (i.e., BBH and MMLU tend to favor λ=0.8, while ME shows a preference for λ=0.85). Additionally, our method has a preference for lower λ values in comparison to FuseLLM (i.e., 0.9), suggesting a stronger emphasis on probabilistic token alignment in our method.

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 HyYP,

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

Thank you for your insightful questions.

Regarding the performance and statistically significant. 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. We further rerun the same hyperparameter settings three times and reported standard deviation error bars.

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.

Regarding using the additional costs to train more data.

This is an excellent question. To further demonstrate the effectiveness of our method, we conducted an additional experiment to compare two distinct settings following your suggestion.

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.

Q2: The explanation of the EM score of GSM.

Sorry for the confusion. The low score reported is attributed to our initial evaluation being conducted under zero-shot settings without employing chain-of-thought prompting, which does not comprehensively reflect the true capabilities of the evaluated models.

To address this issue, we reevaluated our method using an 8-shot setting under a widely adopted evaluation framework, lm-evaluation-harness. The gsm8k-cot evaluation results for Llama-2 7B (i.e., 14.18%) closely align with the result reported in the original paper [ref1] (i.e., 14.6%). From the table below, we observe consistent performance improvements when utilizing probabilistic token alignment.

We have revised the reported results to ensure a fairer and more general comparison. Thank you for your question.

[ref1] Llama 2: Open Foundation and Fine-Tuned Chat Models

(Due to character limitations, we will continue to answer in another comment.)

Thank you for the additional question. This is an excellent observation, and we’d like to provide some explanation here.

First, in the LLaMA2 report [ref1], Table 25 (page 51) presents separate scores for GSM8k (14.6) and MATH (2.5) of LLaMa2-7B. The GSM8k score of 14.6 also aligns with the results listed on their official website (https://www.llama.com/llama2/). To confirm this further, we re-ran LLaMA2-7B on GSM8k using the officially released model [ref2] and obtained a consistent result of 14.18.

Based on this, we suspect there may be a small inconsistent issue with the results reported in Table 3 of the LLaMA2 report. Specifically, the Math results in Table 3 indeed represent the average of the GSM8k and MATH scores from their Table 25, except for the LLaMA2-7B and LLaMA2-14B models.

We greatly appreciate your feedback and hope these clarifications address your concerns. Please let us know if there are any other issues or questions you'd like us to address during the discussion phase. Thank you again for your valuable input!

[ref1] Llama 2: Open Foundation and Fine-Tuned Chat Models (https://arxiv.org/pdf/2307.09288)

[ref2] https://huggingface.co/meta-llama/Llama-2-7b-hf

Dear reviewer biko,

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

Thank you for the insightful suggestion. We have incorporated this limitation into the main manuscript accordingly.

Q2: Regarding the end-to-end training time.

Since the end-to-end training time primarily depends on the dataset size, we conducted an additional study using subsets of MiniPile (which contains 1M training samples). Specifically, we randomly sampled the original MiniPile to create subsets for training.

The results demonstrate a linear relationship between dataset size and training time. Experiments were conducted on 8 NVIDIA A100-80GB GPUs.

To provide further information, we have conducted training on various size datasets and reported their estimated time computed by the framework on 8 NVIDIA A100-80GB GPUs. The same linear relationship can also be observed from the table.

[ref1] https://huggingface.co/datasets/vietgpt/openwebtext_en

[ref2] https://huggingface.co/datasets/wikimedia/wikipedia

Q3: The explanation of the EM score of GSM.

Thank you for the insightful observation. We’d like to provide some clarification here. The EM (Exact Match) score for GSM was intended to measure the accuracy of predicted mathematical results, requiring an exact match with the ground truth answer. The low score reported is attributed to our initial evaluation being conducted under the zero-shot setting without employing chain-of-thought prompting.

To further address your concern, we have reevaluated our method using an 8-shot setting under a widely adopted evaluation framework, lm-evaluation-harness. As shown in the results below, the gsm8k-cot evaluation results for Llama-2 7B (i.e., 14.18%) closely align with the result reported in the original paper [ref1] (i.e., 14.6%). From the table below, we observe consistent performance improvements when utilizing probabilistic token alignment. We have revised the reported results in Sec. 4.2 for completeness.

[ref1] Llama 2: Open Foundation and Fine-Tuned Chat Models

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 vEFt,

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

Thank you for your suggestion. We have revised our writing to adopt a more accurate and modest tone.

Q2: Regarding use of terms like ‘token distribution’ and ‘probability mass’ are a bit misleading.

Sorry for the confusion. We completely agree with the reviewer that 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. In this work, 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 the transport cost. Therefore, the term "transport probability" more accurately describes our method. We have revised accordingly to make it more clear and sincerely appreciate your thoughtful comments.

Q3: The explanation of "surface-level token correspondences”.

Yes, the edit distance is computed at the token level in our method as well. However, besides using edit distance as one metric, our method advances the distance measurement by incorporating the corresponding logit values into the individual cost within the transport framework. Optimization is thus performed at both the "surface-level" and "logit-level".

Specifically, our cost function incorporates both the edit distance cost and the corresponding logit values (i.e., $c_{xy}p_{xy}$), with the overall transport cost being the summation over all alignments: $\sum \sum c_{xy}p_{xy}$. We have revised the content to make it more clear.

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.