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

arXiv:2607.28670 · cs.LG · Submitted 2026-08-17 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Hierarchical Copula-Gumbel-Top-K Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws".

Jane: The paper was written by Richard Yi Da Xu from Hong Kong Baptist University and TadReamk Limited.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Jane: We also have Lu with us today — senior AI researcher at Tsinghua.

Tom: We also have Meng with us today — lead engineer at a mysterious AI startup.

Jane: We also have Lalam with us today — the in-house Large Language Model.

Tom: Alright, let's get started.

Title and Authors: Tom: Welcome back to the show, everybody. Today we're digging into a paper with a title that's a mouthful: "Hierarchical Copula-Gumbel-Top-K Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws." Jane, I need you to translate that for our listeners before my brain melts.

Jane: Happy to, Tom. So imagine you've got a giant language model that doesn't run every expert on every word—it picks just a few experts per token to save compute. That's the "mixture-of-experts" part. The "Top-K routing" is how it picks the top few experts. And this paper is about controlling how those picks relate to each other across different tokens.

Tom: Right, and the key word in that title is "frozen." The model itself isn't being changed at all. The authors found a way to change how tokens coordinate their expert choices without touching a single weight in the base model.

Jane: Exactly. The author is Richard Yi Da Xu from Hong Kong Baptist University and a company called TadReamk Limited. And the core idea is kind of beautiful: each token's individual routing law—the probability it picks any given expert—stays exactly the same. But the joint behavior across tokens changes completely.

Tom: So it's like... every token still rolls the same dice, but the dice are now magnetically linked to each other?

Jane: That's actually a perfect way to put it. The dice are the same, but they're no longer independent. And that matters because right now, in most MoE models, every token rolls its dice completely on its own. That's the default nobody chose—it's just what happens.

Tom: And the paper shows you can do better. You can make related tokens—like words in the same sentence or code phrase—more likely to pick the same experts. That's the "positive coupling" direction. It creates coherence, which could mean better cache usage, fewer distinct experts touched per phrase.

Jane: But here's the twist. If you bunch tokens together, you create bursts of demand on specific experts. So the paper adds a second dial: negative coupling between different groups. If one group gets a random push toward an expert, its paired group gets the opposite push. That smooths out the load.

Tom: Two dials, both directions, and every token's individual routing law is provably unchanged. That's the headline. And it's all done through something called a copula, which is a fancy statistical tool for separating "what each variable does alone" from "how variables move together."

Jane: Right. The marginals stay fixed, the dependence changes. It's a whole new degree of freedom that nobody was touching before. And I have to say, the implications are pretty wild—this could work on any frozen MoE model out there.

Tom: I'm hooked already. Let's get into the actual mechanics of how they pull this off.

Summary and Core Results: Tom: So Jane, we've got the title decoded. Now let's talk about what the paper actually proves. Because it's not just an idea—they've got theorems.

Jane: They do. The central result is Theorem one and it's a guarantee: for every single token, the distribution of its ordered expert list, its selected set, and its mixture weights is identical to what you'd get under independent routing. That's a strong statement.

Tom: And it holds even with all the coupling machinery in place?

Jane: Yes, because of how they build the noise. Each token's noise vector is still i.i.d. Gumbel—that's the distribution you need for the standard Gumbel-Top-K trick. The coupling happens through shared latents, but each token's marginal noise distribution is untouched.

Tom: So the per-token law is exactly preserved. But what does that buy you in practice?

Jane: Well, there's a corollary that's really important: the expected load on each expert is also preserved. If you average over many routing decisions, each expert gets exactly the same expected number of tokens as before. No expert becomes systematically over- or under-used.

Tom: But the variance changes. That's the trade-off they characterize in Proposition one.

Jane: Exactly. Positive coupling within a group can only increase the variance of realized expert loads. That's the cost of coherence. But then the negative coupling between groups can only decrease that variance relative to flat coupling. So you've got two dials pulling in opposite directions, and the paper proves the direction of each effect.

