PaperlayerView exact revision on arXiv

Source-linked AI summary

Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing

Hanchi Sun, Yixin Liu, Yonghui Wu, Lichao Sun

arXiv:2603.11535v1cs.AIcs.CL

TL;DR

MoE routing must balance sparse computation with expert utilization and causal autoregressive inference. The paper proposes Expert Threshold routing, which uses EMA-estimated global thresholds to route tokens independently at training and inference. On FineWeb-Edu, ET outperforms Token Choice by 0.067 cross-entropy loss while achieving near-perfect load balancing and causal generation.

  • Problem

    Sparse routing can collapse onto a few experts, while Token Choice fixes per-token sparsity and Expert Choice depends on future batch tokens during autoregressive inference.

  • Method

    Expert Threshold routing uses EMA-estimated expert-specific quantile thresholds to route each token independently at training and inference.

  • Results

    0.067 cross-entropy loss: ET outperforms TC during 2.4B-parameter FineWeb-Edu pretraining while achieving near-perfect load balancing.

  • Takeaways & Limitations

    ET provides causal autoregressive routing without fixed per-token computation or auxiliary load-balancing losses, while matching EC validation loss at 2.84.

  • Takeaways & Limitations

    Expert Choice remains non-causal for autoregressive generation because its routing depends on all tokens in the batch, including future tokens.

Abstract

from arXiv · show

Token-choice Mixture-of-Experts (TC-MoE) routes each token to a fixed number of experts, limiting dynamic computation allocation and requiring auxiliary losses to maintain load balance. We propose Expert Threshold (ET) routing, where each expert maintains an exponential moving average (EMA) threshold estimated from the global token distribution. At both training and inference, each token is independently routed to an expert if its score exceeds the expert's threshold, enabling dynamic computation allocation while achieving load balance without auxiliary losses. This fully causal mechanism eliminates dependence on other tokens in the batch, making it well-suited for autoregressive language modeling. In pretraining experiments scaling to 2.4B parameters on FineWeb-Edu, ET achieves 0.067 lower cross-entropy loss than TC-MoE, equivalent to reaching the same performance with 1.6$\times$ fewer tokens.

1. Introduction

Sparse routing scales model capacity efficiently but creates load-imbalance and causality challenges. Expert Threshold routing addresses these constraints with causal, threshold-based token routing and improves pretraining loss over Token Choice.

  • Motivation: Mixture-of-Experts models scale capacity efficiently by sparsely activating only a subset of expert networks per token.This decouples parameter count from computational cost, but sparse routing can collapse onto a small subset of experts.
  • Motivation: Load imbalance leaves some experts underutilized and creates hardware bottlenecks, motivating routing that roughly maintains balanced utilization.
  • Prior Routing: Token Choice fixes the number of experts per token, complicating load balancing and motivating auxiliary losses or other heuristics.
  • Expert Threshold: Expert Threshold routes tokens by comparing scores with expert-specific quantile thresholds tracked from global score distributions.The thresholds are used at both training and inference, avoiding train-inference mismatch and enabling fully causal routing.
  • Results: 0.067 cross-entropy loss: ET outperforms TC when pretraining a 2.4B-parameter language model on FineWeb-Edu while achieving near-perfect load balancing.

2. Preliminaries: Routing as Constrained Optimization

MoE routing can be formalized as assigning tokens to experts under computation and load constraints. Token Choice enforces per-token sparsity, whereas Expert Choice removes that constraint and achieves dynamic computation with exact batch-level balancing, but loses causality.

  • Routing Formulation: The router produces token-expert scores and a binary assignment indicating which experts process each token.Selected experts contribute outputs weighted by sigmoid gate values.
  • Routing Formulation: The routing assignment controls both computation allocation and expert load balance, making these competing constraints central to the optimization problem.
  • Token Choice Routing: Token Choice selects exactly G experts for each token and relies on auxiliary or loss-free strategies to address load balancing.Exact optimization under simultaneous sparsity and load constraints is combinatorial.
  • Expert Choice Routing: Expert Choice removes per-token sparsity and selects the top-k tokens for each expert within a batch.This provides exact batch-level balancing while allowing tokens to receive zero, one, or multiple experts.
  • Causality Limitation: Expert Choice is non-causal for autoregressive generation because each expert’s selection depends on all batch tokens, including future tokens.Batch-level top-k only partially alleviates this problem because routing still depends on batch composition.

