Research map / Image generation / Autoregressive and latent generation

Remix-DiT: Mixing Diffusion Transformers for Multi-Expert Denoising

NeurIPS 2024 · Image generation › Autoregressive and latent generation · approach: High-resolution diffusion synthesis · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Remix-DiT crafts N diffusion-transformer denoising experts for different timestep intervals by softmax-mixing K<N learnable basis DiT models, so inference cost matches a standard DiT.

Problem

Multi-expert denoising for diffusion transformers improves quality but training N independent experts is costly, and the number of experts and timestep partition are hand-chosen heuristics.

Approach

Learn K basis DiT models and per-expert softmax mixing coefficients over timesteps so that N experts are weighted averages of the bases; the mixed expert has the architecture of a plain DiT, and the coefficients learn an adaptive timestep partition. A one-hot prior regularizer breaks symmetry when initializing from a pretrained DiT.

Key results

• Remix-B-4-20 (100K fine-tune from DiT-B/2) reaches FID 9.02 vs 9.31 for continual training and 9.28 for 8x multi-expert baseline, at the same 23.04 GFLOPs per step.
• Remix-S-4-20 improves FID from 26.05 (continual training) to 22.84 with IS 69.23 at 6.07 GFLOPs per step.
• Remix-L-4-20 achieves FID 3.22 vs 3.73 for DiT-L/2-1M, with IS 207.54.

Measured effects reported in the paper

Surprising findings

• Raw (unconstrained) mixing coefficients performed worse than softmax-constrained mixing because negative coefficients made learning harder.
• Increasing the number of experts beyond 20 (to 100 or 1000) did not consistently help; 1000 experts degraded FID to 35.60 for DiT-S.

Limitations

• Only one expert is activated per forward/backward pass, so gradients to mixing coefficients are sparse and need momentum updates.
• Evaluated only on ImageNet 256x256 class-conditional generation with DiT backbones; no error bars reported.
• Extra GPU memory is needed to store basis models, and training is somewhat slower.

Experimental setup

Datasets: ImageNet

Benchmarks: ImageNet 256x256

Metrics: FID IS Precision Recall GFLOPs per step steps per second GPU memory latency

Baselines: DiT-S/2 DiT-B/2 DiT-L/2 Multi-expert 8x denoisers (continual training baseline) DiT continual training

Builds on: DiT E-Diff MEME Multi-expert denoising

Closest papers by meaning

• High-Fidelity Diffusion-Based Image Editing (AAAI 2024 · similarity 0.91)
• Denoising Diffusion Bridge Models (ICLR 2024 · similarity 0.90)
• DEFT: Efficient Fine-tuning of Diffusion Models by Learning the Generalised $h$-transform (NeurIPS 2024 · similarity 0.90)
• ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.90)
• FlexiDiT: Your Diffusion Transformer Can Easily Generate High-Quality Samples with Less Compute (CVPR 2025 · similarity 0.90)
• Decouple-Then-Merge: Finetune Diffusion Models as Multi-Task Learning (CVPR 2025 · similarity 0.89)