kyegomez/latent-moe

Latent-MoE is a PyTorch implementation of a research paper from NVIDIA about a more efficient way to build Mixture of Experts layers, a technique used inside large AI models where only a subset of specialized sub-networks, called experts, are activated for each piece of input rather than running the whole model every time. This project packages the idea from that paper into a single file that can be dropped in as a replacement for a standard Mixture of Experts layer. The core idea is that a normal Mixture of Experts layer routes tokens and runs its experts using the model's full internal size. LatentMoE instead first shrinks each token down into a smaller working space using a shared projection step, runs all the expert calculations in that smaller space, and then expands the result back up to full size afterward. Because the expensive parts, communication between experts and loading expert weights, now happen in the smaller space, both cost less. Those savings can then be spent on adding more experts, which the paper argues improves the tradeoff between accuracy and computational cost. The project offers two modes: one that keeps the same number of active experts per token to lower cost at equal accuracy, and one that increases active experts per token to raise accuracy at equal cost, which the README recommends as the better tradeoff. You install it with a Python package manager, then import it into existing PyTorch code as a drop-in layer, supplying settings like model dimension, number of experts, and the compression factor. The project also includes a full example of a small transformer language model built around this layer, plus a ready to run training script that trains on a subset of Wikipedia text, complete with command line flags for adjusting model size and training settings and support for resuming from a saved checkpoint. The README asks that anyone using this implementation in research cite both the original paper and this repository. The project is written in Python and released under the Apache License Version 2.0.