Probabilistic Token Alignment for Large Language Model Fusion

Authors: Runjia Zeng, James Chenhao Liang, Cheng Han, Zhiwen Cao, Jiahao Liu, Xiaojun Quan, Qifan Wang, Tong Geng, Dongfang Liu
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models, Model Fusion

Abstract:
With the drive of scaling laws, the size of large language models (LLMs) continues to grow, causing a burden when constructing stronger baselines from scratch. In response, we propose a more robust baseline by employing knowledge fusion under the model fusion paradigm. However, a key challenge in existing knowledge fusion approaches is their dependence on manually predefined alignment strategies, which may not generalize well across diverse contexts, leading to performance degradation in several zero-shot evaluation tasks. To address this challenge, we draw inspiration from distribution learning and propose the *probabilistic token alignment* method as a general and soft mapping solution for alignment, resulting in **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 task. Our approach is validated across diverse benchmarks and tasks using three prominent LLMs with distinct architectures—Llama-2, MPT, and OpenLLaMA. Empirical results demonstrate that probabilistic token alignment enhances the target model's performance across multiple capabilities.

Decision: reject
Reviewer scores: 8, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730184634439

thread_type:
official_review

<thread_object>
creator: reviewer_qifk

<content>
summary:
This paper introduces a novel approach PTA-LLM to fusing LLMs with different architectures by framing token alignment as an optimal transport problem. This distribution-aware method generalizes token alignment effectively, enhancing performance in zero-shot evaluation tasks and offering interpretability through a distributional lens.

strengths:
This work addresses token alignment, a key challenge of knowledge fusion, by applying the classic mathematical optimization method of optimal transport to achieve a soft alignment. This approach also offers an interpretable analysis for token alignment.

weaknesses:
1.This article focuses on token alignment for models with different structures but only tests three models, with OpenLLaMA and Llama2 having relatively high similarity. The experiment would benefit from including models with greater architectural diversity, such as Mistral and Qwen.

2.In Table 1 of the main experimental results, PTA-LLM shows limited improvement over FUSELLM, with an average absolute gain of less than 0.5 points. Additionally, the experiment employs a zero-shot evaluation approach, which is challenging for the base models and may not fully reflect their capabilities; for example, the GSM dataset typically uses an 8-shot or 4-shot setup. In the zero-shot setting, the scores of the three base models remain very low.

3.In Section 4.3 on STABILITY, the authors selected certain tasks to compare the performance of PTA-LLM and FUSELLM, demonstrating that soft alignment offers greater flexibility and adaptability than hard mapping. However, the article does not analyze the characteristics and shared features of these advantageous tasks or assess the applicability of soft alignment in these cases, making it difficult for the explanation to convincingly support its hypothesis.

4.In the ablation experiments of Section 4.5, the paper examines only two cases each for Optimal Transport Temperature, Combination Weight, and Token Alignment Window Size, which does not adequately reveal the impact trends of each factor.

5.There are some typos in the paper：

Line 223: "Pt" -> "P_t"

Line 274: "Tt" -> "T_t"

Line 417:  "-17/07%" -> "-17.07%"

Line 430: "26.96" -> "36.96"

questions:
See "Weaknesses".
</content>
</thread_object>

</thread>
<thread>
creation date:
1730381026442

thread_type:
official_review

<thread_object>
creator: reviewer_hyyp

<content>
summary:
This paper introduces PTA-LLM, a probabilistic token alignment method designed to align the distributions of various LLMs with different vocabularies. Inspired by the optimal transport problem, PTA-LLM employs distribution-aware learning to enhance more coherent model fusion. Experimental results across multiple benchmarks confirm the effectiveness of PTA-LLM.

strengths:
- The paper is easy to read and presents its ideas clearly.

- The proposed PTA-LLM is general, stable, and interpretable. 

- The paper conducts a series of experiments to demonstrate the effectiveness of PAT.

weaknesses:
- The improvement in effectiveness seems insufficiently significant. In Table 1, the average gap between CLM is only about 1 point. Given the complexity of Model Fusion compared to ordinary CLM, this improvement may not adequately highlight the advantages of PTA. For example, PAT requires obtaining the probability distributions of all teacher models, which means incurring multiple forward pass costs. What would be the effect if these additional costs were used to train more data with CLM? In addition, Multiple experiments are needed for significance testing to demonstrate the effectiveness of PAT.