3. Expert Threshold

Expert Threshold relaxes batch-level routing constraints by estimating global expert thresholds with exponential moving averages. It then routes tokens independently using those thresholds, preserving causal operation while balancing load in expectation.

  • Design: ET replaces per-sequence or per-batch load balancing with stochastic load balancing in expectation.This targets the population-level activation rate rather than enforcing exact utilization for every batch.
  • Threshold Estimation: Each expert’s threshold estimates a global (1 − 1/E)-quantile using an exponential moving average of batchwise kth-largest router logits.
  • Routing Rule: ET routes token t to expert i when its router score exceeds expert i’s threshold, producing a binary routing indicator.The rule uses score and threshold comparisons independently for each token-expert pair.
  • Causality: Because routing depends only on the current score and global threshold, ET is fully causal while satisfying load balancing in expectation.
  • Connection to EC: ET approximates Expert Choice over an infinitely large batch, where each token’s influence on the selection threshold vanishes.
  • Training Behavior: ET fixes thresholds for stable routing decisions but accepts small variance in per-batch expert utilization, unlike Expert Choice.A 4k-step Expert Choice warmup addresses cold-start threshold estimation before switching to ET.

4. Experiments

Experiments compare ET with EC and TC across model scales, routing batch sizes, computation allocation, specialization, and train-evaluation consistency. ET delivers strong language-modeling results while maintaining causal, batch-independent routing and dynamically varying computation.

  • Experiment Setup: The experiments use 575M-parameter d12 and 2.4B-parameter d20 models with 16 routed experts, one shared expert, and FineWeb-Edu training.Each token activates the shared expert and, on average, one routed expert; the models train on 10B and 11.2B tokens, respectively.
  • Main Results: ET consistently outperforms TC in CE loss and CORE across both model scales, improving CE by 0.05 and CORE by 1.89 on d12, and CE by 0.067 and CORE by 2.83 on d20.EC with large batches achieves comparable CE loss, while EC 512k slightly exceeds ET on d12 CORE, 19.94 versus 19.88.
  • Cutoff vs Expert Usage Tradeoff: ET stabilizes the cutoff through an EMA while allowing expert usage to fluctuate around the capacity target, trading fixed hardware usage for train-inference uniformity.EC instead fixes expert usage but has batch-to-batch cutoff variation; ET remains at zero cutoff deviation by design.
  • Dynamic Computation Allocation: Both methods allocate more computation to early positions, but EC 2k shows a first-token fanout near 10 whereas ET rises more mildly to about 2 before declining.When grouped by loss, EC 2k fanout increases with token loss, while ET remains flatter overall.
  • Expert Specialization: ET produces sharper expert-activation patterns than EC on HumanEval and GSM8K, while both methods specialize experts for domain-specific tokens.The analysis measures expert token ratios across code and math evaluation sets; ET achieves comparable specialization without large-batch coordination at inference.
  • Batch Size Scaling: EC improves as routing batches grow from 2k to 64k tokens, then saturates at 512k, while ET reaches comparable performance without batch-size coordination.EC training CE changes from 2.874 at 2k to 2.836 at 64k and 2.840 at 512k; ET records 2.844 CE and 19.876 CORE Eval.
  • Train-Evaluation Gap: Small-batch EC has a substantial train-evaluation gap, whereas ET maintains consistent routing at inference and remains close to large-batch EC across checkpoint pairs.EC 2k reaches 17.91 CORE versus 19.94 for EC 512k, while ET is above EC 2k on every checkpoint pair and close to EC 64k.

