sparse mixture of experts

A mixture-of-experts block for the model you already have¶

Native implementations on PyTorch, TensorFlow, Jax, Flax and Keras, each written the way that framework is written, and one call to add it to Ultralytics YOLO. The router's balancing term reaches backward(), rather than sitting in a config.

router picks 2 of 4 per image

Into any modelIt sizes itself on the first forward pass, so an nn.Sequential takes it as written, and it grafts into a model that already exists by submodule name, without touching its source.
One call to add it to YOLOequip grafts the block into a config, renumbers layer references and wires the loss, on seven official generations from YOLOv5 to YOLO26 and on YOLO-Master's fork.
Five frameworks, one set of argumentskernels sets the experts and their receptive fields, top_k how many run; the code that reads them is shared, so moving framework changes nothing you write.

5frameworks, PyTorch to Keras

7YOLO generations, v5 to 26

1call to wire it into training

anymodel takes the block

Torchthe two knobs, the auxiliary loss, grafting into a model that exists Ultralytics YOLOinstall, the three calls, grafting and renumbering, a comparison you can defend TensorFlowinside tf.function, differentiated by a GradientTape Jaxpure init and apply, under jax.jit and jax.grad Flaxan nnx.Module, under nnx.jit and nnx.grad Keras 3one implementation on tensorflow, jax and torch APIevery argument of ESMoE, and the entry point of each backend and integration Designwhat the block computes, how the package is layered, what a new backend has to write Experimentshow the defaults were chosen, with a figure behind each finding Effect chartper-seed paired deltas across seven generations Judgment linescriteria and predictions first, verdicts after

Open the quick start in Colab: install, equip, train and watch the esmoe_aux column, all on a free GPU.

Install¶

uv add esmoe            # the block, for any torch model
uv add esmoe[yolo]      # ... and the Ultralytics integration

esmoe[ultralytics] is that extra under its other name; esmoe[tensorflow], esmoe[jax], esmoe[flax] and esmoe[keras] each bring one backend. The distribution, the import and the command line are all esmoe.

Use¶

import esmoe
from torch import nn

net = nn.Sequential(nn.Conv2d(3, 64, 3, 2, 1), esmoe.ESMoE(kernels=[3, 5, 7], top_k=1))

In YOLO:

model = esmoe.equip("yolo11n.yaml", weight=0.01)   # register + graft + build + wire
model.train(data="coco8.yaml", epochs=10)

The separate steps, the command line and the hand-written config are in In Ultralytics YOLO; the plain torch path is in In a torch model.

Compatibility¶

backbone	build + forward	grafted config	aux loss in training	protocol runs
YOLOv5	yes	yes	yes	yes
YOLOv8	yes	yes	yes	yes
YOLOv9	yes	yes	yes	yes
YOLOv10	yes	yes	yes	yes
YOLO11	yes	yes	yes	yes
YOLO12	yes	yes	yes	yes
YOLO26	yes	yes	yes	yes
YOLO-Master (fork)	yes	yes	yes	no

Verified by tests/test_ultralytics.py on ultralytics 8.4.101 and the latest release (the CI matrix), which report loss items in two different shapes; both are handled. The training column is backed by real 1-epoch VisDrone runs on four generations and by the 120-epoch protocol runs on all seven, each logging a non-zero train/esmoe_aux; graft and forward run on every row in CI. The YOLO-Master row runs against the fork's vendored ultralytics: scripts/fork_smoke.py grafts their yolo-master-n.yaml, trains one epoch with a non-zero esmoe_aux, and builds their own ES_MOE config alongside ours. The same-configuration comparison trains upstream's own blocks on the fork and trains yolo-master-n with this package's blocks on official ultralytics, hence the last column.

Default configuration¶

ESMoE(num_experts=4, top_k=2) with attach_aux_loss(weight=0.01). Under one budget of a 25% train split, 640 px and 20 epochs, it beat the 2-expert, 8-expert, top-1 and aux-off variants; on the full VisDrone training set, at the same 20 epochs and 640 px, three seeds confirmed it: 3/3 paired wins, +0.0021 mAP50, and 10.4% more parameters on YOLOv8n. Under the protocol budget (800 px, 120 epochs) the default wiring on YOLOv8n is +0.0025, 2/3. The argument is on Selection.