# lucidrains/mixture-of-experts architecture

> Architecture diagram of the GitHub repository lucidrains/mixture-of-experts, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-05-12). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/lucidrains/mixture-of-experts
- Repository: https://github.com/lucidrains/mixture-of-experts

## Overview

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/mixture_of_experts.py`](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py) (MoE block), [`.github/workflows/python-publish.yml`](https://github.com/lucidrains/mixture-of-experts/blob/master/.github/workflows/python-publish.yml) (PyPI publish), [`setup.py`](https://github.com/lucidrains/mixture-of-experts/blob/master/setup.py) (Packaging), [`mixture_of_experts/__init__.py`](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/__init__.py) (Public API)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_core["Core library"]
  node_package_api["Public API<br/>package exports<br/>[__init__.py]"]
  node_moe_impl["MoE module<br/>model implementation"]
  node_moe_block["MoE block<br/>single-layer MoE"]
  node_hier_moe["Hierarchical MoE<br/>hierarchical model"]
  node_gating{{"Gating<br/>token router"}}
  node_dispatch["Dispatch<br/>sparse routing"]
  node_experts["Experts<br/>pluggable network"]
  node_aux_loss["Aux loss<br/>load balancing"]
  node_torch(("PyTorch<br/>runtime dependency"))
end

subgraph group_release["Release pipeline"]
  node_setup["Packaging<br/>setup script<br/>[setup.py]"]
  node_publish["PyPI publish<br/>release workflow<br/>[python-publish.yml]"]
  node_pypi(("PyPI<br/>distribution target"))
end

node_package_api -->|"re-exports"| node_moe_impl
node_moe_impl -->|"defines"| node_moe_block
node_moe_impl -->|"defines"| node_hier_moe
node_moe_block -->|"routes via"| node_gating
node_gating -->|"dispatches"| node_dispatch
node_dispatch -->|"invokes"| node_experts
node_dispatch -->|"recombines into"| node_moe_block
node_moe_block -->|"emits"| node_aux_loss
node_moe_block -->|"built on"| node_torch
node_hier_moe -->|"stacks"| node_gating
node_setup -->|"packages for"| node_pypi
node_publish -->|"publishes to"| node_pypi
node_publish -->|"uses"| node_setup

click node_package_api "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/__init__.py"
click node_moe_impl "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py"
click node_moe_block "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py"
click node_hier_moe "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py"
click node_gating "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py"
click node_dispatch "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py"
click node_experts "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py"
click node_aux_loss "https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py"
click node_setup "https://github.com/lucidrains/mixture-of-experts/blob/master/setup.py"
click node_publish "https://github.com/lucidrains/mixture-of-experts/blob/master/.github/workflows/python-publish.yml"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_package_api,node_moe_impl,node_moe_block,node_hier_moe,node_gating,node_dispatch,node_experts,node_aux_loss,node_torch toneBlue
class node_setup,node_publish,node_pypi toneAmber
```

## Components

### Core library

PyTorch-based mixture-of-experts implementation and public package surface.

- [Public API](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/__init__.py): package exports · Package namespace that exposes the user-facing classes. · `mixture_of_experts/__init__.py`
- [MoE module](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py): model implementation · Core sparsely gated mixture-of-experts implementation and routing logic. · `mixture_of_experts/mixture_of_experts.py`
- [MoE block](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py): single-layer MoE · Single-layer sparse expert block that preserves the input feature shape. · `mixture_of_experts/mixture_of_experts.py`
- [Hierarchical MoE](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py): hierarchical model · Two-level gating composition that expands expert capacity with coarse-to-fine routing. · `mixture_of_experts/mixture_of_experts.py`
- [Gating](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py): token router · Top-2 gating and capacity control that dispatches tokens to experts. · `mixture_of_experts/mixture_of_experts.py`
- [Dispatch](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py): sparse routing · Sparse token-to-expert assignment and recombination path. · `mixture_of_experts/mixture_of_experts.py`
- [Experts](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py): pluggable network · Custom or built-in expert network used for per-token computation. · `mixture_of_experts/mixture_of_experts.py`
- [Aux loss](https://github.com/lucidrains/mixture-of-experts/blob/master/mixture_of_experts/mixture_of_experts.py): load balancing · Auxiliary loss used to encourage balanced expert utilization. · `mixture_of_experts/mixture_of_experts.py`
- **PyTorch**: runtime dependency · Tensor execution backend for model computation.

### Release pipeline

Packaging and publish automation for distribution via PyPI.

- [Packaging](https://github.com/lucidrains/mixture-of-experts/blob/master/setup.py): setup script · Python package metadata and install configuration. · `setup.py`
- [PyPI publish](https://github.com/lucidrains/mixture-of-experts/blob/master/.github/workflows/python-publish.yml): release workflow · GitHub Actions workflow that automates package publishing. · `.github/workflows/python-publish.yml`
- **PyPI**: distribution target · External package registry used for release distribution.

## Connections

- 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.