Probabilistic Token Alignment for Large Language Model Fusion
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.
Primary Area: foundation or frontier models, including LLMs
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 3346
Loading