5. Related Work

Related work addresses MoE scaling, load balancing, dynamic computation, and causal routing through increasingly global statistics and adaptive selection mechanisms. ET combines expert-specific EMA thresholds with causal threshold routing to connect these directions.

  • MoE scaling and load balancing: MoE architectures scale model capacity efficiently by sparsely activating expert subsets, but routers can collapse onto a few experts and create device-level load imbalance.Load imbalance leaves some devices idle and others overloaded, motivating routing mechanisms that maintain approximate balance.
  • Load-balancing methods: Token-choice routing fixes each token’s expert count, while auxiliary losses, PID controllers, and bias updates provide heuristic or loss-free load-balancing corrections.Auxiliary-loss methods can create biased router logits, whereas proportional bias updates scale with load deviation.
  • From batch to population statistics: ET extends population-level statistics to routing by maintaining EMA-based expert cutoffs rather than relying on per-batch top-k selection.This connects ET to momentum-based methods and adaptive optimizers that approximate population distributions with EMAs.
  • Routing stability: Figure 8 reports that ET is more stable than EC 2k and remains close to EC 64k on within-family checkpoint-pair routing consistency.TC is competitive for nearby checkpoints but degrades more across longer ranges.
  • Method taxonomy: Table 3 classifies load-balancing methods by scope, while Table 4 relates ET conceptually to LossFree and GShard.The supplied table passages provide taxonomy and conceptual-connection labels without numerical results.
  • Dynamic computation and causality: Expert Choice enables variable computation by letting each expert select top-k tokens, but its batch-dependent selection creates a causality challenge for autoregressive inference.Prior approaches use predictors, per-expert thresholds, or cross-sequence token pools to approximate causal routing.
  • Dynamic computation and causality: Dynamic-computation alternatives adapt expert counts using differentiable routing, zero-computation experts, cumulative probability thresholds, or token-level confidence.XMoE is closest to ET but uses a fixed probability-mass threshold, whereas ET uses expert-specific EMA cutoffs to causalize expert choice.
  • Causal routing: At inference, ET compares each token’s logit with an EMA cutoff, eliminating the train–inference discrepancy while preserving causal routing.The threshold test depends only on the token and historical cutoff, avoiding future-token access and large-batch coordination.

6. Conclusion

The paper introduces ET routing to address EC’s causality issue while retaining load-balancing advantages. It reports competitive EC performance, improved TC loss, causal generation, and minimal train–inference gaps.

  • Conclusion: ET maintains EMA selection thresholds from historical batches, so token routing depends only on past statistics during training and inference.This makes ET fully causal for autoregressive generation.
  • Conclusion: ET matches EC validation loss at 2.84 and outperforms TC by 0.067 in cross-entropy loss while enabling causal autoregressive generation.The conclusion also reports minimal train–inference gaps and a warmup strategy that stabilizes early training.

Impact Statement

The paper frames Expert Choice’s future-information leakage as a consequence of infinite-precision thresholds, while finite-precision thresholds bound leakage independently of sequence length. It introduces binary-search encoding to show that the infinite-precision upper bound is achievable.

  • Impact Statement: Infinite-precision Expert Choice thresholds can leak at least O(N log N) bits of future information, making the combinatorial upper bound tight.The paper considers both finite- and infinite-precision cutoff representations.
  • Impact Statement: The analysis motivates finite-precision thresholds as a condition under which Expert Choice routing remains causally implementable.The finite-precision proof uses the cutoff itself as the advice variable.
  • Impact Statement: The paper formalizes future-information leakage as the minimum advice-variable range needed for causal functions to reproduce a selection rule.The selection rule maps a full logit sequence to token-selection indicators, while causal decoders access prefixes and finite advice.
  • Impact Statement: Finite-precision cutoff thresholds have constant total future-information leakage, so per-token leakage approaches zero as sequence length increases.The bound is at most b bits when the threshold uses b-bit precision.
  • Impact Statement: Binary-search encoding partitions cutoff space into intervals that represent distinct expert-selection patterns.Each possible pattern is injectively mapped to a unique interval, while midpoint queries reveal successive decision bits.