Tom: And the expected loads stay the same in all schemes. So you're not sacrificing average behavior—you're just reshaping the fluctuations around it.

Jane: Right. And the construction itself is elegant. Within a group, each expert coordinate has a shared Gaussian latent plus a private noise per token. That shared latent creates the positive correlation. Then between paired groups, the latents are antithetically related—one group's push is the other group's counter-push.

Tom: And because the whole thing is built coordinate-by-coordinate, each token still sees independent Gumbel noise across experts. That's the trick that keeps the per-token law intact.

Jane: Exactly. It's a hierarchical copula—hence the title. And the math checks out. They even have a proof in the appendix using the association inequality, which is a classical result about how monotone functions of independent random variables behave.

Tom: So we've got a mechanism that's provably safe for individual tokens, provably changes joint behavior, and gives you a signed trade-off between coherence and load dispersion. That's a solid theoretical foundation. But I'm dying to know—does it actually work in practice?

Jane: That's exactly what the pilot study in Section five tries to answer. Let's bring in Lu and Meng to get their takes.

Improvements and Practical Implications: Tom: So we've got the theory. Now let's talk about what the paper actually did to test it. Lu, you've been quiet—what do you make of the experimental setup?

Lu: I think the pilot is deliberately modest, which I appreciate. They trained a small fifteen point eight-million-parameter model on TinyStories, then froze it completely. The key results are the routing statistics: with fixed coupling at ρ=zero point six, within-window Jaccard similarity—that's how often adjacent tokens pick the same experts—went from zero point two zero nine to zero point three one four. And distinct experts per window dropped from five point two seven seven to four point six zero five.

Jane: So tokens really are clustering onto the same experts. The mechanism works.

Lu: It does. And the validation cross-entropy barely moved—three point one two eight three two to three point one two eight one six. That's a tiny change, but it's not the point. The point is the routing behavior changed dramatically while the loss stayed essentially flat.

Meng: But hold on. From an engineering standpoint, I need to know: does this actually help anything? The load coefficient of variation barely moved in their table—zero point one five nine four six to zero point one five nine five zero. That's noise.

Tom: That's a fair pushback, Meng. The paper itself admits the aggregate load CV isn't a direct test of their variance claims. It's a finite-sample summary across experts.

Meng: Right, and they say the real test would be measuring capacity overflows or actual hardware savings. This pilot doesn't do that. So what's the practical path forward?

Lu: Well, the most exciting part to me is the controller idea in Section three point five. They propose a tiny trainable controller—as few as three parameters—that reads frozen features and sets the coupling strength. It's trained with a score-function estimator, which means the frozen base model is only ever evaluated forward. No backprop through the experts.

Meng: So you're telling me I can adapt routing behavior on a frozen model without touching the weights? That's like... a routing-only adaptation layer. That could be deployed on top of any existing MoE without retraining.

Lu: Exactly. And the paper is honest that the learned controller in their pilot didn't improve validation loss. But that's not surprising—the learning signal for joint routing behavior is weak when you're only optimizing per-token cross-entropy. The signal would come from joint effects, like tokens interacting through later attention layers.

Tom: So the mechanism is proven, but the application is still open. What would make this sing?

Lu: Imagine a model that needs to serve code. You want tokens from the same identifier to hit the same experts for cache locality. This gives you that control without changing the model's behavior on any single token. Or imagine a system where you know certain groups of tokens will arrive together—you can coordinate their routing to minimize expert switching.

Meng: And the negative coupling dial could help with load balancing in real time. If you see a burst coming, you could increase opposition between groups to smooth the demand. That's a systems-level control knob that didn't exist before.

Jane: It's like having a volume knob for coordination. Turn it up for coherence, turn it down for balance, and the model's individual decisions never change.

Conclusion: Tom: Alright, we've covered a lot of ground on "Hierarchical Copula-Gumbel-Top-K Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws." Let's pull it together.

