Heterogeneous Decentralized Diffusion Models

Zhiying Jiang; Raihan Seraj; Marcos Villagra; Bidhan Roy

cs.LGMar 6, 2026

Heterogeneous Decentralized Diffusion Models

Authors: Zhiying Jiang, Raihan Seraj, Marcos Villagra, Bidhan Roy

Organizations: Bagel Labs

Abstract

Training frontier-scale diffusion models often requires substantial computational resources concentrated in tightly-coupled clusters, limiting participation to well-resourced institutions. While Decentralized Diffusion Models (DDM) enable training multiple experts in isolation, existing approaches require 1176 GPU-days and homogeneous training objectives across all experts. We present an efficient framework that dramatically reduces resource requirements while supporting heterogeneous training objectives. Our approach combines three key contributions: (1) a heterogeneous decentralized training paradigm that allows experts to use different objectives (DDPM and Flow Matching), unified at inference time without any retraining; (2) pretrained checkpoint conversion from ImageNet-DDPM to Flow Matching objectives, accelerating convergence and enabling initialization without objective-specific pretraining; and (3) PixArt- $α$ 's efficient AdaLN-Single architecture, reducing parameters while maintaining quality. Experiments on LAION-Aesthetics show that, relative to the training scale reported for prior DDM work, our approach reduces the compute by 16 $\times$ and data by 14 $\times$ . Under aligned inference settings, our heterogeneous configuration achieves better FID and higher intra-prompt diversity than the homogeneous baseline. By eliminating synchronization requirements and enabling mixed DDPM/FM objectives, our framework makes decentralized generative model training accessible to contributors with single GPUs requiring only 24--48GB VRAM.

Figures & tables

Panel of Figure 1 — Figure 1: Text-to-Image Generation with Heterogeneous Decentralized Diffusion. Our framework combines multiple expert models trained with different objectives (DDPM and Flow Matching) in complete isolation to generate high-quality, diverse images from text prompts. All samples are generated at 256 $\times$ 256 resolution using 8 experts trained on LAION-Aesthetics with only 72 A100-equivalent GPU-days of compute.

Panel of Figure 2 — Figure 2: Inference Pipeline for Heterogeneous Expert Fusion. Given noisy input $(x_{t},t,c)$ , the router predicts cluster probabilities $p_{\phi}(k|x_{t},t)$ to weight expert contributions. DDPM experts output epsilon predictions while Flow Matching experts output velocity predictions. Schedule-aware conversion functions deterministically unify all predictions into a common velocity space $v^{(k)}$ without retraining, enabling router-weighted fusion $u_{t}(x_{t})=\sum_{k=1}^{K}p_{\phi}(k|x_{t},t)\cdot v^{(k)}$ for ODE-based sampling.

Panel of Figure 3 — Figure 3: Impact of Pretrained Checkpoint Conversion. Comparison of generated samples after 75K optimization steps using the text prompt “a red car in front of the tree”.

Inference Strategy	FID-50K $\downarrow$
Monolithic (single model)	29.64
Top-1	30.60
Top-2	22.60
Full Ensemble (all experts)	47.89
Improvement vs. Monolithic	7.04

Table 1: FID-50K (lower is better) comparing monolithic training and decentralized multi-expert training with different inference strategies on LAION-Art. All models use the DiT-B/2 architecture (129M parameters per expert).

Method	Data	Compute	FID@50K $\downarrow$
DDM [ 22 ] †	158M	1176 A100 days	5.5–10.5
Ours (Homo)	11M	72 A100 days	12.45
Ours (Hetero)	11M	72 A100 days	11.88

Table 2: Baseline Comparison. Resource comparison against DDM. Both of our rows use the same aligned inference settings (CFG

=7.5

, 50 steps). Compute is in A100-equivalent GPU-days; our experiments use A40 48GB GPUs (120 A40 GPU-days total, normalized via measured training throughput). † The FID range 5.5–10.5 is estimated from results at varying training FLOPs in McAllister et al. [22]