A.5. Formal Proof

The formal proof establishes that infinite-precision Expert Choice selection requires substantial future information, while binary-search decoding realizes the corresponding lower bound.

  • A.5. Formal Proof: Expert Choice routing requires future information because its selection rule depends on a cutoff encoding token-selection patterns.The appendix proves this requirement for all possible top-k assignments.
  • A.5. Formal Proof: For one expert with capacity k = ⌊N/E⌋, any causal mechanism realizing all top-k assignments requires at least log2-sized advice proportional to the selection space.The supplied theorem statement establishes the lower bound, although the displayed continuation is incomplete.
  • A.5. Formal Proof: Binary-search decoding queries the midpoint of the current cutoff interval and updates its lower or upper boundary from each selection bit.After N steps, the procedure returns the decision sequence and the remaining interval.
  • A.5. Formal Proof: Figure 9 depicts the cutoff c partitioning [0, 1] into regions associated with different expert-selection patterns.The regions support encoding N bits of future information in one real-valued threshold.
  • A.5. Formal Proof: Each distinct valid selection pattern requires distinct advice, because identical prefixes and advice would force a causal decoder to produce the same decision for different patterns.The argument uses the first position where two patterns differ.
  • A.5. Formal Proof: The proof sums the single-expert lower bound across GE experts to recover the combinatorial leakage quantity.This establishes that Expert Choice fundamentally requires significant future information for causal implementation.
  • A.5. Formal Proof: The appendix’s proof discussion is distinct from the paper’s nanochat-style architecture description.The architecture passage identifies a separate model-design context rather than a proof step.

B.1. Data and Tokenization

The experiments use FineWeb-Edu with RustBPE tokenization, 2048-token sequences, and nanochat-style model configurations.

  • B.1. Data and Tokenization: Training uses the FineWeb-Edu 100B shuffle dataset, a high-quality educational web corpus.Sequences contain 2048 tokens.
  • B.1. Data and Tokenization: RustBPE tokenization uses a 65,536-token vocabulary aligned to a power of two for GPU efficiency.The vocabulary size is 64k tokens.
  • B.1. Data and Tokenization: The model configurations follow nanochat’s naming convention, with d* denoting layer count and nembd = d×64.The attention head dimension is fixed at 128.
  • B.1. Data and Tokenization: Attention uses grouped-query attention with head dimension 128, head count nhead = nembd/128, and QK normalization.The d12 and d20 configurations use 6 and 10 heads, respectively.
  • B.1. Data and Tokenization: MoE variants use 16 routed experts plus one shared expert, with each token activating one routed expert on average.The shared expert processes every token, matching the dense model’s active parameter count.

C. Training Setup Details

The training setup combines nanochat-style initialization and hybrid optimization with custom MoE implementations, expert parallelism, and shared architecture configurations.

  • C. Training Setup Details: Initialization uses aspect-ratio-scaled weight variance, reducing to standard 1/√din initialization for square or tall matrices.Wide matrices receive lower variance when dout ≫ din.
  • C. Training Setup Details: The architecture tables cover shared nanochat-style configurations for d12 and d20 models, including total and active parameter accounting.MoE variants use G=1 and E=16 with 16 routed and one shared expert.
  • C. Training Setup Details: Component-specific initialization uses standard-normal embeddings, zero output projections, small router weights, and specialized expert initialization.Attention weights use the aspect-ratio-scaled scheme.
  • C. Training Setup Details: Muon optimizes attention, MLP, and expert matrices, while AdamW optimizes embeddings and the output head.Muon uses Newton-Schulz orthogonalization; AdamW uses learning-rate scaling.
  • C. Training Setup Details: No weight decay is used because Muon supplies implicit regularization and language models benefit from memorization.
  • C. Training Setup Details: Training runs on one node with 8 NVIDIA B200 GPUs, each providing 180GB of memory.
  • C. Training Setup Details: TC uses ScatterMoE, whereas EC and ET use custom PyTorch MoE implementations with padding for variable expert loads.
  • C. Training Setup Details: Custom expert-parallel all-to-all communication enables maximum routing batch size, improving the usage–cutoff variance trade-off while saving memory.The implementation uses Nanochat’s distributed AdamW and Muon with ZeRO-2-style gradient synchronization.