Jane: The core idea is that every token's routing law—its individual probability of picking each expert—can stay exactly fixed while the joint behavior across tokens changes dramatically. That's the copula insight: marginals and dependence are separate things.

Tom: And the paper proves it. Theorem one guarantees per-token invariance. Corollary one guarantees expected loads are preserved. Proposition one signs the trade-off: positive coupling increases load variance, negative coupling decreases it relative to flat coupling.

Lu: The construction is elegant—shared Gaussian latents within groups, antithetic latents between paired groups, all built coordinate-by-coordinate so each token still sees independent Gumbel noise. The math is clean.

Meng: And the practical potential is real, even if the pilot is small. A controller that sets these dials on a frozen model, trained with a score-function estimator, could give systems-level control over routing behavior without any base-model retraining.

Jane: The pilot validates the mechanism—routing statistics change as predicted, per-token laws hold to within Monte Carlo error. But it doesn't yet prove task-level gains. That's the honest limitation.

Tom: So where does this leave us? We've got a new degree of freedom in MoE routing that nobody was exploiting. It's provably safe for individual tokens, it gives you two complementary dials, and it works on frozen models. The open questions are about real-world impact: does coherence actually speed up inference? Does opposition actually prevent capacity overflow?

Lu: Those are exactly the experiments that should come next. On a larger pretrained model, with capacity-aware measurements, I think we'd see real benefits.

Meng: I'd want to see it deployed in a production MoE serving stack. That's where the rubber meets the road.

Jane: And that's what makes this paper exciting. It's not claiming to solve everything—it's opening a door. A new dial on a frozen model, with proofs that you're not breaking anything. That's a rare combination.

Tom: Well said, Jane. We'll be watching for follow-ups on this one. Thanks to Lu and Meng for joining us today. And to our listeners—if you're working on MoE systems, this paper is worth your time. We'll see you next episode.

Jane: Take care, everyone.

Richard Yi Da Xu

Hong Kong Baptist University · TadReamk Limited

cs.LG

Submitted: 2026-08-17

Updated: 2026-08-18

License: http://creativecommons.org/publicdomain/zero/1.0/

Importance score: 60/100

Key concepts

Mixture-of-Experts (MoE): A model architecture that saves compute by selecting only a few specific experts to process each token instead of running every expert in the model for every word.
Copula: A statistical tool used to separate the individual behavior of variables from their joint behavior. It allows for changing how tokens coordinate their expert choices without altering each token's individual routing probabilities.
Coupling: The coordination of routing decisions between tokens. Positive coupling makes related tokens more likely to pick the same experts to improve coherence, while negative coupling between groups helps balance the load by smoothing out expert demand.

Terminology

Copula — 44× in this paper · explained in 2 episodes
Adaptation — 16× in this paper · explained in 1 episode
Frozen model — 10× in this paper · explained in 3 episodes
Fine-Tuning — 5× in this paper · explained in 5 episodes
Core — 4× in this paper · explained in 2 episodes
Mixture-of-Experts (MoE) — 3× in this paper · explained in 7 episodes
Accuracy — 2× in this paper · explained in 1 episode
Benchmark — 2× in this paper · explained in 2 episodes
Mixture of Experts — 2× in this paper · explained in 2 episodes
Clustering — 1× in this paper · explained in 2 episodes
Domain Shift — 1× in this paper · explained in 4 episodes
Embeddings — 1× in this paper · explained in 2 episodes
Parameter-Efficient Fine-Tuning — 1× in this paper · explained in 2 episodes
Reinforcement Learning — 1× in this paper · explained in 4 episodes
Score Function — 1× in this paper · explained in 4 episodes
Sparsity — 1× in this paper · explained in 4 episodes
Transformer — 1× in this paper · explained in 2 episodes

Summary

Summary

