adu2021 DynaMoE: Adaptive Expert Capacity Allocation in Mixture-of-Experts
Standard Mixture-of-Experts (MoE) architectures use fixed Top-K routing, where exactly K experts activate per token across all layers. This constraint wastes capacity on simple inputs and under-scales complex ones. DynaMoE relaxes both constraints: enabling dynamic per-token expert counts and layer-dependent expert capacity scheduling to match the representational structure of deep networks.
The core insight is that different layers have different optimal expert counts based on input diversity. Early layers process heterogeneous raw inputs requiring diverse expert specialization; deeper layers exhibit representational convergence where fewer experts suffice. By jointly optimizing token routing flexibility and layer capacity distribution, DynaMoE achieves superior efficiency-performance trade-offs.
Core Concept
DynaMoE makes two related innovations:
Token-Level Threshold Routing: Replace fixed Top-K with a percentile threshold where expert count varies per token based on input-specific routing logits. Simple inputs use 1-2 experts; complex inputs use more.
Layer-Wise Capacity Scheduling: Distribute total expert budget across layers using scheduling strategies (descending, ascending, pyramid, wave) to match how representational complexity evolves through the network.
This two-level decomposition aligns model capacity with actual computational needs rather than fixed architectural assumptions.
Architecture Overview
- Input: Token sequences with per-token routing scores from gating network
- Dynamic Threshold Extraction: For each token, compute expert activation threshold based on percentile of routing logits
- Expert Selection: Activate all experts exceeding the threshold (variable count per token)
- Capacity Budget Distribution: Assign expert counts to layers via scheduling strategy
- Output: Efficient sparse forward/backward passes with adaptive activation patterns
Implementation Steps
Step 1: Compute per-token routing logits and dynamic thresholds
Standard MoE gating networks output logits for each expert. Instead of selecting top-K, use percentile-based thresholding to determine active experts adaptively.
# Forward pass through MoE gating network
batch_size, seq_len, num_experts = tokens.shape[0], tokens.shape[1], 128
logits = gating_network(tokens) # Shape: [batch, seq, num_experts]
# Dynamic threshold routing: for each token, activate experts above percentile p
percentile = 50 # Can vary by task; higher = fewer active experts
thresholds = np.percentile(logits, percentile, axis=-1, keepdims=True)
# Binary mask for expert activation
active_mask = (logits >= thresholds).astype(np.float32)
num_active = active_mask.sum(axis=-1) # Varies per token
Step 2: Implement capacity scheduling strategy
Distribute expert budget across layers based on representational diversity patterns. Choose scheduling strategy aligned with task structure.
def layer_capacity_schedule(layer_idx, num_layers, total_experts, strategy='descending'):
"""Allocate expert count per layer."""
if strategy == 'descending':
# More experts early (diverse raw input), fewer deep (converged representations)
# Quadratic decay for smooth transition
ratio = (num_layers - layer_idx) / num_layers
return int(total_experts * (ratio ** 2))
elif strategy == 'ascending':
# Fewer experts early, more deep
ratio = layer_idx / num_layers
return int(total_experts * (ratio ** 2))
elif strategy == 'pyramid':
# Peak capacity in middle layers
mid = num_layers / 2
dist_to_mid = abs(layer_idx - mid) / mid
ratio = 1.0 - dist_to_mid
return int(total_experts * (ratio ** 2))
else: # uniform
return total_experts
# Allocate capacity per layer for 24-layer network with 256 total experts
expert_counts = [
layer_capacity_schedule(i, num_layers=24, total_experts=256, strategy='descending')
for i in range(24)
]
Step 3: Select experts and route tokens within capacity constraints
Route tokens to selected experts, handling overflow when activated experts exceed layer capacity.
def route_with_overflow_handling(logits, active_mask, layer_expert_capacity):
"""Route tokens respecting dynamic activation and layer capacity."""
batch_size, seq_len, num_experts = logits.shape
# Get indices of active experts per token
active_expert_indices = np.where(active_mask)
# Apply load-balancing: if too many experts active, keep top-scoring ones
expert_scores = logits * active_mask - 1e9 * (1 - active_mask)
# For each token, select up to layer_expert_capacity experts
selected_experts = []
for token_idx in range(batch_size * seq_len):
scores = expert_scores.reshape(batch_size * seq_len, -1)[token_idx]
active = np.where(active_mask.reshape(batch_size * seq_len, -1)[token_idx])[0]
# Prioritize by score within active set
active_scores = [(idx, scores[idx]) for idx in active]
active_scores.sort(key=lambda x: x[1], reverse=True)
selected = [idx for idx, _ in active_scores[:layer_expert_capacity]]
selected_experts.append(selected)
return selected_experts
Step 4: Execute sparse forward pass with routing weights
Compute token-to-expert assignments and weighted combination of expert outputs.
# Compute routing weights via softmax over active expert logits
expert_logits_masked = logits * active_mask - 1e9 * (1 - active_mask)
routing_weights = softmax(expert_logits_masked, axis=-1)
# Expert forward pass: dispatch tokens to selected experts
expert_outputs = []
for expert_idx in range(num_experts):
# Collect tokens assigned to this expert
token_mask = (active_mask[:, :, expert_idx] > 0).astype(np.float32)
expert_input = tokens * token_mask.reshape(batch_size, seq_len, 1)
# Expert forward pass (can be parallelized across experts)
expert_output = expert_networks[expert_idx](expert_input)
expert_outputs.append(expert_output)
# Weighted combination of expert outputs
combined_output = np.zeros_like(tokens)
for expert_idx in range(num_experts):
weight = routing_weights[:, :, expert_idx].reshape(batch_size, seq_len, 1)
combined_output += weight * expert_outputs[expert_idx]
Step 5: Backward pass with gradient accumulation across routing decisions
Support gradient flow through both gating network and expert networks to enable end-to-end training.
# Gradient accumulation through routing and expert selection
# (Simplified; real implementation uses automatic differentiation)
# Loss backprop to expert outputs
dL_dexpert_output = loss.backward()
# Gradient through weighted combination
for expert_idx in range(num_experts):
dL_drouting_weight = (dL_dexpert_output * expert_outputs[expert_idx]).sum()
dL_dlogits[:, :, expert_idx] = dL_drouting_weight * (routing_weights - routing_weights**2)
# Load balancing auxiliary loss to prevent expert collapse
auxiliary_loss = load_balance_loss(routing_weights, active_mask)
total_loss = task_loss + 0.01 * auxiliary_loss
Practical Guidance
Hyperparameter Selection:
- Percentile for threshold: 30-70 (higher = fewer active experts). Start at 50; adjust based on latency budgets.
- Capacity scheduling strategy: For image tasks use descending (early diversity dominates); for language tasks try ascending or uniform.
- Layer expert counts: Range from 64-256 depending on model scale; total budget ~4-8% of dense equivalent.
When to Use:
- Large-scale models (7B+ parameters) where efficiency gains compound
- Tasks with variable computational demands across samples (e.g., document length varies)
- Scenarios prioritizing inference latency and memory over peak training performance
When NOT to Use:
- Small models (<1B) where MoE overhead dominates benefits
- Tasks requiring uniform compute per sample (real-time control systems)
- Environments with strict memory constraints per device
Common Pitfalls:
- Imbalanced expert utilization: Some experts become idle. Add explicit load-balancing losses.
- Capacity mismatches: Layer capacity too low causes token overflow; too high wastes resources. Validate with profiling.
- Routing instability during training: Thresholds can jitter. Use exponential moving averages of logits for stability.
- Neglecting backward pass complexity: Gradient computation through dynamic routing can be expensive; use gradient checkpointing selectively.