SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs - ICML 2026

This repository contains the official implementation of SoftMoE, a soft and differentiable routing method for Mixture-of-Experts layers in large language models, as presented in our ICML 2026 paper:

"SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs"
arXiv preprint | OpenReview

Repository Structure

.
├── external/Megatron-LM/        # Megatron-LM submodule pinned to the paper baseline
├── patches/                    # SoftMoE patch applied on top of Megatron-LM
├── scripts/                    # Setup, preprocessing, and training launch scripts
│   └── utils/                  # Dataset and tokenizer utilities
├── train_configs/              # Paper training configurations
├── img/                        # Paper figures
├── requirements.txt            # Python dependencies
└── README.md                   # You are here

Tutorials and Scripts

For setup, OpenWebText preprocessing, tokenizer training, Megatron preprocessing, and launch scripts, see ./scripts/README.md.

Utility scripts used by the pipeline are documented in ./scripts/utils/README.md.

Paper training configurations are available in ./train_configs/:

• hard_topk.sh
• soft_topk.sh
• learn_k.sh

Tested Environment

Our experiments were run on NVIDIA GH200 systems. The training environment used Python 3.11 and a CUDA 12.4-based software stack.

For reproducibility, make sure that the low-level GPU packages are compiled or installed for the exact Python/CUDA/PyTorch combination used on your machine, in particular:

• PyTorch
• Triton
• Apex
• FlashAttention
• Transformer Engine

The requirements.txt file lists the Python package set used by the project, but cluster-specific CUDA wheels and compiled extensions may need to be installed from your system-provided packages or rebuilt for your environment.

License

This code is released for research use. The Megatron-LM submodule remains governed by its original license in external/Megatron-LM/LICENSE.

Citation

If you use SoftMoE in your research, please cite our arXiv preprint. The paper has been accepted to ICML 2026, but the conference proceedings are not publicly probably until conference in July.

@inproceedings{
zasada2026softmoe,
title={SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in {LLM}s},
author={Miko{\l}aj Zasada and {\L}ukasz Struski and Jacek Tabor and Marcin Kurdziel},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=vGTFOp3jLO}
}