Permutation-invariant hyperedge load balancing

Integrate this into the router-loss path of a structured MoE. Define a finite set E of admissible expert bundles; each bundle e contains the class strata that are allowed to share or jointly activate that bundle. For every minibatch, attach each token t to a stratum c_t, such as domain, language, modality, label bucket, or sequence-position group. The router produces logits z_{t,e} and probabilities p_{t,e}=softmax_e(z_t/ au). Compute expected counts q[v,e]=sum_t 1[c_t=v]p[t,e] for every v in e. Then compute the paper's all-pairs quantity L2[e] by looping over unordered pairs (u,v) in e and accumulating (q[u,e]-q[v,e])^2. Add lambda times the normalized sum over bundles to the task loss; use epsilon in the denominator to avoid instability for unused bundles. Pseudocode: logits=router(x); p=softmax(logits); q=zeros(|V|,|E|); for t,e: q[c[t],e]+=p[t,e]; balance=0; for e in E: for u<v in e: balance+=(q[u,e]-q[v,e])2/(q[:,e].sum()+eps)2; loss=task_loss+lambda*balance; backprop(loss). The formula and zero-at-equal-load property come from the paper; the soft expected-count relaxation and normalization are engineering adaptations. First test on a small 8-expert Switch-Transformer-like model on CIFAR-100 or a 100M-token language subset, comparing ordinary importance/load balancing against this penalty at equal lambda sweeps. Measure router entropy, per-class expert utilization variance, overflow/drop rate, task loss, and equal-FLOP validation accuracy. Success is lower worst-class routing imbalance and fewer capacity overflows without increasing task loss.