PaperlayerView exact revision on arXiv

Source-linked AI summary

Token-Level LLM Collaboration via FusionRoute

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

arXiv:2601.05106v5cs.AIcs.CLcs.LG

TL;DR

General-purpose LLMs can be costly, while specialized models may not generalize beyond their training distributions, and token-level collaboration can be unreliable when experts or selection fail. FUSIONROUTE combines token-level expert selection with complementary logits and trains the router through supervised and preference optimization. The paper reports broad empirical gains and theoretically characterizes limits of purely expert-based routing.

  • Problem

    General-purpose LLMs can be prohibitively costly, specialized models may struggle beyond their training distributions, and existing token-level collaboration can be unreliable when experts or selection fail.

    Source: §1 · pp. 1–2

  • Method

    FUSIONROUTE uses a lightweight router to select an expert at each decoding step and add complementary logits, with preference optimization training the corrective signal.

    Source: §1 · p. 2, §3.2.2 · p. 3

  • Results

    FUSIONROUTE outperforms token-level collaboration, model merging, and fine-tuned models across specialized and general-purpose evaluations, while the theory identifies limits of purely expert-based routing.

    Source: §1 · p. 2, §4.3 · p. 6, §7 · p. 9

  • Takeaways & Limitations

    The reported experiments support FUSIONROUTE’s effectiveness across domain-specific and general datasets, while its complementary generator addresses a theoretical limitation of pure expert routing.

    Source: §7 · p. 9

    Even when a near-optimal trajectory exists and value gaps are small on good prefixes, pure token-level routing cannot guarantee a near-optimal policy under the stated setting.

    Source: §4.3 · p. 6

Abstract

from arXiv · show

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.

1. Introduction

Large general-purpose LLMs are costly, while efficient specialized models can underperform outside their training distributions. FUSIONROUTE addresses the coordination challenge with token-level expert selection and complementary generation.

1. Introduction

  • Large general-purpose LLMs can perform across domains but may be too costly, while efficient specialized models can struggle beyond their training distributions.The paper identifies robust coordination of specialized models as a central challenge.Source: §1 · pp. 1–2
  • Token-level collaboration remains vulnerable when experts perform poorly or the selection strategy is incorrect for particular tasks or tokens.This motivates a more robust and automatic collaboration paradigm.Source: §1 · p. 2
  • FUSIONROUTE combines token-level expert selection with a complementary logit that can refine generation.Its lightweight router selects an expert at each decoding step and contributes its own generation signal.Source: §1 · p. 2
  • The paper reports that FUSIONROUTE outperforms token-level collaboration, model merging, and fine-tuned models across specialized and general-purpose evaluations.The reported domains include mathematics, coding, and instruction following.Source: §1 · p. 2

2. Preliminaries

The preliminaries model language-model decoding as autoregressive sampling from a policy. Greedy decoding chooses the most probable next token under that policy at each step.

2. Preliminaries

  • An autoregressive policy assigns next-token probabilities conditioned on the prompt and preceding tokens, whose product defines a response probability.Greedy decoding selects the token with the highest conditional probability at each step.Source: §2 · p. 2

3. FUSIONROUTE

FUSIONROUTE routes each token to a selected specialist and adds a router-generated logit to refine that expert’s prediction. Its training separates expert routing from complementary generation, then mixes supervised and preference-optimization data to preserve both capabilities.

3.1. Overall Design

  • At each decoding step, FUSIONROUTE selects the expert with the highest routing weight and adds the router’s logits to the expert’s logits.The router produces both routing weights and complementary logits; this preserves expert strengths while allowing correction of expert behavior.Source: §3.1 · p. 3

3.2. Training FUSIONROUTE

  • Training uses an SFT phase followed by preference optimization because jointly optimizing routing and complementary generation can degrade either objective.The two coupled goals are reliable token-level expert selection and a complementary logit contribution that corrects expert failures.Source: §3.2 · p. 3

3.2.1. SFT PHASE

  • SFT jointly trains the base model and routing projection for next-token prediction and expert selection, restricting routing supervision to positions where experts disagree.The disagreement filter avoids letting tokens with similar expert predictions dominate routing supervision.Source: §3.2.1 · p. 3

