DLLG: Dynamic Logit-Level Gating of LLM Experts

Technical Summary: DLLG: Dynamic Logit-Level Gating of LLM Experts

Problem Statement
The proliferation of specialized Large Language Models (LLMs) optimized for narrow domains (e.g., mathematical reasoning, code generation) presents a challenge: how to integrate independently trained "black-box" experts into a unified system that leverages their complementary strengths without prohibitive retraining costs. Existing approaches suffer from a persistent tension between adaptivity, robustness, and practicality:

• Routing-based methods rely on hard, early expert selection. This leads to premature commitment; if the chosen expert fails on a specific sub-task within a single response, the error is unrecoverable. Furthermore, many routers require extensive benchmark-specific calibration.
• Heuristic Ensembling (e.g., perplexity weighting, token voting) uses inference-time proxy signals. These metrics are often empirical and misaligned with true task correctness, leading to suboptimal outputs.
• Parameter Merging (e.g., model souping, task arithmetic) statically fuses weights. This sacrifices modularity and often induces destructive interference when combining experts with conflicting gradient directions, resulting in performance degradation.
• Mixture-of-Experts (MoE) architectures typically require end-to-end joint training of gates and experts, limiting their ability to utilize off-the-shelf, pre-trained specialists.

Methodology: DLLG
The authors propose DLLG (Dynamic Logit-Level Gating), a framework that performs dynamic, token-level logit fusion using a lightweight gating model trained on sparse, response-level supervision.

1. Architecture:

   • Frozen Experts: $K$ specialized LLMs remain frozen. At each decoding step $t$, given a prompt $x$ and prefix $y_{<t}$, each expert $E_i$ produces a hidden state $h^{(i)}_t$ and next-token logits $\ell^{(i)}_t$.
   • Gating Model: A lightweight model ($g_\phi$) processes the concatenated hidden states of all experts. It projects these representations into a shared embedding space (using low-rank projection) and processes them through an autoregressive stem (e.g., Qwen2.5-0.5B-Instruct with LoRA adapters).
   • Weight Prediction: A prediction head outputs token-wise fusion weights $w_{t,i}$ via a sigmoid transformation. These weights are applied to the experts' logits: $\ell_t = \sum_{i=1}^K w_{t,i} \ell^{(i)}_t$. The weights are not constrained to sum to one, allowing for soft, flexible contributions.

2. Training (Response-Level Supervision):

   • The framework addresses the lack of token-level annotations by using response-level correctness signals ($s_i \in \{0, 1\}$) derived from automatic verifiers (e.g., unit tests for code, answer checking for math).
   • Teacher Forcing: During training, experts generate hidden states conditioned on the ground-truth prefix $y_{<t}$. The gating model predicts weights $w_{t,i}$ for each token position.
   • Objective: The model is trained to minimize the Mean Squared Error (MSE) between the predicted weights and the broadcast response-level correctness signal: $\mathcal{L} = \frac{1}{T} \sum_{t=1}^T \sum_{i=1}^K (w_{t,i} - s_i)^2$. This teaches the gate to align token-level fusion with the overall success of the trajectory.

3. Inference:

   • DLLG operates autoregressively. At each step, the gating model predicts fusion weights based on current expert hidden states, logits are combined, and the next token is selected. This allows for recoverable expert utilization; if an expert makes an error early on, the gating model can dynamically shift weight to a more competent expert in subsequent steps.

Key Contributions

1. Framework Proposal: Introduction of DLLG, a dynamic logit-level ensembling framework that bridges token-level generation with sparse response-level supervision, enabling fine-grained expert utilization without token-level annotation or online reinforcement learning.
2. Mechanism Analysis: Demonstration that the learned gating mechanism dynamically adapts expert contributions during rollout, adjusting weights based on the evolving context (e.g., shifting from math reasoning to code generation within a single response).
3. Empirical Validation: Evaluation across diverse reasoning and code benchmarks (GSM8K, MATH, HumanEval, etc.) and model scales (0.5B, 1.5B), showing consistent gains over routing, heuristic ensembling, and parameter-merging baselines.

Results

• Performance: Across both 0.5B and 1.5B model scales, DLLG consistently achieves the highest average performance compared to strong baselines, including RouterDC, EmbedLLM, various logit ensembling methods (GaC, UniTe), and parameter merging techniques (Linear, SLERP, Task Arithmetic).
• Robustness: In low-capacity (0.5B) settings where individual experts are weak, DLLG outperforms brittle routing and heuristic methods, demonstrating its ability to mitigate early mistakes through soft fusion.
• Generalization: The method shows strong performance on out-of-domain benchmarks (e.g., HumanEval, MBPP, BBH), indicating that the learned fusion rules generalize beyond the specific training distribution.
• Ablation: Experiments confirm that a sufficiently expressive gating architecture (incorporating an auxiliary LLM and low-rank projection) is necessary for effective token-level fusion, outperforming simple MLP heads.

Significance and Claims
The paper claims that DLLG offers a robust and scalable paradigm for integrating specialized experts. By learning logit-level fusion from weak supervision, it achieves MoE-like adaptivity without the cost of retraining experts or the risk of parameter interference. The authors emphasize that DLLG preserves expert modularity and enables fine-grained, recoverable expert utilization, overcoming the limitations of premature commitment in routing and the static nature of parameter merging. The work suggests that learning fine-grained expert utilization from weak supervision is a practical path forward for leveraging increasingly specialized language models.