Skip to content

Tutorial

Install it, add the block to a model, confirm its loss reaches the optimiser, then measure it properly.

pip install esmoe

What the block does

The block sends each image to a few of several convolutional experts and sums what they return. Each expert is a depthwise-separable convolution with its own kernel size (3, 5, 7, 9, ...), so the experts differ in receptive field rather than merely in weights. A lightweight router scores them per image, the top-k are kept and renormalised, and the block returns their weighted sum. The block preserves channel count, which is what lets a stock parse_model size it.

flowchart LR
    X["feature map<br/>(n, c, h, w)"] --> R["router<br/>pool -> linear -> SiLU -> linear"]
    R --> P["softmax over experts"]
    P --> T["top-k, renormalised"]
    X --> E1["expert k=3"]
    X --> E2["expert k=5"]
    X --> E3["expert k=7"]
    X --> E4["expert k=9"]
    T --> S(("weighted sum"))
    E1 --> S
    E2 --> S
    E3 --> S
    E4 --> S
    S --> Y["output<br/>(n, c, h, w)"]
    P -.-> A["auxiliary loss"]
    T -.-> A
    A -.-> L["training loss"]

The default balancing objective is the Switch-Transformer load-balancing form; times its weight, it is the auxiliary loss:

\[ \mathcal{L}_{\text{aux}} = E \sum_{i=1}^{E} \bar{p}_i \cdot f_i , \]

where \(E\) is the expert count, \(\bar{p}_i\) the mean routing probability of expert \(i\) over the batch and \(f_i\) the fraction of samples that actually activated it. It is minimised when routing mass and realised load are spread evenly. Measured, what it prevents is an expert dying, not the dispatch concentrating (judgment lines, round six).

The three calls

inject_esmoe makes the block nameable in a config, graft puts it there and fixes the layer references, attach_aux_loss wires the auxiliary loss into training. equip does all four steps at once: register, graft, build, wire.

import esmoe

model = esmoe.equip("yolo11n.yaml", weight=0.01)
model.train(data="coco8.yaml", epochs=3, imgsz=320)

Take them apart when you need to:

esmoe.inject_esmoe()
esmoe.graft("yolov8n.yaml", out="v8-esmoe.yaml", at=[4, 6])
model = YOLO("v8-esmoe.yaml")
esmoe.attach_aux_loss(model, weight=0.01)

or from the shell:

esmoe graft yolo11n.yaml -o yolo11n-esmoe.yaml -e 4 -k 2 --at 4,6

Written by hand, the grafted layer is one line:

[-1, 1, ESMoE, [4, 2]]   # num_experts, top_k

Grafting and renumbering

A YOLO config addresses earlier layers by absolute index:

- [[-1, 12], 1, Concat, [1]]

Insert a layer at position 10 and every reference at or past 10 now points one layer too early. The model still builds, still trains, and is quietly wrong. graft therefore renumbers every reference that sits at or after each insertion point, and a unit test compares the rewritten head against the original one reference by reference.

Renumbering moves references; it does not retarget them. A head branch that names the old backbone end by index - YOLOv8's P5 lateral [-1, 9] Concat does - therefore keeps reading SPPF after the insertion, and the block reaches P5 only through the top-down path. To have the block take over every consumer of the backbone end, pass rewire=True (--rewire on the CLI):

esmoe.graft("yolov8n.yaml", out="v8-esmoe.yaml", rewire=True)

It is off by default to keep existing run records comparable. Under one budget (YOLOv8n, imgsz 800, 120 epochs, three seeds) the default wiring gives +0.0025 mAP50 (2/3 wins) with APl βˆ’0.0104 (0/3); rewire gives +0.0036 mAP50 (3/3 wins) with APl +0.0063 (2/3) - on v8n, bypassing the P5 lateral is where the large-object loss came from. That mechanism does not pin down across generations: on 26n the default wiring consistently loses small objects instead (APs βˆ’0.0045, 0/3) and on 12n the direction is unstable; rewire pulls 12n and 26n back to parity and trails the default on v5n, v9t, v10n and 11n. Seven-generation verdicts: judgment lines.

Checking the auxiliary loss

A configuration key named aux_loss proves nothing. What proves it:

  1. results.csv gains an esmoe_aux column, non-zero and moving;
  2. a unit test asserts total(with aux) == total(without) + aux * batch_size;
  3. the same test asserts the router receives gradient.

If you want the number yourself in a custom loop:

esmoe.clear_aux_loss()
task_loss = criterion(model(images), targets)
aux = esmoe.collect_aux_loss(model)
(task_loss + 0.01 * aux).backward()

collect_aux_loss reads without clearing: it sums the value each block last published to the registry. Call clear_aux_loss() before each forward, so a block that does not run in a step leaves no value from the step before; the loss patch attach_aux_loss installs does this itself.

Running a comparison

uv run python scripts/capture_env.py             # freeze versions and hardware into results/env/
EPOCHS=20 FRACTION=1.0 SEEDS="0 1 2" uv run bash scripts/sweep.sh
uv run python scripts/report.py                  # results/summary.md

Each run writes one JSON record: model config, dataset and fraction, hardware, budget, seed, metrics, artifact path, status, limitation. report.py groups arms by backbone, block config and budget, never averaging two different budgets into one row, then prints per-seed paired deltas against the baseline of the same seed.

Read the paired table, not the two means. Per-arm standard deviations overlap in this kind of experiment; what carries the claim is that the same seed, same data and same schedule moved in the same direction three times. At the selection budget (full VisDrone, 640 px, 20 epochs) the default configuration wins 3/3 seeds by +0.0021 mAP50, and one of those seeds is nearly a tie; at the protocol budget (800 px, 120 epochs) YOLOv8n gives +0.0025, 2/3. A small, consistent effect, not a reliable per-run improvement. Experiments has the full measurements.

