Hierarchical Copula-Gumbel-Top-\texorpdfstring{$K$}{K} Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws

Technical Summary: Hierarchical Copula-Gumbel-Top-K Routing

Problem Statement
In sparse Mixture-of-Experts (MoE) models, the standard stochastic routing mechanism (Gumbel-Top-K) defines a "routing law" for each token: a specific probability distribution over ordered lists of experts and mixture weights, conditioned on the token's logits. While nearly all existing work modifies this per-token law (e.g., via fine-tuning logits or auxiliary losses), the joint distribution of routing choices across different tokens remains largely unexamined. The paper asks: given that every individual token's routing law is held exactly fixed (as in a frozen model), what joint distributions over the routing choices of different tokens are reachable? Specifically, can one control the dependence between tokens to improve expert coherence or load balancing without altering the marginal behavior of any single token?

Methodology: Hierarchical Copula-Gumbel-Top-K (H-CGA)
The author proposes H-CGA, a routing-side dependence layer that operates on the Gumbel noise vector before the Top-K selection. The method relies on Sklar's theorem to separate marginal distributions from dependence structures. The construction proceeds in four stages:

1. Pre-Routing Grouping: Tokens are partitioned into disjoint groups (e.g., fixed windows of adjacent tokens). A controller reads frozen pre-routing features (e.g., hidden states, gate entropy) to determine coupling strengths.
2. Within-Group Positive Coupling: For tokens within a group $g$, the Gumbel perturbations for each expert coordinate are generated via an exchangeable Gaussian copula. Tokens share a latent variable $\zeta_{ge}$, inducing positive correlation in their random perturbations. This increases the likelihood that related tokens select the same experts (coherence).
3. Cross-Group Antithetic Coupling: Disjoint pairs of groups $(g, g')$ are matched. A tunable opposition strength $\alpha_{g,g'} \in [0, 1]$ is introduced. The shared latent of the second group is constructed as a linear combination of the first group's latent and an independent noise term. At $\alpha_{g,g'} = 1$, the groups receive fully opposite shared pushes; at $\alpha_{g,g'} = 0$, they are independent. This introduces negative dependence between groups to counteract load burstiness.
4. Preservation of Marginals: The transformation from shared latents to Gumbel variates ensures that for any single token $t$, the vector of perturbations $(\gamma_{t1}, \dots, \gamma_{tE})$ remains i.i.d. Gumbel distributed, identical to independent routing.

Key Contributions

1. H-CGA Architecture: A hierarchical copula layer that controls cross-token dependence in both directions (positive within-group, tunable negative cross-group) without modifying router logits, expert weights, or any per-token routing law.
2. Theoretical Invariance: The author proves (Theorem 1) that the entire hierarchy preserves the conditional distribution of every token's ordered Top-K list and mixture weights. Consequently, the conditional expected expert inclusion counts (Corollary 1) are identical to those of independent routing.
3. Coherence–Dispersion Trade-off Characterization: The paper characterizes the variance trade-off (Proposition 1):
   • Positive within-group coupling increases the variance of realized expert loads relative to independent routing (a cost of local coherence).
   • Nonnegative cross-group opposition reduces this variance relative to "flat" coupling (independent groups) at the same within-group strength, though it does not necessarily reduce variance below the independent routing baseline.
4. Routing-Only Adaptation: The author formulates an adaptation method where a small controller sets the dependence dials ($\rho_g, \alpha_{g,g'}$) based on frozen features. This controller is trained using a score-function estimator (REINFORCE), requiring only forward passes through the frozen base model and gradients confined to the controller.

Results and Pilot Study
The paper presents a small-scale pilot on a 15.8M parameter decoder-only MoE trained on TinyStories, with the base model frozen during evaluation.

• Invariance Checks: Synthetic tests confirmed that the empirical frequency of ordered Top-K lists and expert inclusion probabilities under H-CGA matched independent routing within Monte Carlo error margins, validating the theoretical invariance.
• Statistical Effects:
  • Increasing within-group coupling ($\rho$) significantly increased the Jaccard similarity of expert sets within windows and reduced the number of distinct experts used per window.
  • Increasing cross-group opposition ($\alpha$) significantly reduced the Jaccard similarity at group boundaries, confirming the ability to decouple group demands.
• Learning: A learned scalar controller (3 parameters) trained stably but showed no improvement in validation cross-entropy compared to the independent baseline. The learned coupling strengths were low, consistent with the paper's observation that marginal preservation limits the signal available for optimization when only per-token loss is used.
• Load Variance: The pilot observed minimal changes in aggregate load coefficient of variation (CV), noting that this summary metric does not directly test the conditional variance bounds proposed in Proposition 1.

Significance and Claims
The paper claims that H-CGA exposes a new degree of freedom in MoE routing: the joint dependence structure of routing noise. Its primary significance is theoretical and mechanistic:

• Safety: It provides a "safety guarantee" that modifying cross-token dependence does not alter the marginal routing behavior of any token, ensuring that the base model's learned preferences are untouched.
• Control: It offers two complementary "dials" to trade off local expert-set coherence against global load dispersion.
• Efficiency: It enables routing-only adaptation where the base model is evaluated only in the forward direction, avoiding the computational cost of updating expert weights.

Limitations and Modesty
The author is explicit about the scope and limitations:

• Layer-Local Guarantee: The invariance results are conditional on the hidden states and logits entering a specific routing layer. They do not guarantee end-to-end invariance in a multi-layer network, as coupling may alter downstream hidden states.
• No Task-Level Gains: The pilot does not establish that routing-only adaptation improves performance on downstream tasks. The author notes that marginal preservation implies that objectives depending solely on single-token distributions cannot identify the controller; useful signals must arise from joint effects (e.g., attention interactions), which may be weak.
• Variance vs. Overflow: While the method theoretically characterizes load variance, the pilot did not measure conditional load variance or capacity overflows directly. The paper does not claim to solve capacity overflow but suggests the opposition dial is a mechanism to address it.
• Overhead: The method requires generating and coordinating Gumbel perturbations for all experts, which may introduce computational overhead compared to standard routing.

In summary, the paper establishes a mathematically rigorous framework for controlling cross-token routing dependence in frozen MoEs, proving that coherence and load dispersion can be manipulated without altering per-token routing laws, while remaining cautious about the immediate practical benefits for task performance.