Very Simple MoE Intro

I've written a notebook containing a model and training code for a deliberately-simple character-level Mixture-of-Experts Transformer model. I've removed as many complications as I can from it while still retaining better performance on a per-FLOP basis than a similar dense Transformer. The only non-default dependencies are PyTorch, and WanDB. The following explains how this model works, and how Mixture-of-Experts work in general. (Note: The implementation is closest to a Switch Transformer. But it isn't really that. For instance, my implementation is trained solely on an autoregressive language modeling task, while Switch Transformer uses a masked token language modeling task; I use a different position encoding method; and so on. And a big chunk of the contribution of Switch Transformer is the routing across multiple GPUs in a particular way, while I completely ignore multi-GPU training or inference.) Why would we care about MoEs? Well, MoE Transformers generally take less compute th

Very Simple MoE Intro 1 A 3 O R N 1 A 3 O R N Very Simple MoE Intro Created: 2024-02-03 Wordcount: 3k Tags: essays machine-learning effort-post I've written a notebook containing a model and training code for a deliberately-simple character-level Mixture-of-Experts Transformer model. I've removed as many complications as I can from it while still retaining better performance on a per-FLOP basis than a similar dense Transformer. The only non-default dependencies are PyTorch, and WanDB. The following explains how this model works, and how Mixture-of-Experts work in general. ( Note: The implement

related reading

• Papers I’ve read this week, Mixture of Experts editionfinbarrtimbers.substack.com
• Mixture of experts - Wikipediaen.wikipedia.org
• Monet: Mixture of Monosemantic Experts for Transformers Explained — LessWronglesswrong.com
• [2202.08906] ST-MoE: Designing Stable and Transferable Sparse Expert Modelsar5iv.labs.arxiv.org
• A Mathematical Framework for Transformer Circuitstransformer-circuits.pub
• All the Transformer Math You Need to Know | How To Scale Your Modeljax-ml.github.io
• Transformer Explainer: LLM Transformer Model Visually Explainedpoloclub.github.io
• frontier model training methodologies | Alex Wa's Blogdjdumpling.github.io
• An Alternative to Test-Time Scalingrentry.org
• [1701.06538] Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layerarxiv.org
• MiLo: Efficient Quantized MoE Inference with Mixture of Low-Rank Compensatorsbeichenhuang.github.io
• The Annotated Transformernlp.seas.harvard.edu