Sampling Method	LPIPS $\uparrow$	FID $\downarrow$	CLIP $\uparrow$
Native DDPM	0.787	27.04	0.316 ±0.030
FM	0.752	20.23	0.324 ±0.034
DDPM $\rightarrow$ FM	0.761	25.61	0.319 ±0.032
Combined (same schedule)	0.782	32.67	0.312 ±0.035
Combined (diff. schedules)	0.777	33.29	0.312 ±0.034

Table 3: Sampling Quality Comparison. DDPM

\rightarrow

FM conversion improves over native DDPM and enables mixed-objective sampling. Combined experts achieve higher diversity (LPIPS) than single FM, though at a FID cost, reflecting the quality–diversity trade-off of heterogeneous fusion.

Model	CFG	Steps	FID-50K $\downarrow$
Homogeneous (8FM)	7.5	50	12.45
Heterogeneous (1DDPM:7FM)	6.0	75	19.75
Heterogeneous (2DDPM:6FM)	6.0	75	15.09
Heterogeneous (2DDPM:6FM)	7.5	50	11.88

Table 4: Homogeneous vs Heterogeneous Comparison. The first and last rows are directly comparable (same CFG and steps).

Panel of Figure 4 — Figure 4: Qualitative comparison: Homogeneous vs. Heterogeneous models. Images generated from identical prompts and random seeds. Homogeneous models (left, trained with Flow Matching only) often appear smoother in texture. Heterogeneous models (right, combining FM and DDPM experts) often preserve sharper local details and richer texture variation.

Panel of Figure 5 — Figure 5: Training Pipeline for Decentralized Heterogeneous Experts. LAION dataset $\mathcal{D}$ is partitioned into $K$ semantic clusters $\{S_{1},S_{2},\ldots,S_{K}\}$ using DINOv2 feature extraction and hierarchical clustering. Each expert trains independently on its assigned cluster with heterogeneous objectives: DDPM experts predict noise $\epsilon_{\theta_{k}}(x_{t},t)$ while Flow Matching experts predict velocity $v_{\theta_{k}}(x_{t},t)$ . The router network $\phi$ trains on all data to predict cluster assignments via cross-entropy loss. Crucially, there is zero gradient synchronization, parameter sharing, or activation passing between experts during training.

Figure 6: Validation Loss: Pretrained vs. Scratch Initialization. Average validation loss over early training steps for a Flow Matching expert initialized from a converted ImageNet-DDPM checkpoint versus training from scratch. Pretrained initialization yields 1.2 $\times$ faster loss reduction (0.030 vs. 0.025 drop over 5K steps in the steady-improvement regime) and consistently lower loss throughout training.

Panel of Figure 7 — Figure 7: Diverse Generation Examples (Set 1). Representative samples demonstrating our framework’s ability to generate high-quality, diverse images across multiple semantic categories. Images are produced by our heterogeneous ensemble combining DDPM and Flow Matching objectives.

Panel of Figure 8 — Figure 8: Diverse Generation Examples (Set 2). Additional samples showcasing consistent generation quality across various prompts and content types.

Panel of Figure 9 — Figure 9: Diverse Generation Examples (Set 3). Further examples demonstrating the robustness and versatility of our heterogeneous framework.

Panel of Figure 10 — Figure 10: Diverse Generation Examples (Set 4). Additional samples showcasing the framework’s ability to generate high-quality images across various prompts and configurations.

Panel of Figure 11 — Figure 11: Heterogeneous vs. Homogeneous Comparison (Part 1). Direct side-by-side comparison showing that our heterogeneous approach (columns 1-2) maintains or improves visual quality compared to homogeneous baseline (columns 3-4) across diverse categories including animals and nature.

Panel of Figure 12 — Figure 12: Heterogeneous vs. Homogeneous Comparison (Part 2). Continued comparison across vehicles and food categories, demonstrating consistent quality maintenance across different semantic domains.

Figure 13: Impact of Router Threshold on Generation Quality. Different thresholds affect quality-diversity trade-offs.