3.2.2. COMPLEMENTED DIRECT PREFERENCE OPTIMIZATION (CDPO)

  • CDPO trains the router’s complementary logits through preference optimization while treating expert outputs as fixed.The router is trained to refine expert outputs that are locally suboptimal; its corrective signal is intended to compensate for expert weaknesses.Source: §3.2.2 · pp. 3–4

3.2.3. MIX TRAINING

  • Mixed training applies routing and language-model losses to SFT samples, CDPO to preference samples, and excludes the routing projection from preference updates.The procedure mixes supervised and preference-optimization data to preserve expert selection while learning complementary logits.Source: §3.2.3 · pp. 4–5

4. Theoretical Analysis

The analysis models decoding as sequential token decisions and shows that selecting among fixed expert outputs is generally insufficient without strong coverage assumptions. A complementary router logit expands the policy class, allowing approximation of the optimal policy even when no single expert suffices.

4.1. Token-Level Markov Decision Process

  • The decoding process is a token-level MDP whose state is the prompt and generated prefix, whose action is the next token, and whose reward scores the resulting text.A language model is represented as a policy over next-token choices given the preceding text.Source: §4.1 · p. 5

4.2. Token-Level Expert Selection through the Lens of Performance Difference Lemma

  • The Performance Difference Lemma implies that choosing an expert whose next token approximately maximizes optimal value at each step yields a near-optimal final policy.This connection assumes SFT data are generated by the optimal policy; global coverage over all prefixes, including unseen ones, suffices when the value gap is o(1/T).Source: §4.2 · p. 5

4.3. Limitations of Purely Token-Level Collaboration

  • Pure token-level routing cannot guarantee a near-optimal policy from optimal-value observations, even when near-optimal trajectories exist and value gaps are bounded on good prefixes.The analysis attributes this limitation to an identifiability failure: optimal values along optimal-policy trajectories do not determine which expert actions realize them.Source: §4.3 · p. 6

4.4. Benefits of Router Training

  • Adding the router’s logits to the selected expert’s logits removes the requirement that any one expert be near-optimal, provided the final policy stays close to the optimal policy.The adaptive logits expand the action space, while a candidate already aligned with the optimal policy lets the router focus corrections on states where it deviates.Source: §4.4 · p. 6

5. Experiments

Across mixed-domain benchmarks, FUSIONROUTE outperforms collaboration, model-merging, and fine-tuning baselines while matching or exceeding specialists on their strongest tasks. Its GPT-4o win rate over a fine-tuned baseline is higher on general prompts, and its advantage over other baselines grows in the larger Llama-3 setting.

5.1. Experimental Setup

  • Experiments use math, coding, and instruction-following experts from Llama-3 and Gemma-2, with sequence-level, token-level, model-merging, and fine-tuning baselines.The setup draws on PerfectBlend for training and testing and applies SFT followed by a mixed training procedure.Source: §5.1 · pp. 6–7

5.2. Main Results

  • FUSIONROUTE consistently outperforms sequence-level collaboration, Collab, model merging, and directly fine-tuned models across Llama-3 and Gemma-2 benchmarks.The comparison covers cross-domain performance across the two model families.Source: §5.2 · p. 7, Table 1 · p. 7
  • FUSIONROUTE matches or exceeds domain experts on their strongest math and coding benchmarks while remaining robust across mixed domains.The reported examples include math experts on GSM8K and MATH500 and coding experts on MBPP and HumanEval.Source: §5.2 · p. 8

5.2.2. PERFORMANCE ON GENERAL DATASET

  • On general prompts, FUSIONROUTE achieves a substantially higher GPT-4o win rate than the fine-tuned baseline on both model families.The evaluation uses 500 prompts sampled from the PerfectBlend test set.Source: §5.2.2 · p. 8, Figure 3 · p. 8

5.2.3. FUSIONROUTE IS MORE BENEFICIAL AT A LARGER SCALE

  • The performance gap over other baselines becomes markedly larger for 8B Llama-3, where Collab and Sequence Selection lose substantial win rate.In the 2B Gemma-2 family, the gap over pure token-level collaboration is smaller; the authors suggest less room for complementary correction at that scale.Source: §5.2.3 · p. 8

