How lucidrains/mixture-of-experts works
1. This repository is a small Python library, not an application. Its architecture is centered on one PyTorch implementation of Sparsely Gated Mixture of Experts, with packaging and release automation around it.
2. The main runtime dependency and execution environment is PyTorch. All compute happens in tensor form, and the public API is intended to be used directly inside a larger model rather than as a standalone service.
3. The core implementation lives in mixture_of_experts/mixture_of_experts.py. The package root mixture_of_experts/__init__.py exposes the public classes, primarily MoE and HeirarchicalMoE, so consumers import from the package namespace instead of the implementation file.
4. The primary data flow is: input tensor → gating network → sparse expert routing → expert computation → recombination into an output tensor of the same feature dimension. The README examples show inputs shaped like [batch, sequence, dim], and outputs preserve that shape.
5. MoE is the single-layer mixture-of-experts block. It supports configurable expert count, expert hidden size, activation choice, capacity factors for train/eval, and auxiliary load-balancing loss. That loss is part of the model output and is meant to be added to the caller’s training objective.
6. Expert routing is designed around top-2 gating behavior, with separate train and eval policies for whether a second expert contributes. This is important architecturally because it determines how tokens are dispatched across experts and how much computation is spent per token.
7. HeirarchicalMoE is the second major subsystem. It composes two gating levels, matching the GShard-style hierarchy described in the README. Conceptually, the first gate selects a coarse group, and the second gate selects a specific expert within that group, multiplying expert count without changing the overall external interface.
8. The expert computation itself is intentionally pluggable. The README shows that users can provide a custom experts module, meaning the MoE layer is split into a routing subsystem and an arbitrary expert-network subsystem. This is a key boundary for diagramming: gating/dispatch on one side, expert network on the other.
9. The library is distribution-friendly: setup.py defines packaging for installation via pip, and .github/workflows/python-publish.yml automates publishing to PyPI. So the repo has a clear build/release path, but no runtime backend, database, or queue infrastructure.
10. The repository’s external surface area is minimal: PyTorch for model execution, PyPI for distribution, and GitHub Actions for release automation. There are no signs of web endpoints, persistent storage, or service-to-service communication.
11. From an architecture-diagram perspective, the only meaningful internal components are:
• public package API
• MoE routing layer
• hierarchical MoE composition layer
• expert network implementation
• auxiliary balancing-loss path
• packaging/release pipeline
12. The design goal is capacity scaling: increase parameter count by adding experts while keeping per-token computation sparse and roughly constant. That principle is the main architectural driver behind the tensor routing and capacity constraints.
Main entry points: mixture_of_experts.py MoE block · python-publish.yml PyPI publish · setup.py Packaging · __init__.py Public API
Components 12
Core library
PyTorch-based mixture-of-experts implementation and public package surface.
- Public API · Package namespace that exposes the user-facing classes.
mixture_of_experts/__init__.py - MoE module · Core sparsely gated mixture-of-experts implementation and routing logic.
mixture_of_experts/mixture_of_experts.py - MoE block · Single-layer sparse expert block that preserves the input feature shape.
mixture_of_experts/mixture_of_experts.py - Hierarchical MoE · Two-level gating composition that expands expert capacity with coarse-to-fine routing.
mixture_of_experts/mixture_of_experts.py - Gating · Top-2 gating and capacity control that dispatches tokens to experts.
mixture_of_experts/mixture_of_experts.py - Dispatch · Sparse token-to-expert assignment and recombination path.
mixture_of_experts/mixture_of_experts.py - Experts · Custom or built-in expert network used for per-token computation.
mixture_of_experts/mixture_of_experts.py - Aux loss · Auxiliary loss used to encourage balanced expert utilization.
mixture_of_experts/mixture_of_experts.py - PyTorch · Tensor execution backend for model computation.
Release pipeline
Packaging and publish automation for distribution via PyPI.
- Packaging · Python package metadata and install configuration.
setup.py - PyPI publish · GitHub Actions workflow that automates package publishing.
.github/workflows/python-publish.yml - PyPI · External package registry used for release distribution.
Connections 13
- Public API → MoE module: re-exports · Public package namespace exposes the main model classes.
- MoE module → MoE block: defines · Implementation file contains the primary MoE layer.
- MoE module → Hierarchical MoE: defines · Implementation file also contains the hierarchical composition.
- MoE block → Gating: routes via · Tokens are scored and assigned through sparse gating.
- Gating → Dispatch: dispatches · Selected experts receive routed token subsets.
- Dispatch → Experts: invokes · Experts perform the per-token transformation.
- Dispatch → MoE block: recombines into · Expert outputs are merged back to the original tensor shape.
- MoE block → Aux loss: emits · Training objective includes auxiliary load-balancing loss.
- MoE block → PyTorch: built on · All compute is expressed as PyTorch tensor operations.
- Hierarchical MoE → Gating: stacks · Hierarchical variant composes two routing levels.
- Packaging → PyPI: packages for · Build metadata prepares the library for PyPI distribution.
- PyPI publish → PyPI: publishes to · Workflow automates package release to the registry.
- PyPI publish → Packaging: uses · Release workflow relies on package metadata and build configuration.