Mixture-of-Experts layer with sigmoid top-k routing.
Uses loss-free load balancing: expert bias terms are adjusted after each PPO update based on expert utilisation, without an auxiliary loss term.
References
- https://kexue.fm/archives/10699 (Math details)
- https://kexue.fm/archives/10757 (Auxiliary-Loss-Free Load Balancing)
- https://github.com/ambisinister/lossfreebalance (Loss-Free impl)
Classes:
| Name | Description |
|---|---|
MoELayer |
Mixture-of-experts layer with sigmoid top-k routing. |
MoELayer(input_size: int, output_size: int, num_experts: int = 4, k: int = 2, bias_update_speed: float = 0.001)
Bases: Module
Mixture-of-experts layer with sigmoid top-k routing.
Uses unbiased sigmoid gate values for expert weighting and adds a
learned bias to routing logits for load balancing. Bias corrections
are accumulated in :attr:bias_updates and applied externally
(e.g. by :class:PPO_MoE).
Initialize experts, gate network, and load-balancing bias.
Methods:
| Name | Description |
|---|---|
forward |
Compute routed expert outputs. |
forward(x: Tensor) -> torch.Tensor
Compute routed expert outputs.
Returns:
| Type | Description |
|---|---|
Tensor
|
torch.Tensor: Weighted combination of top-k expert outputs. |