Extending

Experts and the balancing objective are plain callables:

class ThinExpert(nn.Sequential):
    def __init__(self, c1, c2, k):
        super().__init__(nn.Conv2d(c1, c2, k, 1, k // 2, groups=c1), nn.SiLU())

def entropy_balance(probs, gate):
    return -(probs * probs.clamp_min(1e-9).log()).sum(dim=1).mean()

block = esmoe.ESMoE(num_experts=3, top_k=2, expert=ThinExpert, balance=entropy_balance)

esmoe.blocks(model) iterates every block in a model, which is how the collector and the tests find them.

Built-in balancing objectives

The default is Switch (switch_balance), as in every earlier release and in the default arms under results/:

objective formula reads
switch_balance (default) E * sum(p_i f_i) mean probabilities x realised load
gshard_balance N * sum(usage_i^2) the gated weights (what upstream's ES_MOE reads)
master_balance (1/E) * sum((mu_i - 1/E)^2) the gated weights (the paper's eq. 13; an affine map of the row above)
gshard_probs_balance N * sum(usage_i^2) the raw probabilities, to isolate which tensor is read

What separates them is not a coefficient but what they read. Where the mean probabilities are near uniform and the top-k dispatch concentrates on a few experts, the two that read the probabilities sit at fixed values (k for switch, 1 for gshard_probs) and cannot tell that apart, while the two that read the gate (gshard, master) can. The measured checkpoints are close to that shape: routing entropy runs at 90% to 100% of its maximum while the leading expert's top-1 share runs from 0.47 to 1.00. Upstream's code and its paper agree here, differing only by an affine map (L_paper = (L_upstream - 1) / E^2).

Telling a collapse apart is not the same as pushing against it. The gate holds only the scores inside the top-k subset, so an objective that reads the gate has exactly zero gradient for an expert outside the top-k: once an expert drops out of the top-k on every image, the term can no longer reach it. Measured, the gate-reading objectives lost an expert in five of six checkpoints, and in all six same-configuration B checkpoints trained with upstream's recipe, while Switch lost none in 66, which is why Switch is the default (judgment lines, rounds six to eight). To use upstream's objective instead:

esmoe.equip("yolo11n.yaml", balance="gshard")

A custom objective goes into the config too: define the function at module level in an importable module and pass it to equip or graft. The config stores module:qualname, and the trainer and every DDP worker import the same function back from that name when they rebuild the model. A lambda, a nested function or a function defined in __main__ cannot be imported back by name and is refused when grafting. Custom experts work the same way (expert=MyExpert). The measurements are on Experiments and Judgment lines.

Matching upstream

Upstream's ES_MOE differs from this package's defaults in the first five rows below: four affect training, and pruning affects inference only. They all go through equip, and a run compared against upstream adds its training recipe, recipe="upstream" (see API):

model = esmoe.equip(
    "yolo11n.yaml",
    at="backbone_stages",     # one block per stage, four in all
    balance="gshard",         # the objective that reads the gate
    out_norm=True,            # BatchNorm + SiLU after the weighted sum
    dense_training=True,      # every expert runs while training
    dynamic_threshold=0.4,    # pruning at inference
    recipe="upstream",        # what upstream's trainer does to a routed model
)
item upstream / paper this package's default how to turn it on
balancing objective N*sum(u^2) on the gate Switch (switch_balance) balance="gshard"
blocks one per backbone stage, four in all one, at the backbone end at="backbone_stages"
output norm BatchNorm + SiLU (the paper's eq. 2 Norm) none out_norm=True
training forward every expert runs, unrouted ones weighted zero unrouted experts skipped dense_training=True
inference pruning dynamic_threshold=0.4 drops low-share experts, keeps the leader, renormalises 0.0, nothing dropped dynamic_threshold=0.4
sparse inference use_sparse_inference=True same sparse_inference=False runs all

The block matches upstream's remaining parameters too: out_channels (the input's width by default), top_k=None meaning every expert, even kernels stepped down to odd and capped at max_kernel_size so a pruned checkpoint's kernels reload, and the same validation of num_experts, reduction, dynamic_threshold and max_kernel_size -- refused at construction rather than halfway through a run.

out_norm, dense_training and dynamic_threshold are off by default so the runs of this package's recipe in results/ still reproduce; the first two each have an arm on the judgment lines page. On the esmoe graft command line: --at, --balance, --out-norm, --dense-training; the pruning threshold goes through equip or the config.

These settings have to travel in the config; setting them on the blocks afterwards does not work. The trainer rebuilds the model from model.yaml, and anything set after YOLO(cfg) returns disappears with the instance that is discarded -- no error, no trace. equip and graft write them into the config, so a rebuild keeps them; ESMoE.configure(...) is for the paths that never reach a trainer -- inference, export, unit tests. To read back what was in force from any checkpoint: uv run python scripts/blockspec.py.

Things to know

  • Channels are inferred on the first forward: script, export or load a state_dict before any forward and the block has no expert weights to load into yet.
  • attach_aux_loss keeps its weight and recipe per process, so a process trains one auxiliary setting at a time; a checkpoint loaded in a process that never called it trains without the auxiliary loss.
  • loss_items changed shape around ultralytics 8.4.13x; both shapes are handled.
  • How the block computes in training and inference is on ES-MoE and YOLO; how every number was measured is on Experiments.