This paper introduces Hierarchical Copula-Gumbel-Top-K (H-CGA), a method for controlling the joint distribution of routing choices across tokens in a frozen mixture-of-experts (MoE) model, while holding each individual token’s marginal routing law exactly fixed. The core problem is stated as: “Under independent stochastic routing, the joint routing distribution of a batch is a product measure over tokens—a default nobody chose.” The paper asks “which joint distributions over the routing choices of different tokens are reachable while every individual token’s complete routing law is held exactly fixed.”

The construction is based on Sklar’s theorem, which separates a joint distribution into marginals and a copula (dependence structure). The paper states: “by Sklar’s theorem (Sklar, 1959), every joint distribution factors as F (x1,..., xn) = C(F1 (x1),..., Fn (xn)) with C a copula on [0, 1]n, uniquely when the marginals are continuous. Independence is exactly the copula density c ≡ 1; changing C never touches the Fi.”

The method operates on the Gumbel noise of a stochastic Gumbel-Top-K router. For each token t and expert e, the router produces logits lte; standard stochastic routing adds independent Gumbel noises and takes the top K perturbed logits. H-CGA replaces the independent noise with a hierarchical copula construction:

Within-group positive coupling: Tokens are partitioned into disjoint groups (e.g., fixed windows of m adjacent tokens). For each group g and expert coordinate e, a shared latent ζge ∼ N(0,1) and private noises ϵte ∼ N(0,1) are drawn. The construction is: yte = √ρg ζge + √(1−ρg) ϵte, then ute = Φ(yte), and γte = −log[−log(ute)]. This yields a standard Gumbel marginal for each token while inducing positive correlation across tokens in the same group at each expert coordinate. The correlation strength ρg = ag ρmax is set by a controller, with ρmax < 1 to keep the construction nondegenerate.
Cross-group negative coupling (antithetic): Groups are matched into disjoint pairs (g, g′). For each pair and expert coordinate e, the shared latents are related by: ζg′e = −αg,g′ ζge + √(1−α2g,g′) ηg′e, where ηg′e ∼ N(0,1) is independent, and αg,g′ ∈ [0,1] is a tunable opposition strength. At αg,g′ = 1, the groups receive fully opposite shared random pushes; at αg,g′ = 0, they are independent.

The paper proves the central invariance result:

Theorem 1 (Per-token Top-K routing-law invariance): Condition on all frozen hidden states, router logits, group memberships, and controller outputs. Suppose H-CGA uses the construction above and that the copula draws are independent across expert coordinates. For every token t, (γt1,..., γtE) is distributed as i.i.d. Gumbel(0,1). Consequently, the ordered list (rt1,..., rtK), selected set St, and weights (wt1,..., wtE) have exactly the same conditional distribution as under the independent Gumbel-Top-K router.

Corollary 1 (Expected inclusion load is invariant): Let Ne = Σt 1[e ∈ St] be the number of tokens that include expert e. Under the conditions of Theorem 1, E[Ne lt t, ag, ρg g] = Σt P(e ∈ St lt; independent Gumbel-Top-K). The right-hand side does not depend on the coupling strengths, groups, or fixed pairing once the layer’s incoming logits are fixed.

Corollary 2 (Hierarchical routing-law invariance): Under the paired-latent rule, for any fixed αg,g′ ∈ [0,1], the conclusion of Theorem 1 and Corollary 1 continue to hold for every token.

The paper also characterizes the trade-off between coherence and load dispersion:

Proposition 1 (Coherence–dispersion trade-off): Condition on all logits, group memberships, pairings, within-group coupling strengths, and opposition strengths αg,g′ ∈ [0,1], and fix an expert e. Write Xg = Σt∈g 1[e ∈ St] and Ne = Σg Xg. Then the conditional expected loads E[Ne] are identical under independent routing, flat coupling (independent ζge across groups), and the tunable paired construction, and: (i) under flat coupling with any strengths ρg ≥ 0, Var(Ne) is at least its value under independent routing; (ii) under paired coupling with the same ρg and any αg,g′ ∈ [0,1], Var(Ne) is at most its value under flat coupling.