questions:
- Why does llamda2 perform only 0.66 on GSM, while I saw that the original llamda2 paper did not report such a low score?

- Why is direct training for next token prediction not ideal? In theory, all models fit the training corpus through next-token prediction. Why does distilling the distribution of different models on the corpus outperform the next-token-prediction training objective itself?

- Knowledge fusion also keeps the CLM loss for training, using a parameter called lambda to control weight. How does the value of lambda affect the performance? Whether lambda need to be fine-tuned elaborately for improvement？
</content>
</thread_object>

</thread>
<thread>
creation date:
1730611126396

thread_type:
official_review

<thread_object>
creator: reviewer_biko

<content>
summary:
This paper proposes PTA-LLM, a knowledge fusion-based approach for language model fusion. Compared with FuseLLM, a pioneer knowledge fusion work, PTA-LLM changes the hard mapping of source and target vocabulary into soft mapping and uses an optimal transport method to calculate the optimal token alignment. Overall, the paper offers a sound and well-motivated solution to the core problem, token alignment, in knowledge fusion over heterogeneous source LLMs.

strengths:
- Token alignment is the core issue of heterogeneous knowledge fusion. The paper offers a sound and well-motivated solution to this problem.
- The overall experimental results look good, though the improvements are not significant.

weaknesses:
- In practice, the Sinkhorn algorithm seems to be time-consuming. The limitations section should be better included in the main pages for a fair demonstration. I also have some concerns about the end-to-end training time on corpus other than MiniPile.
- What does the EM score for GSM stand for? Is it 1.90%? The score seems extremely low and does not align with other reported results.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691559704

thread_type:
official_review

<thread_object>
creator: reviewer_veft

<content>
summary:
This paper proposes a new method for aligning tokens when fusing models with different vocabularies. The method uses optimal transport to transfer the information from a source model’s logits to the target model’s logits for a given text. Dynamic token pairing allows a single token from the shorter token sequence to be mapped to multiple tokens in the other’s sequence. The transport plan minimizes the cost of transferring mass between the two logit spaces according using edit distance of tokens to determine transport cost. The most probable alignment is selected, and the token logits of the target model are adjusted toward their aligned source token logits. 

Positive evidence for the value of this method is provided through experiments on a variety of standard LLM evaluation datasets compared to reasonable baselines and related work. Analysis shows that the proposed method provides more stability than other fusion methods. A qualitative analysis of the proposed method shows a more compact and consistent representation compared to a competitive fusion method. Ablations provide information about the robustness and effect of hyperparameters.

strengths:
This paper provides an interesting solution to the problem of mixing information from language models with different vocabularies. The performance gains from the proposed method are consistent versus the competing methods. Optimal transport is very fashionable right now, and the ICLR community may appreciate seeing it applied here.

weaknesses:
The paper’s writing is at times grandiose. The introduction begins with a question about how to build better models without resorting to scaling, and then insists we reframe this question as one about model fusion. Some hedging here would help i.e. “one method for addressing question 1 is knowledge fusion”. 

Sometimes, the paper is unclear. For instance: while optimal transport is general and can be applied to moving mass between a variety of spaces, it is most often applied in ML to moving probability mass. The current paper uses the standard ML terminology of “transferring probability” between distributions (e.g. line 272), but seems to be moving mass between logit spaces. Since logits are not probability distributions, the use of terms like “token distribution” and “probability mass” are a bit misleading. I’m not certain that moving mass in logit-space is the same as moving mass between probability spaces since the normalization terms in the source and target space are different. More careful writing could address this problem.

questions:
The paper argues the proposed method improves on earlier methods which rely on “surface level token correspondences” (line 153), but doesn’t the use of edit distance as a cost function for the transport plan also rely on surface level correspondences? Is this edit distance computed at the token level?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732189555454

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revision

comment:
To All Reviewers:

Thank you for your thorough review and insightful comments. We have revised our paper according to the suggestions. The major changes are summarized as follows:

