Table of Contents
Fetching ...

Token-Level LLM Collaboration via FusionRoute

Nuoya Xiong, Yuhang Zhou, Hanqing Zeng, Zhaorun Chen, Furong Huang, Shuchao Bi, Lizhu Zhang, Zhuokai Zhao

TL;DR

FusionRoute presents a token-level collaboration framework that jointly selects domain-specific experts and supplies complementary logits at each decoding step, addressing the limitations of purely expert routing. The router is trained in stages—SFT to learn token-level routing, followed by CDPO to learn corrective logits—with mix training to stabilize learning. Theoretical analysis via a token-level MD P and the Performance Difference Lemma shows purely token-based routing cannot guarantee near-optimal decoding without strong coverage, while the complementary generator expands the policy class to approach optimality. Empirically, FusionRoute improves cross-domain performance on math, coding, and instruction tasks across Llama-3 and Gemma-2 and scales favorably with model size, outperforming sequence- and token-level baselines, model merging, and fine-tuning. These results demonstrate FusionRoute's robustness, efficiency, and practicality for deploying specialized LLMs in a general-purpose setting.

Abstract

Large language models (LLMs) exhibit strengths across diverse domains. However, achieving strong performance across these domains with a single general-purpose model typically requires scaling to sizes that are prohibitively expensive to train and deploy. On the other hand, while smaller domain-specialized models are much more efficient, they struggle to generalize beyond their training distributions. To address this dilemma, we propose FusionRoute, a robust and effective token-level multi-LLM collaboration framework in which a lightweight router simultaneously (i) selects the most suitable expert at each decoding step and (ii) contributes a complementary logit that refines or corrects the selected expert's next-token distribution via logit addition. Unlike existing token-level collaboration methods that rely solely on fixed expert outputs, we provide a theoretical analysis showing that pure expert-only routing is fundamentally limited: unless strong global coverage assumptions hold, it cannot in general realize the optimal decoding policy. By augmenting expert selection with a trainable complementary generator, FusionRoute expands the effective policy class and enables recovery of optimal value functions under mild conditions. Empirically, across both Llama-3 and Gemma-2 families and diverse benchmarks spanning mathematical reasoning, code generation, and instruction following, FusionRoute outperforms both sequence- and token-level collaboration, model merging, and direct fine-tuning, while remaining competitive with domain experts on their respective tasks.

Token-Level LLM Collaboration via FusionRoute

TL;DR

FusionRoute presents a token-level collaboration framework that jointly selects domain-specific experts and supplies complementary logits at each decoding step, addressing the limitations of purely expert routing. The router is trained in stages—SFT to learn token-level routing, followed by CDPO to learn corrective logits—with mix training to stabilize learning. Theoretical analysis via a token-level MD P and the Performance Difference Lemma shows purely token-based routing cannot guarantee near-optimal decoding without strong coverage, while the complementary generator expands the policy class to approach optimality. Empirically, FusionRoute improves cross-domain performance on math, coding, and instruction tasks across Llama-3 and Gemma-2 and scales favorably with model size, outperforming sequence- and token-level baselines, model merging, and fine-tuning. These results demonstrate FusionRoute's robustness, efficiency, and practicality for deploying specialized LLMs in a general-purpose setting.

Abstract

Large language models (LLMs) exhibit strengths across diverse domains. However, achieving strong performance across these domains with a single general-purpose model typically requires scaling to sizes that are prohibitively expensive to train and deploy. On the other hand, while smaller domain-specialized models are much more efficient, they struggle to generalize beyond their training distributions. To address this dilemma, we propose FusionRoute, a robust and effective token-level multi-LLM collaboration framework in which a lightweight router simultaneously (i) selects the most suitable expert at each decoding step and (ii) contributes a complementary logit that refines or corrects the selected expert's next-token distribution via logit addition. Unlike existing token-level collaboration methods that rely solely on fixed expert outputs, we provide a theoretical analysis showing that pure expert-only routing is fundamentally limited: unless strong global coverage assumptions hold, it cannot in general realize the optimal decoding policy. By augmenting expert selection with a trainable complementary generator, FusionRoute expands the effective policy class and enables recovery of optimal value functions under mild conditions. Empirically, across both Llama-3 and Gemma-2 families and diverse benchmarks spanning mathematical reasoning, code generation, and instruction following, FusionRoute outperforms both sequence- and token-level collaboration, model merging, and direct fine-tuning, while remaining competitive with domain experts on their respective tasks.
Paper Structure (46 sections, 2 theorems, 53 equations, 4 figures, 3 tables, 1 algorithm)

This paper contains 46 sections, 2 theorems, 53 equations, 4 figures, 3 tables, 1 algorithm.

Key Result

Lemma 4.2

Denote $\rho \in \Delta(\mathcal{X})$ as a distribution over the prompt space. Suppose each response $y$ has length $T$. Then, for any two policies $\pi, \pi^*$, we have

Figures (4)

  • Figure 1: Sequence-level collaboration is coarse and inefficient, while prior token-level methods are unstable. FusionRoute achieves fine-grained, efficient, and robust token-level collaboration through complementary routing.
  • Figure 2: Top: Overall architectural design of FusionRoute. FusionRoute enables multiple specialized LLMs to work collaboratively through a carefully designed router. The router outputs both the weights for token-level routing and the logits to complement the experts' output. Bottom: The training process is decomposed into two phases. In the SFT phase (§\ref{['subsubsec:sft']}), the router learns token-wise mixture weights over the expert models and is jointly fine-tuned to become a good response generator. In the CDPO phase (§\ref{['subsubsec:cdpo']}), FusionRoute refines the final policy by using preference optimization to improve upon the experts’ outputs.
  • Figure 3: GPT-4o winrate on general datasets compared to fine-tuned model.
  • Figure 4: GPT-4o winrate comparison: FusionRoute with vs. without CDPO training on both Llama-3 and Gemma-2 Family.

Theorems & Definitions (3)

  • Lemma 4.2: Performance Difference Lemma kakade2002approximately
  • Theorem 4.3
  • proof