The proof uses the association inequality of Esary et al. (1967). Part (i) states that positive within-group coupling cannot make an expert’s conditional load less variable than under independent routing. Part (ii) states that adding cross-group opposition cannot make that load more variable than flat coupling at the same strengths.

The paper also formulates routing-only adaptation as an application. A small controller ϕ reads frozen, pre-routing features sg (e.g., mean hidden state, mean gate entropy, within-group gate similarity) and sets the dependence dials: ag = σ(ϕ(sg)) ∈ [0,1], ρg = ag ρmax, and αg,g′ = σ(ϕpair(stopgrad(sg), stopgrad(sg′))) ∈ [0,1]. The controller is trained with a score-function estimator: ∇ϕ E[L x] = E[(L − b(x)) ∇ϕ log pϕ(Y x) x], where b(x) is a stop-gradient baseline. The frozen base model is evaluated only in the forward direction; gradients are confined to the controller. The paper notes: “The stopgrad(·) operation guarantees that the controller’s inputs provide no gradient path into the backbone or router, so the routing-only property would survive even if the base parameters were trainable.”

The paper includes an initial small-scale pilot. A 15.8M-parameter, six-layer decoder-only Top-2 MoE language model was trained on 10M TinyStories tokens, with eight experts in each of three MoE layers and sequence length 256. Groups are fixed, non-overlapping windows of m = 4 adjacent tokens. The main fixed-coupling comparison uses ρ = 0.6 and α = 0. The learned-controller experiment learns one constant ρ per MoE layer (three parameters total) using the score-function route, ρmax = 0.95, and three seeds.

Key pilot results: In a synthetic fixed-logit test with four tokens, 60,000 draws, and ρ = 0.8, the largest difference between the empirical frequency of an ordered Top-2 list under independent and copula routing was 0.0035; the largest difference in an expert-inclusion frequency was 2.2 × 10−4. Adjacent-token Top-2 Jaccard overlap rose from 0.53 to 0.77 in this synthetic setting. On the full frozen model, fixed copula (ρ = 0.6) raised within-window Jaccard from 0.209 to 0.314 and reduced distinct experts per window from 5.277 to 4.605, with validation cross-entropy changing from 3.12832 to 3.12816. The learned scalar controller achieved validation CE of 3.12837 ± 0.00003, within-window Jaccard of 0.228 ± 0.004, and distinct experts per window of 5.132 ± 0.027. A Router-LoRA reference (1,584 trainable parameters) achieved validation CE of 3.12439 ± 0.00034, within-window Jaccard of 0.208 ± 0.001, and distinct experts per window of 5.282 ± 0.005. The paper explicitly cautions: “The pilot is not a task-adaptation benchmark” and “It is not evidence that routing-only adaptation improves a pretrained MoE on a downstream task.”

The paper also tests the higher-level α dial. On identical synthetic logits at ρ = 0.6, paired-boundary Jaccard overlap decreases from 0.389 at α = 0 to 0.303 at α = 0.5 and 0.218 at α = 1, while within-window overlap remains approximately constant (0.601, 0.601, and 0.600). An α = 1 law check gives maximum ordered-list and inclusion-frequency deviations of 0.0028 and 0.0034. On the full frozen model, a single-seed evaluation-only sweep over ρ ∈ 0.6, 0.9 and α ∈ 0, 0.5, 1 lowers paired-window-boundary Jaccard from 0.210 to 0.170 to 0.137 at ρ = 0.6, and from 0.210 to 0.152 to 0.110 at ρ = 0.9, while within-window Jaccard is unchanged to within 3 × 10−4 and validation cross-entropy varies by at most 3.1 × 10−4 across all six cells.

