Go to ICLR 2025 Conference homepageProbabilistic Token Alignment for Large Language Model Fusion
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.