Spend Experts Where You Are Unsure: Confidence-Adaptive Routing for Mixture-of-Experts LoRA

Overview§

CARE (Confidence-Adaptive Routing of Experts) is a drop-in replacement for the fixed top-k routing in Mixture-of-Experts (MoE) Low-Rank Adaptation (LoRA). The key insight is that the router's output distribution (softmax over experts) inherently captures per-token uncertainty: a peaked distribution indicates high confidence, while a flat distribution indicates ambiguity. Instead of routing every token to the same number of experts $k$, CARE adaptively selects a variable number of experts per token, allocating more compute to hard tokens and less to easy ones.

Method§

CARE proceeds in two steps: 1. Nucleus Routing: Given router weights $\mathbf{w} \in \mathbb{R}^E$ (where $E$ is the number of experts), sort them descending: $w_{(1)} \ge w_{(2)} \ge \dots \ge w_{(E)}$. Then select the smallest set of top experts such that their cumulative mass exceeds a threshold $\tau$: $\min\{m: \sum_{i=1}^{m} w_{(i)} \ge \tau\}$. Additionally, if the selected experts disagree (i.e., their outputs have high variance), one more expert is added (a small extension). This disagreement signal improves handling of ambiguous tokens. 2. Budget Thermostat: To control the average number of active experts across a batch, the threshold $\tau$ is adjusted per step. Let $B$ be the target average budget (e.g., 2.5 experts per token). The thermostat maintains a running estimate of the current average $\bar{k}$ and adjusts $\tau$ via a simple update: $\tau \leftarrow \tau + \eta (\bar{k} - B)$, where $\eta$ is a learning rate. This ensures that over time, the average active expert count matches the desired budget without per-token tuning.

No additional parameters are needed; CARE uses only the router logits and the outputs of the initially activated experts.

Implementation§

The code for CARE is straightforward. Given router logits logits of shape (batch, num_experts) and a budget target budget:

import torch

def care_routing(logits, budget, expert_outputs, tau_init=0.5, eta=0.01):
    # logits: (B, E)
    # expert_outputs: (B, E, D)  -- outputs from each expert (computed after selection for efficiency)
    probs = torch.softmax(logits, dim=-1)
    sorted_probs, indices = torch.sort(probs, descending=True, dim=-1)
    cumsum = torch.cumsum(sorted_probs, dim=-1)
    # cumulative mass >= tau
    mask = cumsum >= tau_init
    # find first index where mask is True
    _, first_exceed = torch.max(mask, dim=1)  # (B,)
    k = first_exceed + 1  # number of experts for each token
    # disagreement check (compute variance of selected expert outputs)
    # ... (simplified) if variance > threshold, add one more expert
    # adjust tau via thermostat
    avg_k = k.float().mean()
    tau = tau_init + eta * (avg_k - budget)
    # gather selected indices
    selected = [indices[i, :k[i]] for i in range(len(k))]
    # compute weighted sum of expert outputs
    # ... (uses selected weights)
    return output, tau

Note: In practice, CARE computes expert outputs only for the candidates before selection, but can use a two-pass strategy to avoid activating all experts.

Results§

Across eight commonsense benchmarks on LLaMA-3.1-8B and Qwen2.5-7B, CARE improves over fixed top-k MoE-LoRA at matched compute. For example, with average budget 2.5, CARE matches the performance of fixed top-4 while activating fewer experts on average. The confidence and disagreement signals also improve out-of-distribution detection over baselines like maximum softmax probability (MSP) and entropy.

Key Equations§

• Nucleus selection: $\mathcal{S}_t = \text{top-m experts s.t.} \sum_{i=1}^{m} w_{(i)} \ge \tau$
• Budget thermostat update: $\tau \leftarrow \tau + \eta (\bar{k} - B)$
• Disagreement extension: if $\text{Var}(\text{expert outputs}) > \delta$, add one more expert.