The paper’s stated contributions are: (1) introducing H-CGA, a hierarchical copula layer over the routing noise of a frozen stochastic Gumbel-Top-K MoE that controls cross-token dependence in both directions without touching router logits, experts, or any per-token routing law; (2) proving full per-token routing-law invariance for the entire hierarchy (Theorem 1, Corollaries 1 and 2); (3) characterizing the coherence–dispersion trade-off (Proposition 1); and (4) formulating routing-only adaptation as an application with a score-function estimator.

The paper explicitly lists limitations: H-CGA is exactly plug-compatible only with a stochastic Gumbel-Top-K base router; capacity clipping, token dropping, and expert-choice allocation act after sampling and are outside the invariance results; Proposition 1 signs but does not quantify either variance effect; the pilot does not measure conditional load variance, capacity overflows, or a systems benefit; and marginal preservation is a safety and identifiability property, not an accuracy theorem. The paper concludes: “These guarantees are conditional and routing-layer-local, and do not by themselves render a multi-layer MoE end-to-end invariant.”

Improvements for AI systems

Based on the paper, here are the specific improvements I can make to an AI system, along with what the improved system can do:

What I implement:

I add a hierarchical copula layer (H-CGA) between the frozen router’s logits and the Gumbel-Top-K sampling step. This layer does not modify any per-token routing law—it only changes the joint distribution of routing noise across tokens.

What the improved system can do:

For a group of related tokens (e.g., tokens in the same code phrase, entity mention, or sentence), the system can increase the probability that they select the same experts, without changing any individual token’s expert-selection probabilities.
This reduces the number of distinct experts touched per group (e.g., from 6 to 3 in a Top-2 setting), improving expert locality and cache-friendliness in inference systems that benefit from batch-level expert reuse.
The system can do this with zero change to the base model’s weights, router logits, or per-token mixture weights.

The improved system is a frozen MoE language model with a controllable routing-dependence layer. It can:

Increase local expert coherence (fewer distinct experts per group) without changing any token’s routing law.
Decrease load variance via cross-group opposition, reducing capacity overflows.
Adapt to new tasks with a tiny trainable controller, using only forward passes through the base model.
Guarantee per-token routing-law invariance at the layer level.
Quantify the coherence–dispersion trade-off and set dials accordingly.
Operate as a drop-in addition to existing stochastic Gumbel-Top-K MoEs, with no changes to the base weights or router logits.

These improvements are directly implementable from the paper’s construction (Sections 3.1–3.5), and the pilot (Section 5) confirms the mechanism works in practice, though it does not yet claim task-level gains.

Abstract

A stochastic Gumbel-Top- K router defines, for every token of a mixture-of-experts (MoE) model, a routing law: a distribution over ordered expert lists and mixture weights. We ask which joint distributions over the routing choices of different tokens are reachable while every individual token's complete routing law is held exactly fixed. We give a two-sided construction, Hierarchical Copula-Gumbel-Top- K. Within a group of related tokens, an exchangeable Gaussian copula positively correlates the Gumbel perturbations at each expert coordinate, which can increase within-group expert-set coherence. Across disjoint pairs of groups, a tunable antithetic construction introduces a selectable amount of negative dependence. We prove that both operations leave each token's ordered Top- K sample, mixture weights, and inclusion probabilities identical in distribution to independent routing at a routing layer conditioned on its pre-routing logits; conditional expected expert traffic is preserved as a consequence. We characterize the resulting trade-off: positive within-group coupling can only inflate the variance of realized expert loads relative to independent routing, while nonnegative cross-group opposition can only reduce it relative to flat coupling at the same within-group strength. Coherence and load dispersion are thus controlled by two complementary dependence dials on the invariance constraint surface. Because the base model is untouched, the dials can be driven by a small controller over frozen features, trainable with a score-function estimator: the frozen network is evaluated only in the forward direction, and gradients are confined to the controller. An initial small-scale pilot validates the mechanism and the training route, but does not establish task-level fine-tuning gains.