6. Ablation Study

Ablations indicate that both complementary logits and preference optimization contribute to FUSIONROUTE’s reported gains. Routing alone performs better than Collab across most benchmarks, but the full framework more consistently outperforms routing without complementary logits.

6.1. Ablation on Complementary Logit Contribution

  • FUSIONROUTE outperforms its routing-only variant across nearly all benchmarks and both model families, with especially pronounced gaps on coding and instruction-following tasks.The authors attribute the gap to selected experts sometimes producing locally suboptimal or misaligned tokens that need correction.Source: §6.1 · p. 8
  • The routing-only variant already outperforms Collab across most benchmarks, while the authors report greater instability for approaches relying solely on external reward signals.They identify direct training of the routing component on expert data as important for accurate and stable routing.Source: §6.1 · p. 9

6.2. Ablation on FUSIONROUTE Training Procedure

  • With CDPO, the full FUSIONROUTE model significantly outperforms its SFT-only counterpart in GPT-4o win rate on Llama-3-8B and Gemma-2-2B.The authors report that SFT provides a reasonable initialization, while preference optimization substantially improves general response quality.Source: §6.2 · p. 9, Figure 4 · p. 9

7. Conclusion

FUSIONROUTE combines flexible token-level expert selection with complementary logits, and the conclusion reports effectiveness on domain-specific and general datasets. The related-work sections distinguish its approach from token-level collaboration and model integration methods.

7. Conclusion

  • The conclusion reports that extensive experiments on domain-specific and general datasets verify FUSIONROUTE’s effectiveness.It combines token-level collaboration with complementary logits and avoids requiring each model to generate candidate tokens or full responses for selection.Source: §7 · p. 9

A.1. Multi-LLM Collaboration

  • FUSIONROUTE differs from prior token-level collaboration by using multiple specialized LLMs, evaluating cross-domain performance and general win rate, and training the router as a complementary generator.The passage contrasts this with Collab’s controlled decoding and CO-LLM’s reliance on one strong math expert assisting a small fine-tuned model.Source: §A.1 · p. 12

A.2. Multi-LLM Integration

  • Unlike MoE and model merging, FUSIONROUTE does not jointly train or merge experts and requires neither architectural compatibility nor gradient access.It uses a lightweight router for token-level integration of heterogeneous, off-the-shelf LLMs.Source: §A.2 · p. 12

B.1. Statement and Proof of Theorem 4.3

Theorem 4.3 establishes a limitation of token-level routing: even when a near-optimal path exists, routing cannot guarantee near-optimal value for all constructed reward settings. The proof uses indistinguishable observations early in decoding to show that some required expert-selection path can be missed.

B.1. Statement and Proof of Theorem 4.3

  • Under the theorem’s stated setting, no token-level routing algorithm can guarantee value at least V* − T/2 + 2 across all possible rewards.The theorem assumes deterministic policies and transitions and a fixed response length T; the excerpt also states that a near-optimal path exists.Source: §B.1 · pp. 12–13
  • The constructed MDPs give the learner identical observations during the first T/2 − 1 steps, making them indistinguishable during that part of decoding.The learner observes expert values along the visited trajectory and possible next-token continuations.Source: §B.1 · pp. 15–16
  • Because the same observations induce a distribution over possible expert-selection paths, at least one path has probability at most 1/|P_T/2|.The proof obtains this bound by counting the possible paths.Source: §B.1 · p. 16
  • For the MDP associated with that low-probability path, the routing algorithm’s expected value is at most V* − T/2 + 2.The proof concludes that the algorithm cannot guarantee approximately optimal value for every constructed reward setting.Source: §B.1 · p. 16

B.2. Proof of Eq. (6) in Section 4.4

The proof of Eq. (6) invokes the Performance Difference Lemma and describes a per-token policy that selects an expert.

B.2. Proof of Eq. (6) in Section 4.4

  • The proof invokes the Performance Difference Lemma to derive the result for Eq. (6).The excerpt does not state the resulting bound.Source: §B.2 · p. 16
  • At each token, the policy chooses an expert πi according to a criterion stated in the proof.The criterion itself is not included in this excerpt.Source: §B.2 · p. 16

C. Theoretical Discussion of Prior Token-Level Approaches

