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

TL;DR

Remix-DiT addresses the high training cost of transformer-based diffusion models by constructing N denoising experts from K basis models using learnable mixing coefficients. The method formulates expert parameters as $\boldsymbol{\Theta}=\boldsymbol{\alpha}\boldsymbol{\beta}$, extends DiT architecture to support mixing, and trains with a single active expert per step while updating all basis models. Experiments on ImageNet-256×256 show Remix-DiT can match or exceed the performance of independent multi-expert baselines within similar budgets, with coefficients that reveal adaptive allocation—more capacity at early timesteps and ensemble-like mixtures at later steps. The approach also supports leveraging pre-trained DiT weights via a prior coefficient regularization, enabling efficient fine-tuning and offering substantial practical gains for scalable diffusion modeling.

Abstract

Transformer-based diffusion models have achieved significant advancements across a variety of generative tasks. However, producing high-quality outputs typically necessitates large transformer models, which result in substantial training and inference overhead. In this work, we investigate an alternative approach involving multiple experts for denoising, and introduce Remix-DiT, a novel method designed to enhance output quality at a low cost. The goal of Remix-DiT is to craft N diffusion experts for different denoising timesteps, yet without the need for expensive training of N independent models. To achieve this, Remix-DiT employs K basis models (where K < N) and utilizes learnable mixing coefficients to adaptively craft expert models. This design offers two significant advantages: first, although the total model size is increased, the model produced by the mixing operation shares the same architecture as a plain model, making the overall model as efficient as a standard diffusion transformer. Second, the learnable mixing adaptively allocates model capacity across timesteps, thereby effectively improving generation quality. Experiments conducted on the ImageNet dataset demonstrate that Remix-DiT achieves promising results compared to standard diffusion transformers and other multiple-expert methods. The code is available at https://github.com/VainF/Remix-DiT.

Figures (6)

• Figure 1: (a) Re-training a larger DiT incurs significant training and inference costs. (b) Multi-expert denoising trains multiple expert models to improve generation quality while maintaining low inference overhead. However, training multiple experts still results in significant training costs. (c) This work introduces RemixDiT, a learnable method to craft any number of experts by mixing basis models.
• Figure 2: An example of mixing 4 linear layers basis into 6 expert layers. Each expert linear layer is a weighted averaging of the basis layers. At each denoising interval, only one expert is activated for inference or training. To increase the number of experts, we increase the number of coefficients ${\boldsymbol{\alpha}}$, which is more efficient than independently training new experts.
• Figure 3: Comparision to existing methods
• Figure 4: (a) The learned coefficients for a -S-4-20, which mixes 4 basis DiT-S to obtain 20 expert models, each associated with a 50-step interval. The x-axis shows the corresponding timestep intervals associated with each mixed model. The value in each grid refers to the coefficients that will be used to weight the corresponding basis models. At early timesteps $T\rightarrow999$, tend to use an ensembled model for inference, which is averaged over all basis models. And at the late timesteps $T\rightarrow 0$, more specialized models, such as the first basis model are picked for fine-grained prediction. (b) The learned coefficients for a Remix-DiT-S-4-8. (c) The loss value of mixed experts on different timestep intervals, where each expert shows lower errors on their assigned timesteps.
• Figure 5: Visualization of generated samples from DiT-B and Remix-DiT-B.