MR-DiT: Adaptive Material Refinement Transformer

• The paper introduces MR-DiT, which employs adaptive multi-expert denoising with learnable mixing coefficients, achieving superior benchmarks such as FID 3.22 and IS 207.54.
• MR-DiT uses a set of K basis models to dynamically create expert models via softmax-normalized mixing, reducing training cost to O(K) while handling heterogeneity in material data.
• The model shows practical improvements in tasks like PBR map synthesis and microstructure recovery, reducing processing time to around 12 seconds compared to previous minute-long approaches.

Material Refinement Diffusion Transformer (MR-DiT) refers to the application of advanced diffusion-based transformer architectures, particularly leveraging multi-expert or mixture-of-expert strategies, to tasks requiring the synthesis, denoising, or refinement of material-related data representations. Such applications commonly arise in scientific and engineering contexts, including physically-based rendering (PBR) map generation, texture/structure synthesis, and crystalline/microstructure reconstruction. MR-DiT approaches seek to combine the strengths of Diffusion Transformers (DiT) with adaptive expert allocation to efficiently address the variable complexity inherent in material refinement denoising trajectories.

1. Foundational Architectures: DiT and Multi-Expert Mixing

DiT architectures are transformer-based diffusion models that replace traditional convolutional UNets with scalable, token-based processing, well-suited to high-resolution image or multi-channel outputs. In standard DiT, a single model is trained to perform denoising across the entire diffusion chain, resulting in efficient, high-fidelity generation; however, capacity allocation is static and may be suboptimal for trajectories with heterogeneous complexity.

Remix-DiT (Fang et al., 7 Dec 2024) introduces a principled mechanism for crafting $N$ expert models for different denoising timesteps, without incurring the training and storage costs of $N$ independent networks. Instead, $K$ basis models (with $K \ll N$) are trained concurrently, and learnable mixing coefficients $\mathbf{A}_{N \times K}$ are used to construct parameters for each timestep or interval expert as soft mixtures of the basis parameters:

$$\mathbf{W}_{N \times P}^{\text{experts}} = \mathbf{A}_{N \times K} \cdot \mathbf{W}_{K \times P}$$

where $\mathbf{W}_{K \times P}$ denotes basis parameters and each row of $\mathbf{A}$ indexes softmax-normalized mixing weights per expert. This architecture supports adaptive specialization throughout the denoising trajectory, allocating model capacity in accordance with local complexity.

2. Multi-Expert Denoising and Parameter Mixing

MR-DiT leverages the multi-expert paradigm in diffusion transformers, in which different timesteps of the reverse diffusion process are handled by dedicated experts. Each expert is not separately trained but is dynamically created from $K$ shared basis models via mixing. At training, an expert corresponding to the active timestep interval is sampled, and its parameters formed as a linear combination of the basis sets according to the learned coefficients. The loss associated with each expert involves the standard denoising objective:

$$\mathcal{L}(\theta_i) := \mathbb{E}_{t \in \mathcal{T}_i} \mathbb{E}_{x_0, \epsilon} \left[\left\| \epsilon - \epsilon_{\theta_i}( \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, t ) \right\|^2 \right]$$

where $\theta_i$ are the mixed parameters for expert $i$ operating on interval $\mathcal{T}_i$. Regularization on the mixing coefficients may be employed to encourage initial specialization, using a one-hot prior that is gradually relaxed.

3. Efficiency, Specialization, and Model Capacity Allocation

The MR-DiT approach inherits Remix-DiT’s efficiency: only $K$ basis models are stored and trained, with $O(K)$ training cost that does not scale with the number of experts $N$. Inference at each denoising step is performed using the mixed weights of a single expert, with the same computational cost as a standard DiT model. Empirically, the mixing coefficients reflect the changing demands of the denoising schedule, manifesting as more ensemble-like utilization of bases during early (high noise, global content) steps and sharper, nearly one-hot specialization during later (fine detail, low noise) steps.