Panel of Figure 14 — Figure 14: Expert Ordering and Router Threshold Effects. Comparison of 2-expert heterogeneous configurations showing the impact of expert ordering (DDPM $\rightarrow$ FM vs. FM $\rightarrow$ DDPM) and router threshold values. All experiments use the same unified schedule. FM $\rightarrow$ DDPM ordering produces more stable, coherent results, while DDPM $\rightarrow$ FM shows higher sensitivity to threshold selection.

Panel of Figure 15 — Figure 15: Noise Schedule Comparison. Controlled comparison using 2-expert configuration with different training strategies. Columns 1-2: Different schedules — DDPM expert trained with cosine schedule and Flow Matching expert trained with linear interpolation schedule (their natural configurations). Columns 3-4: Same schedule baseline — both DDPM and Flow Matching experts trained with cosine schedule. The different-schedules approach shows qualitatively sharper details in these examples, although the same-schedule setting achieves slightly better FID in Table 3 of the main paper.

Panel of Figure 16 — Figure 16: Sampling Strategy Comparison. Visual comparison of expert selection strategies on identical prompts. Top-1 commits to a single expert per timestep, producing sharp but potentially limited samples. Top-2 achieves optimal FID by balancing complementary information with prediction coherence. Full ensemble (K=8) can over-smooth details when averaging divergent expert predictions. All strategies use 75 Euler steps with CFG scale 6.

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,176 papers

Improved Distributional Diffusion Models
Sep 29, 2026Tommaso Martorella, Alexandre Galashov, Felix Krause +4Few-Step Diffusion Sampling Diffusion-Based Image Generation
Munich Center for Machine Learning (MCML) · Google DeepMind
Accelerating Diffusion via Hybrid Data-Pipeline Parallelism Based on Conditional Guidance Scheduling
Feb 25, 2026Euisoo Jung, Byunghyun Kim, Hyunjin Kim +2Diffusion Model Inference Acceleration Diffusion Model Guidance
School of Computing, KAIST
PixelMoE: Scaling Pixel-Space Diffusion with LLM-Style Mixture-of-Experts
Oct 13, 2026Ruixiao Dong, Jianmin Bao, Bozhi Luan +6Diffusion-Based Image Generation Pixel-Space Diffusion Models
University of Science and Technology of China · KlingAI Research · Peking University

Sep 29, 2026J/K move · Enter open · S save

Improved Distributional Diffusion Models

Tommaso Martorella, Alexandre Galashov, Felix Krause, Stefan Andreas Baumann+3

Munich Center for Machine Learning (MCML) · Google DeepMind

Few-Step Diffusion Sampling Diffusion-Based Image Generation T2I Diffusion Models Diffusion Models

Open paper AskPDF arXiv

Distributional Diffusion Models (DDMs) replace the standard mean-prediction denoiser with a distributional denoiser trained via a scoring rule objective, learning a stochastic approximation to

p(x_1 \mid x_t)

rather than its conditional mean. However, scaling DDMs to modern image-generation settings faces two obstacles: (i) multi-particle training incurs overhead that scales with the number of particles, (ii) DDMs use globally fixed scoring rule hyperparameters, forcing a single trade-off across sampling budgets. We mitigate these limitations by deferring particle expansion to late transformer layers, and the hyperparameter trade-off by introducing time-dependent scoring rule schedules informed by the dynamical regimes of~\citet{Biroli2024}. Combined with a DiT-based latent setup, these changes make DDM training practical on class-conditional ImageNet-

256^2

, achieving 4.48 FID at 4 steps and 2.38 at 50 steps with DiT-XL/2, from a single model trained from scratch in one stage, without a teacher, self-distillation or JVPs. The result is a stochastic few-step generator whose FID does not degrade as the sampling budget grows from 4 to 50 NFE, and the same recipe transfers to text-to-image generation. Code and pre-trained models available at https://github.com/CompVis/iDDM.

Abstract

Diffusion Models

Explore similar work

Improved Distributional Diffusion Models

Accelerating Diffusion via Hybrid Data-Pipeline Parallelism Based on Conditional Guidance Scheduling

PixelMoE: Scaling Pixel-Space Diffusion with LLM-Style Mixture-of-Experts