- As suggested by Reviewer vEFt, we have revised our writing in Sections 1 and 2 to adopt a more accurate and modest tone. Additionally, we added more discussion about the terminology of "transport probability" in Appendix A.
- Following the comments from Reviewers biko, HyYP, and qifk, we revised the reported results in Section 4.2 to ensure a fairer and more general comparison. Moreover, we supplemented experiments with more hyperparameter settings in Appendix E and expanded the discussion in Appendix F to further clarify why knowledge fusion outperforms the CLM objective.
- As suggested by Reviewer biko, we have incorporated the limitation into Section 5 to ensure a fair demonstration.
- Based on the suggestions by Reviewer qifk, we have corrected typos in Section 3.2 and provided an expanded discussion of case studies in Appendix D. Furthermore, we reported the p-value and standard deviation in Appendix C to demonstrate the statistical significance of our results.
- We updated Figures 1-3 to enhance their interpretability and alignment with the revised text. Additionally, all equations (1-5) were reformulated for improved clarity and precision.

All modifications have been marked in ${\color{blue} blue}$ in our revised submission.

Sincerely yours,\
Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732330969930

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Looking forward to the discussion

comment:
Dear Reviewers,

We sincerely appreciate the time and effort you've devoted to reviewing our work. We understand that your schedule may be quite busy, and we are truly grateful for your valuable feedback. As the Author-Reviewer discussion phase is ending soon, we would greatly value the opportunity to engage in further discussion with you. Our aim is to gain insights into whether our responses effectively address your concerns and to ascertain if there are any additional questions or points you would like to discuss.

We look forward to the opportunity for further discussion with you. Thank you for your thoughtful consideration.

Best regards,\
Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733205145089

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Author-Reviewer Discussion

comment:
Dear Area Chair and Reviewers,

We would like to express our sincere gratitude for your efforts in facilitating the discussion regarding our paper. As the discussion is coming to an end, we would like to provide a brief summary of the key points that have been discussed:

- We have clarified the terms of "token distribution" and "probability mass" in the revised paper to enhance their clarity, as suggested by **Reviewer vEFt**.

- In response to **Reviewer biko**'s suggestions, we have included an analysis of training time with varying data sizes. Additionally, we have elaborated on the evaluation of the EM score on GSM and provided supplementary experimental results under 8-shot settings.

- As recommended by **Reviewer HyYP**, we have included significance test results and conducted further experiments on cost-effective training under identical resource constraints. These results further validate the superior performance of our approach.

- Following **Reviewer qifk**'s suggestions, we have conducted experiments integrating additional heterogeneous LLMs, showcasing the robustness of our method in fusing diverse source models. We have also added supplementary case studies to illustrate the stability and effectiveness of our approach, along with detailed ablation studies analyzing the impact of each hyperparameter.

In summary, we are grateful to all reviewers for their constructive feedback and acknowledgment of our responses. We particularly appreciate that **Reviewers HyYP and qifk** have increased their scores and that **Reviewers vEFt and biko** have maintained their positive evaluations.

Finally, we deeply value the reviewers' insightful comments, which have helped us refine our work. We believe that our revisions significantly enhance the clarity and impact of our contributions. We hope our work provides meaningful insights and advances the development of the knowledge fusion domain.

Sincerely,\
Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper tackles a technical challenge in knowledge fusion for LLMs: how do we align the outputs of different LLMs during the fusion step? When tokenization differs, the sequences produced by two LLMs may be different in length and use different vocabulary sizes / different tokens at each step.  The proposed approach uses a dynamic programming solution to find a potentially many-to-many monotonic alignment between the sequences. Once two tokens are aligned, an optimal transport problem is defined to yield a transport plan to align the distributions themselves.

The paper tackles an interesting (if somewhat niche) question and presents a novel and effective algorithm. The paper is generally well-written. The results show consistent improvement over FuseLLM.

One main weakness of this paper is the narrowness of the setting: this kind of fusion across LLMs with different tokenizers is a fairly particular setting.  Another major question is whether the accuracy improvements in Table 1 merit the complexity of the approach (HyYP, qifk). I am not convinced by relative gains. On most of the given tasks, even if there is an improvement over FuseLLM, the gain is not large enough to be practically significant: it may not transfer to noticeable performance improvements for LLMs in practice.

Runtime information is buried in the appendix and is left in theoretical terms except for a vague reference to 13.75% aligning delay. As biko points out, this should be included in the main body of the paper with more discussion and analysis, particularly as the accuracy improvements are small.