Valuation-Threshold Congestion Router
Implementation & benchmark of arXiv:2607.03625 — Congestion Games with Heterogeneous Valuations: An Optimal Transport Approach
Source paper: Congestion Games with Heterogeneous Valuations: An Optimal Transport Approach arXiv:2607.03625 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Give every token a valuation vector over experts and route it to the expert maximizing utility after subtracting a dynamically updated congestion price. The router is a neural analogue of the paper's measurable valuation-space partition: tokens with different valuations are assigned to different experts, while expert prices discourage overload without requiring a generic global load-balancing penalty.
Formulas
Mathematical statement
The paper forms the joint valuation measure space as the disjoint union V := disjoint_union_i V_i, with measure mu(disjoint_union_i A_i) = sum_i mu_i(A_i). For an agent with valuation v in V_i, destination j in D_i, and congestion state tau_i = (tau_i,j), the choice region is V_i,j(tau_i) := {v in V_i : j belongs to argmax over j-hat of U_i(v_j-hat, tau_i,j-hat)}. Here v_j-hat is the token's valuation for destination or expert j-hat, U_i is utility as a function of valuation and congestion, and tau_i,j-hat is the congestion price or state of that destination. The paper argues that pairwise indifference sets are null when U_i is strictly monotone in valuation and the valuation measure mu_i is absolutely continuous, so the regions form an almost-everywhere partition. We use this partition as a hard router and update congestion prices from observed expert loads using a projected dual ascent rule.
Implementation notes
Integrate this at the router logits of a sparse MoE Transformer. For each token representation h_b, compute a valuation vector v_b = W_v h_b in R^E, where E is the number of experts. Maintain one nonnegative congestion price tau_j per expert, shared across the minibatch and updated with an exponential moving average of recent loads. Use this pseudocode: v = Wv @ h; score[b,j] = v[b,j] - tau[j]; route[b] = argmax_j score[b,j]; dispatch tokens to selected experts up to capacity C_j; compute load[j] = number of assigned tokens; update load_ema = beta * load_ema + (1-beta) * load; update tau[j] = relu(tau[j] + eta_tau * (load_ema[j] / C_j - 1)). In the forward pass use the one-hot argmax route. During training use a straight-through estimator whose backward pass substitutes the gradient of softmax(score divided by temperature T), or a softmax over the top two experts. Use U(v_j,tau_j) = v_j - lambda * tau_j, with lambda either fixed to one or tuned over a small grid. Estimate loads directly from each batch; the first implementation does not require a valuation-space density estimate or an optimal-transport solver. Compare against top-1 routing with Switch-style auxiliary load balancing on a 12-layer, 8-expert Transformer trained on a small language-model corpus. Measure validation perplexity, overflowed or dropped-token fraction, expert-load coefficient of variation, router assignment churn, and training throughput at equal FLOPs. The main falsifiable target is fewer overflowed tokens and lower load variance at the same auxiliary-loss weight, or equal load variance with less routing regularization and improved perplexity.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.