The theoretical discussion connects SFT-based router training to approximating optimal expert selection, while contrasting this with prior token-level methods. It also describes a constructed setting where a mismatch can scale linearly with the horizon.

C. Theoretical Discussion of Prior Token-Level Approaches

  • Router training on SFT trajectories generated by the optimal policy can be viewed as approximating an optimal expert-selection strategy, with the Performance Difference Lemma linking generated responses to optimal responses.FUSIONROUTE directly fine-tunes the router on the expert dataset, which the passage says helps bound final error using the lemma.Source: C. Theoretical Discussion of Prior… · p. 17
  • Collab evaluates candidate tokens using each expert’s Qπi function and selects the token with the maximum corresponding value, rather than maximizing the optimal Q function.The discussion says this action choice can mismatch the generated and optimal responses.Source: C. Theoretical Discussion of Prior… · p. 17
  • In a constructed setting with two experts that disagree at every prefix, the optimal policy uses π1 in the first half and π2 in the second half, attaining Q*(x) = H.The discussion says the mismatch error is linear in horizon H.Source: C. Theoretical Discussion of Prior… · p. 17

D. Ablation Study

This appendix block describes the ablation comparison and evaluation setup, then presents examples spanning coding, instruction-following, and math prompts. It also specifies how router behavior is color-coded.

D. Ablation Study

  • Table 2 compares purely token-level routing with FUSIONROUTE across model families, highlighting the best and second-best results within each family.The text points to the ablation results but supplies no table values.Source: Table 2 · p. 17, D. Ablation Study · p. 17

E.1.1. MIXED-DOMAIN GENERATION

  • Because some tasks have many prompts and computation is constrained, evaluation uses dataset subsets; Table 3 lists sample counts and maximum new-token settings.Source: §E.1.1 · p. 18, Table 3 · p. 18

E.1.2. GPT EVALUATION PROMPT

  • The GPT-4o evaluation prompt requests correctness, fluency, and detail ratings on a 1–10 scale, with distinct scores for the two responses.Source: §E.1.2 · p. 18

E.2. Router’s Behavior

  • Router-behavior examples cover domain-specific and general datasets; red, blue, and soft green denote the math, coding, and instruction-following experts.Source: §E.2 · p. 18

E.2.1. EXAMPLE OF CODING PROBLEM

  • The coding example asks for a Python function to find the minimum-cost path through a cost matrix; its accompanying explanation describes a dynamic-programming solution.The explanation initializes the first row and column, then uses the cheaper path from above or left plus the current cell cost.Source: §E.2.1 · p. 19

E.2.2. EXAMPLE OF INSTRUCTION FOLLOWING PROBLEM

  • The instruction-following example asks for a resume template for a fresh high-school graduate, with at least 12 square-bracketed placeholders.Source: §E.2.2 · p. 19

E.2.3. EXAMPLE OF MATH PROBLEM

  • The math example asks to convert (0, 3) from rectangular to polar coordinates and gives (3, π/2) as the answer.Source: §E.2.3 · p. 20

E.2.4. EXAMPLE OF GENERAL PROBLEM IN PERFECTBLEND DATASET

The examples illustrate a general arithmetic problem, while the discussion reports that routing patterns vary with token type and domain. On general-purpose data, reasoning and discourse tokens tend to go to different experts.

E.2.4. EXAMPLE OF GENERAL PROBLEM IN PERFECTBLEND DATASET

  • The example problem asks how many pens were held initially after receiving 22, doubling the total, giving away 19, and ending with 65; its worked answer is 20.Source: §E.2.4 · p. 20

E.2.5. DISCUSSION ON THE QUALITATIVE EXAMPLES

  • FUSIONROUTE assigns the majority of tokens associated with coding, mathematical reasoning, and instruction following to their corresponding expert models in domain-specific examples.Source: §E.2.5 · p. 20
  • On general-purpose datasets, mathematical or logical reasoning tokens predominantly go to math and code experts, while discourse, formatting, and instruction-phrasing tokens more often go to the instruction-following model.The passage characterizes this as adaptive routing and says it improves generation quality across diverse tasks.Source: §E.2.5 · p. 20
Loading 2601.05106v5…