Degree-Capacity Regularizer for Sparse Routing

Apply this at the router output of a sparse-attention or MoE layer. Let M_{ve} be a soft assignment probability from token v to group or expert e, with top-k dispatch converted to a differentiable soft matrix during the forward-backward pass. Compute d(v)=sum_e M_{ve}; for hard top-k routing this is the number of selected groups per token, while for MoE it can instead be the number of routed interaction groups associated with token v. At each minibatch, form soft degree masks using sigmoids around thresholds r and 2r: b_v=sigmoid((r+0.5-d(v))/tau), c_v=sigmoid((d(v)-r-0.5)/tau)sigmoid((2r-0.5-d(v))/tau), and a_v=1-b_v-c_v. Replace |B| and |C| in the paper's inequality by sum_v b_v and sum_v c_v. Add the displayed capacity residual, a penalty on excessive A-degree, and optionally the pair-overlap penalty when groups represent attention subsets. Stop gradients through the masks in an initial ablation to avoid unstable threshold dynamics, then anneal tau from 1 to 0.1. The exact degree partition and coefficients come from the paper; the smooth masks and residual are engineering adaptations. First test a 4-layer, 8-expert top-2 MoE on WikiText-2 or CIFAR-100, comparing standard Switch-style load balancing with this regularizer at equal capacity. Log expert load variance, maximum token degree, dropped-token rate, router entropy, training loss, and throughput. Success means lower maximum load and fewer routing overflows without increasing perplexity or reducing tokens-per-second; a negative result is increased routing collapse or no improvement over standard load balancing.