KO | EN

GitLite

Minimalist Open Source Repository Search Engine

gitlite — search

#dsh #dsh-plugin #developer-tools #ai-agent #claude-code #codex #typescript #ai-agents #deepseek-harness #mcp #cli #rust

DiffMoE

★ 176 Open GitHub ↗

[Arxiv 2025] Official PyTorch implementation of DiffMoE, TC-DiT, EC-DiT and Dense DiT

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers

### [<a href="https://huggingface.co/papers/2503.14487" target="_blank">arXiv</a>] [<a href="https://shiml20.github.io/DiffMoE/" target="_blank">Project Page</a>] [<a href="https://huggingface.co/datasets/KwaiVGI/DiffMoE/" target="_blank">Code</a>]

_**[Minglei Shi1*](https://github.com/shiml20/), [Ziyang Yuan1*](https://scholar.google.ru/citations?user=fWxWEzsAAAAJ&hl=en), [Haotian Yang2](https://scholar.google.com/citations?user=LH71RGkAAAAJ&hl=en), [Xintao Wang2†](https://xinntao.github.io/), [Mingwu Zheng2](https://scholar.google.com.hk/citations?user=MdizB60AAAAJ&hl=en), [Xin Tao2](https://www.xtao.website/), [Wenliang Zhao1](https://wl-zhao.github.io/), [Wenzhao Zheng1](https://wzzheng.net/), [Jie Zhou1](https://www.imem.tsinghua.edu.cn/info/1330/2128.htm),[Jiwen Lu1†](https://www.au.tsinghua.edu.cn/info/1078/3156.htm), [Pengfei Wan2](https://scholar.google.com/citations?user=P6MraaYAAAAJ&hl=en), [Di Zhang2](https://openreview.net/profile?id=~Di_ZHANG3), [Kun Gai2](https://scholar.google.com/citations?user=PXO4ygEAAAAJ&hl=zh-CN)**_
 
(*equal contribution, †corresponding author.)

1Tsinghua University, 2Kuaishou Technology.

</div>

## 🔥 Updates
- __[2025.3.20]__: Release the [code](https://github.com/KwaiVGI/DiffMoE/) of DiffMoE
- __[2025.3.19]__: Release the [project page](https://shiml20.github.io/DiffMoE/) of DiffMoE

## 📖 Introduction

**TL;DR:** DiffMoE is a dynamic MoE Transformer that outperforms 3× larger dense models in diffusion tasks, using global token pool and adaptive routing while keeping 1× parameter activation.

This repo contains:

* 🪐 A simple PyTorch implementation of [Dense-DiT](models/models_DiT.pu), [EC-DiT](models/models_ECDiT.py), [TC-DiT](models/models_TCDiT.py), [DifffMoE](models/models_DiffMoE.py)

* 🛸 An DiffMoE [training and evaluation script](train.py) using PyTorch DDP

### To-do list

- [x] training / inference scripts
- [ ] huggingface ckpts

## ✨ Key Points

Token Accessibility and Dynamic Computation. (a) Token accessibility levels from token isolation to crosssample interaction. Colors represent tokens in different samples, ti indicates noise levels. (b) Performance-accessibility analysis across architectures. (c) Computational dynamics during diffusion sampling, showing adaptive computation from noise to image. (d) Class-wise computation allocation from hard (technical diagrams) to easy (natural photos) tasks. Results from DiffMoE-L-E16-Flow (700K).

DiffMoE Architecture Overview. DiffMoE flattens tokens into a batch-level global token pool, where each expert maintains a fixed training capacity of $C^{E_i}_{train} = 1$. During inference, a dynamic capacity predictor adaptively routes tokens across different sampling steps and conditions. Different colors denote tokens from distinct samples, while ti represents corresponding noise levels.

## Preparation

### Dataset
Download [ImageNet](http://image-net.org/download) dataset, and place it in your `IMAGENET_PATH`.

### Installation

Download the code:
```
git clone https://github.com/KwaiVGI/DiffMoE.git 
cd DiffMoE
```

A suitable [conda](https://conda.io/) environment named `diffmoe` can be created and activated with:

```
conda env create -f environment.yaml
conda activate diffmoe
```

Pre-trained models are Coming soon, stay tuned !

| Model Name                 | # Avg. Act. Params | Training Step | CFG | FID-50K↓  | Inception Score↑ |
|----------------------------|-------------------------|-------------|-----|---------|----------------|
| TC-DiT-L-E16-Flow   | 458M               | 700K        | 1.0 | 19.06   | 73.49           |
| EC-DiT-L-E16-Flow    | 458M               | 700K        | 1.0 | 16.12   | 82.37           |
| Dense-DiT-L-Flow     | 458M               | 700K        | 1.0 | 17.01   | 78.17           |
| Dense-DiT-XL-Flow   | 675M               | 700K        | 1.0 | 14.77   | 86.82           |
| DiffMoE-L-E16-Flow   | 454M               | 700K        | 1.0 | 14.41   | 88.19           |
| Dense-DiT-XL-Flow   | 458M                | 7000K       | 1.0 | 9.47 | 115.58         |
| DiffMoE-L-E8-Flow   | 458M                | 7000K       | 1.0 | 9.60 | 131.45         |
| Dense-DiT-XL-DDPM   | 458M                | 7000K       | 1.0 | 9.62 | 123.19         |
| DiffMoE-L-E8-DDPM   | 458M                | 7000K       | 1.0 | 9.17 | 131.10         |

<!-- ### (Optional) Caching VAE Latents

Given that our data augmentation consists of simple center cropping and random flipping, 
the VAE latents can be pre-computed and saved to `CACHED_PATH` to save computations during DiffMoE training:

```
torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_cache.py \
--img_size 256 --vae_path stabilityai/sd-vae-ft-mse \
--batch_size 128 \
--data_path ${IMAGENET_PATH} --cached_path ${CACHED_PATH}
``` -->

## Usage

### Training
Script for the default setting:
```
CUDA_VISIBLE_DEVICES=0,1,2,3 \
torchrun --nproc_per_node=4 --nnodes=1 --node_rank=0 \
train.py --config ./config/000_DiffMoE_S_E16_Flow.yaml
```

### Evaluation (ImageNet 256x256)

Evaluate DiffMoE:

```
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
torchrun --nnodes=1 --nproc_per_node=8 \
sample_ddp_feature.py --image-size 256 \
    --per-proc-batch-size 125 --num-fid-samples 50000 --cfg-scale 1.0 --num-sampling-steps 250 --sample-dir samples \
    --ckpt exps/EXPNAME/checkpoints/xxxxxx.pt
```

## Acknowledgements
We thank [Zihan, Qiu](https://scholar.google.com/citations?user=24eVHiYAAAAJ&hl=en) for helpful discussion. A large portion of codes in this repo is based on [MAR](https://github.com/LTH14/mar), [DiT](https://github.com/facebookresearch/DiT), [DeepSeekMoE](https://github.com/deepseek-ai/DeepSeek-MoE)

## 🌟 Citation

```
@misc{shi2025diffmoedynamictokenselection,
      title={DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers}, 
      author={Minglei Shi and Ziyang Yuan and Haotian Yang and Xintao Wang and Mingwu Zheng and Xin Tao and Wenliang Zhao and Wenzhao Zheng and Jie Zhou and Jiwen Lu and Pengfei Wan and Di Zhang and Kun Gai},
      year={2025},
      eprint={2503.14487},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2503.14487}, 
}
```

DiffMoE

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README