This adaptive allocation allows MR-DiT models to concentrate capacity where material refinement is most challenging — for example, in late denoising steps correlated with microstructure recovery or texture details.

Aspect	Remix-DiT Approach	Standard Multi-Expert	Plain DiT
Number of experts	$N$ (large, e.g., 20)	Yes	No
Trainable sets (storage)	$K$ (e.g., 4, 8)	$N$ (e.g., 20, 1000)	1
Training cost	$O(K)$	$O(N)$	$O(1)$
Capacity per step	Adaptive, learned	Fixed	Fixed
Inference cost	$O(1)$	$O(1)$	$O(1)$
Generation quality	Highest, adaptive	Good	Baseline

4. Application to Material Synthesis and Refinement

MaterialPicker (Ma et al., 4 Dec 2024) exemplifies the practical use of DiT variants for the synthesis of high-quality PBR material maps via text, photographs, or joint conditioning. In this context, the model recasts the task as a multi-channel video prediction problem, stacking image, segmentation mask, and material map frames along the temporal axis. This allows the DiT backbone to exploit cross-channel relationships and global receptive field attention, facilitating distortion correction and diversity in output.

A plausible implication is that MR-DiT, by incorporating multi-expert denoising and adaptive capacity allocation, can further benefit material refinement tasks. For cases involving complex microstructural data, distribution shift, or multi-scale requirements, the ability to specialize model behavior at different diffusion steps (via Remix-DiT mixing) can enhance both refinement fidelity and efficiency. Such frameworks reduce the need for pixel-aligned inputs or explicit user annotation, as the self-attention mechanism and learned expert specializations handle misalignment and semantic ambiguity.

5. Empirical Performance and Benchmarking

On benchmark datasets such as ImageNet 256×256, Remix-DiT with 4 basis models and up to 20 experts demonstrates superior quantitative performance relative to both plain DiT and independently trained ensembles. For example, Remix-L achieves FID $3.22$ and IS $207.54$, outperforming both the DiT-L baseline and Multi-Expert variants. Similar trends are observed in generation quality, with improved object structure and detail retention at critical denoising steps.

In material generation, models like MaterialPicker report substantial improvements in distortion correction (automatic perspective/occlusion handling), output diversity (robust to non-stationary and highly textured materials), and workflow efficiency (material synthesis in $\sim12$s, compared to several minutes for previous approaches).

6. Mathematical Formulation and Implementation Considerations

The core mixing operation in multi-expert MR-DiT models is matrix-multiplication based:

$$\mathbf{W}_{N \times P}^{\text{experts}} = \mathbf{A}_{N \times K} \cdot \mathbf{W}_{K \times P}$$

where after softmax normalization, each expert's parameters are a convex combination of the $K$ bases. Efficient implementation entails that the mixed weights for each expert can be precomputed; at inference, only one expert's parameters are used per denoising step. Training involves sampling timesteps and corresponding experts, applying the standard denoising loss, and updating basis and mixing coefficients jointly.

For material-specific applications, latent encoders such as VAE are often employed to reduce spatial resolution before tokenization, facilitating DiT training. Conditioning inputs consist of CLIP-encoded text tokens and patchwise image tokens, processed jointly with positional and mask encodings.

7. Significance, Generalization, and Future Directions

MR-DiT represents a convergence of efficient transformer-based diffusion modeling with adaptive, multi-expert specialization, directly addressing the challenges of material refinement, texture synthesis, and multi-channel image generation. It leverages global self-attention, learns cross-channel relationships, and adaptively allocates architectural capacity to match denoising complexity over trajectory steps. This suggests strong prospects for MR-DiT frameworks to generalize across diverse materials, support inverse rendering, and scale efficiently to complex scientific domains with heterogeneous data.

Future directions include further integration of multi-modal conditioning, extensions to hierarchical refinement schedules, and application to domains where material properties, structure, and generative diversity are critical. The basis-mixing principle offers flexibility for targeted adaptation, transfer learning across material classes, and efficient deployment in resource-constrained environments.

---