D. CORE Evaluation Details

The evaluation details describe CORE’s centered-accuracy protocol and analyze ET’s training stability, normalization, capacity behavior, and routing-consistency metrics. Results indicate warmup and shared experts improve training behavior, while capacity constraints rarely create train-inference mismatch.

  • Evaluation protocol: CORE evaluates language models on multiple-choice, schema-matching, and language-modeling tasks using centered accuracy averaged across tasks.Centered accuracy maps random guessing near zero and perfect accuracy to 1.0.
  • Training dynamics: ET warmup is necessary because slow cutoff-EMA updates make early threshold routing unreliable and frequently trigger the capacity lower bound.The method uses TopK selection for the first 4,000 steps before switching to threshold routing.
  • Routing stability: ET and EC stabilize cutoff-EMA quickly, whereas DeepSeek’s loss-free controller drifts upward during training, especially in early layers.The authors characterize this as exploratory because the behavior may depend on hyperparameters and gating parameterization.
  • Ablations: Shared experts improve loss by roughly 0.02 in both ET no warmup and EC configurations.The comparison includes matched routed-expert settings with and without a shared expert.
  • Normalization: No normalization outperforms fanout normalization by 0.04 in cross-entropy loss in this setting.The authors suspect normalization makes expert contributions unpredictable.
  • Capacity constraints: After warmup, raw expert usage stabilizes around 6.5%, while saturation and starvation remain low, indicating minimal capacity-induced train-inference mismatch.ET enforces per-batch capacity constraints during training even though inference lacks those constraints.

F.3. Activation Dynamics Sweep Across Routing Variants

The activation-dynamics sweep compares routing variants across loss, fanout, routing consistency, and domain specialization. ET develops specialization comparable to large-batch EC while avoiding batch-size coordination at inference, whereas several activation relationships remain setup dependent.

  • Activation dynamics: EC 8k is flatter than EC 2k in inverse layerwise activation views, while loss–fanout relationships vary substantially across routing variants and depth.The sweep includes ET no warmup, EC 64k, and EC 512k as additional variants.
  • Token-level routing: Content-bearing numerical, mathematical, and code-specific tokens receive higher fanout than function words, punctuation, boilerplate text, and whitespace.The visualizations cover GSM8K and HumanEval passages.
  • Routing consistency: ET routing remains clearly separated from EC 2k and close to EC 64k in joint-JSD checkpoint-pair consistency.Lower joint JSD indicates more stable routing.
  • Batch-size effects: Increasing EC batch size sharpens specialization, with concentrated expert activation replacing diffuse and noisy activation patterns.The heatmaps compare EC batch sizes 2k, 8k, 64k, and 512k with ET.
  • Routing variants: ET achieves specialization comparable to EC at 512k, with concentrated expert-domain associations across layers without batch-size coordination at inference.The result is attributed to ET’s population-level threshold mechanism capturing routing structure found in large-batch top-k selection.
  • Domain specialization: Experts show different activation patterns for HumanEval code and GSM8K math, indicating domain-level specialization across routing configurations.Experts heavily activated for code can show low activation for math, and vice versa.
Loading 2603.11535v1…