Quantile Balancing derives the next expert bias directly from a batch of raw sigmoid router scores. For token i, raw score vector si, current bias b, and selected count k, current routes are the Top-k entries of si+b. Mixture weights use only the selected raw scores:
pi,j=∑r∈Tisi,rsi,j.Use the (k+1)-th biased score as token cutoff αi. For m tokens and n experts, target load q=mk/n is integral. For each expert j, take the (q+1)-th largest value of si,j−αi, negate it, then subtract the mean across experts. The resulting centered bias applies to the next batch only. Return a dictionary with selected_experts, mixture_weights, expert_loads, and next_bias. Selected indices and loads use integer tensors; the other tensors preserve the router-score dtype and device.
Examples
Input: router_scores = [[0.9,0.1],[0.2,0.8]], current_bias = [0,0], selected_count = 1
Output: {"selected_experts": tensor([[0], [1]]), "mixture_weights": tensor([[1.0], [1.0]]), "expert_loads": tensor([1, 1]), "next_bias": tensor([0.0, 0.0])}
Explanation: Each token selects a different expert. Both loads already equal the target, so the centered next bias remains zero.
Input: router scores shape (4, 4), current bias shape (4,), selected count = 1
Output: {"selected_experts": tensor of shape (4, 1), "mixture_weights": tensor of shape (4, 1), "expert_loads": tensor of shape (4,), "next_bias": tensor of shape (4,)}
Input: router scores shape (4, 4), current bias shape (4,), selected count = 2
Output: {"selected_experts": tensor of shape (4, 2), "mixture_weights": tensor of shape (4, 2), "expert_loads": tensor of shape (4,), "next_bias": tensor of shape (4,)}
Hint 1
Use a stable descending sort so tied route scores have deterministic expert indices.
Hint 2
Gather mixture weights from raw router scores, not the bias-adjusted scores.
Hint 3
The next bias uses the expert-wise order statistic of raw-score margins against each token cutoff.
Requirements
- Return a dictionary with exactly selected_experts, mixture_weights, expert_loads, and next_bias.
- Selected indices use integer dtype, and loads count every selected assignment.
- Mixture weights sum to one within every token.
- The next bias has zero mean within numerical tolerance.
Constraints
- Router scores have shape (tokens, experts) and lie strictly between zero and one.
- Selected count is positive and smaller than the expert count.
- Total assignments are divisible by the expert count.
- Relevant biased scores and quantile margins contain no ties.
All steps in Kimi K3
Accepts: array
Accepts: array
Accepts: number
Quantile Balancing
Quantile Balancing derives the next expert bias directly from a batch of raw sigmoid router scores. For token i, raw score vector si, current bias b, and selected count k, current routes are the Top-k entries of si+b. Mixture weights use only the selected raw scores:
pi,j=∑r∈Tisi,rsi,j.Use the (k+1)-th biased score as token cutoff αi. For m tokens and n experts, target load q=mk/n is integral. For each expert j, take the (q+1)-th largest value of si,j−αi, negate it, then subtract the mean across experts. The resulting centered bias applies to the next batch only. Return a dictionary with selected_experts, mixture_weights, expert_loads, and next_bias. Selected indices and loads use integer tensors; the other tensors preserve the router-score dtype and device.
Examples
Input: router_scores = [[0.9,0.1],[0.2,0.8]], current_bias = [0,0], selected_count = 1
Output: {"selected_experts": tensor([[0], [1]]), "mixture_weights": tensor([[1.0], [1.0]]), "expert_loads": tensor([1, 1]), "next_bias": tensor([0.0, 0.0])}
Explanation: Each token selects a different expert. Both loads already equal the target, so the centered next bias remains zero.
Input: router scores shape (4, 4), current bias shape (4,), selected count = 1
Output: {"selected_experts": tensor of shape (4, 1), "mixture_weights": tensor of shape (4, 1), "expert_loads": tensor of shape (4,), "next_bias": tensor of shape (4,)}
Input: router scores shape (4, 4), current bias shape (4,), selected count = 2
Output: {"selected_experts": tensor of shape (4, 2), "mixture_weights": tensor of shape (4, 2), "expert_loads": tensor of shape (4,), "next_bias": tensor of shape (4,)}
Hint 1
Use a stable descending sort so tied route scores have deterministic expert indices.
Hint 2
Gather mixture weights from raw router scores, not the bias-adjusted scores.
Hint 3
The next bias uses the expert-wise order statistic of raw-score margins against each token cutoff.
Requirements
- Return a dictionary with exactly selected_experts, mixture_weights, expert_loads, and next_bias.
- Selected indices use integer dtype, and loads count every selected assignment.
- Mixture weights sum to one within every token.
- The next bias has zero mean within numerical tolerance.
Constraints
- Router scores have shape (tokens, experts) and lie strictly between zero and one.
- Selected count is positive and smaller than the expert count.
- Total assignments are divisible by the expert count.
- Relevant biased scores and quantile margins contain no ties.
All steps in Kimi K3
Accepts: array
Accepts: array
Accepts: number