MoE Routing Policies: Top-1, Top-2, and Dynamic-k

What this part does

Policies turn calibrated router scores into an execution decision. Selection chooses the anchor or specialists, combination determines how approved specialist work contributes to the output, and capacity rules preserve a safe fallback when selected experts are unavailable.

Anchored top-1 is integrated through the package-to-model lifecycle. The top-2, adaptive, combination, and capacity implementations are validated and composable, but still require explicit assembly. Hierarchical, compatibility, and dynamic-k selection remain research primitives.

Where it fits

The policy boundary sits between evidence and execution:

1. a router produces calibrated scores in stable registry order;
2. a selection policy compares those scores with the fixed-zero anchor;
3. a capacity policy rewrites unavailable routes without adding new experts;
4. inference executes only the effective paths;
5. a combination policy mixes approved outputs or anchor-relative deltas.

The policies package does not extract features, fit scores, install packages, estimate serving capacity, run MLPs, or record telemetry. Those responsibilities belong to routers, calibration, registry, inference, and evaluation.

Implemented approaches

Core selection policies

AnchoredTop1Selection conceptually prepends the anchor’s zero score to the specialist score bank. A specialist runs only when it has strictly positive evidence. Exact zero ties select the anchor, while exact positive expert ties follow stable package append order.

AnchorVetoTop2Selection ranks the anchor and all specialists together. It approves a pair only when the top two candidates are both specialists, which is equivalent to requiring two strictly positive specialist scores. One strong specialist is not enough.

AdaptiveTop1Top2Selection ranks specialists without the anchor, then chooses route arity. Non-positive best evidence means anchor; one positive specialist means top-1; a second score strictly above the configured mixture threshold enables top-2.

Advanced selection policies

Hierarchical selection chooses one package against the anchor before selecting a local pair. Compatibility-constrained top-2 searches only explicit same-group pairs. Dynamic-k retains a confidence-qualified prefix up to a fixed per-token cap. These strategies expose deterministic, validated decisions, but package-driven model attachment or arbitrary-k dispatch is not yet implemented.

Combination policies

Score softmax converts the selected pair’s calibrated scores into non-negative weights that sum to one. Equal weighting and normalized positive evidence use the same pair-shaped contract, so they can be compared without changing selection or dispatch.

A learned mixture head deliberately separates pair choice from interpolation: selection scores still choose the experts, while a separately fitted frozen affine head produces their weights. The always-on-anchor mode instead executes the dense anchor for every token and adds zero, one, or two selected deltas. Parameter-space LoRA merging combines effective projection updates before one logical MLP execution.

Capacity policies

Capacity rules consume availability supplied by the caller; they do not estimate device load. Top-1 overflow returns directly to the anchor. For an approved top-2 route, both available specialists remain active, one survivor is promoted to weight one, and no survivors means exact anchor fallback. Tokens are never dropped and the policy never searches outside the pair that selection approved.

How it works

Preserve stable candidate meaning

Candidate zero always means the immutable anchor. Candidate k means expert position k - 1 in the stable expert_ids tuple. Decisions retain selected indices, selected scores, masks, and route arity so execution and evaluation can audit the same choice.

Separate choice from weighting

Top-2 selection returns an ordered pair and its original scores. Combination receives only the pair that selection approved. This prevents a weighting strategy from silently changing which experts execute and lets the dispatcher swap score softmax for a baseline or learned head without changing token grouping.

Apply capacity before execution

Availability is composed with the route before any candidate MLP runs. Unavailable experts therefore cannot execute accidentally. Effective slots remain pair-shaped for stable reconstruction, with zero used only as inactive anchor padding.

Mathematical intuition

Minimal usage

Select an anchored route, then apply explicit availability before dispatch:

from anchored_moe import AnchoredTop1Selection, OverflowToAnchorCapacity

decision = AnchoredTop1Selection()(scores, expert_ids=expert_ids)
effective = OverflowToAnchorCapacity()(
    decision,
    expert_availability=availability,
)
anchor_tokens = effective.effective_routes.anchor_mask

The input scores are already calibrated. Producing them, estimating availability, and executing the effective path remain separate steps.

Verified implementation details

decision = AnchoredTop1Selection()(
    scores,
    expert_ids=router.expert_ids_for_layer(1),
)

# Candidate zero is anchor; equal positive scores keep append order.
assert torch.equal(
    decision.selected_indices,
    torch.tensor([[1, 0, 2, 1]]),
)
assert decision.expert_ids == ("code:0", "math:0")

decision = AnchorVetoTop2Selection()(
    scores,
    expert_ids=router.expert_ids_for_layer(1),
)
mixed = ScoreSoftmaxCombination()(
    ranked_outputs,
    selected_scores=decision.selected_scores[
        decision.expert_pair_mask
    ],
)

assert mixed.output.shape == (2, hidden_states.shape[-1])
assert torch.allclose(mixed.weights.sum(-1), torch.ones(2))

result = dispatcher.dispatch(
    hidden_states,
    expert_scores=scores,
    expert_ids=("code:0", "math:0"),
    anchor_mlp=anchor,
    experts=(code, math),
    expert_availability=availability,
)

assert torch.equal(result.capacity.route_arity, torch.tensor([[2, 1, 0]]))
assert torch.equal(result.mixture_weights[0, 1], torch.tensor([1.0, 0.0]))
assert torch.equal(result.output[0, 2], hidden_states[0, 2])

Current limitations

• Package-driven model synchronization currently configures anchored top-1; top-2, adaptive, hierarchical, compatibility, and dynamic-k model configuration remain explicit.
• Capacity policies consume caller-provided availability. They do not measure queue depth, memory pressure, token quotas, or distributed load.
• Top-2 degradation searches only the selected pair; it does not try a third specialist when both members are unavailable.
• Hierarchical routing requires rectangular package topology and explicit package/local assembly. Compatibility routing requires a separately built immutable group index.
• Dynamic-k produces selection metadata only. Arbitrary-k sparse dispatch, combination, compatibility, and capacity are not implemented.
• Learned mixture weights require an offline fitting plan and are not attached through package manifests.
• Always-on deltas and parameter-space LoRA merging have explicit routed-layer correctness paths, but no package-driven LoRA materialization or optimized production kernels.
• Token-aligned auxiliary arguments, telemetry, multi-device execution, and distributed capacity coordination remain future work.

References

POL-001 through POL-005, COM-001 through COM-003, BAL-001, and BAL-002 are project-specific framework infrastructure rather than paper-derived training algorithms. Their exact tie, shape, fallback, and gradient contracts follow the linked source implementations and PyTorch operations.

The policy theory README contains the complete algorithms, tensor shapes, tests, and documented research differences. The architecture blueprint defines the fixed-zero anchor, reference inference policies, and capacity invariants shared across the framework.

Related pages

• Read routers for the calibrated score banks consumed here.
• Read experts for complete specialist and anchor-relative delta representations.
• Return to configuration and modeling for the current package-driven anchored top-1 synchronization boundary.
• Read inference for sparse token dispatch, reconstruction, and telemetry.
• Read calibration for score temperatures, thresholds, utility labels, and learned mixture fitting.