Decentralized Diffusion Models (2501.05450v2)

Abstract: Large-scale AI model training divides work across thousands of GPUs, then synchronizes gradients across them at each step. This incurs a significant network burden that only centralized, monolithic clusters can support, driving up infrastructure costs and straining power systems. We propose Decentralized Diffusion Models, a scalable framework for distributing diffusion model training across independent clusters or datacenters by eliminating the dependence on a centralized, high-bandwidth networking fabric. Our method trains a set of expert diffusion models over partitions of the dataset, each in full isolation from one another. At inference time, the experts ensemble through a lightweight router. We show that the ensemble collectively optimizes the same objective as a single model trained over the whole dataset. This means we can divide the training burden among a number of "compute islands," lowering infrastructure costs and improving resilience to localized GPU failures. Decentralized diffusion models empower researchers to take advantage of smaller, more cost-effective and more readily available compute like on-demand GPU nodes rather than central integrated systems. We conduct extensive experiments on ImageNet and LAION Aesthetics, showing that decentralized diffusion models FLOP-for-FLOP outperform standard diffusion models. We finally scale our approach to 24 billion parameters, demonstrating that high-quality diffusion models can now be trained with just eight individual GPU nodes in less than a week.

• The paper introduces a decentralized diffusion model framework that trains independent expert models on data subsets and ensembles their outputs to match centralized performance.
• It eliminates high-bandwidth communication by isolating training phases and using a lightweight routing mechanism for efficient inference.
• Experiments on ImageNet and LAION Aesthetics validate scalability, training a 24B-parameter model on eight GPU nodes in under a week.

The paper "Decentralized Diffusion Models" (2501.05450) introduces a novel framework for training diffusion models across independent clusters or datacenters, removing the reliance on centralized, high-bandwidth networking. The method involves training expert diffusion models on partitions of the dataset in isolation and ensembling them at inference through a lightweight router.

Core Methodology of Decentralized Diffusion Models

The core innovation of Decentralized Diffusion Models (DDMs) lies in its approach to dataset partitioning and distributed training. Unlike traditional methods that require synchronous gradient updates across all GPUs, DDMs operate on a principle of complete isolation during the training phase. Each expert model is assigned a subset of the data and trained independently, thereby eliminating the need for high-bandwidth communication.

The objective function optimized by the ensemble of experts collectively approximates that of a monolithic model trained on the entire dataset. This is achieved through a carefully designed routing mechanism that directs inference requests to the most appropriate expert based on input characteristics. The router's architecture is lightweight, ensuring minimal overhead during inference.

The mathematical underpinning of DDMs involves ensuring that the aggregation of expert models converges to the same solution as a single, centrally trained model. This is typically achieved by:

1. Data Partitioning: Dividing the dataset $D$ into $n$ disjoint subsets $D_i$, where $\bigcup_{i=1}^{n} D_i = D$.
2. Expert Training: Training each expert model $\theta_i$ on its respective dataset $D_i$ to optimize the local objective function $L(\theta_i, D_i)$.
3. Ensemble Aggregation: Combining the predictions of the expert models through a router $R$ such that the ensemble prediction approximates the prediction of a monolithic model $\theta$ trained on the entire dataset $D$, i.e., $$E[\theta(D)] \approx E[R(\theta_1(D_1), \theta_2(D_2), ..., \theta_n(D_n))]$$.

Experimental Results

The paper provides empirical validation of DDMs on ImageNet and LAION Aesthetics datasets. The results demonstrate that DDMs achieve competitive or superior performance compared to standard diffusion models when evaluated FLOP-for-FLOP.

A key finding is the scalability of DDMs. The authors successfully trained a model with 24 billion parameters using only eight GPU nodes in under a week. This highlights the potential for DDMs to democratize access to large-scale AI model training by reducing infrastructure costs and reliance on specialized hardware.

The optimal number of experts was empirically determined to be around eight, balancing model parameterization and computational efficiency. Different configurations were tested to validate the scalability and adaptability of the approach.

Implications and Scalability

DDMs have several implications for the future of AI model training and deployment:

• Cost Reduction: By enabling training on smaller, more readily available compute resources, DDMs lower the barrier to entry for researchers and organizations with limited budgets.
• Enhanced Resilience: The decentralized nature of DDMs makes them more resilient to localized GPU failures. If one expert model fails, the remaining models can still provide reasonable performance.
• Privacy Preservation: DDMs facilitate training on decentralized data sources, potentially preserving privacy and sovereignty in sensitive domains such as medical imaging.
• Federated Learning: The approach can be combined with federated learning techniques to enable collaborative model training without sharing raw data.

The scalability of DDMs is a significant advantage. The ability to train models with billions of parameters on a small number of GPU nodes opens up new possibilities for deploying AI in resource-constrained environments. Additionally, DDMs can be combined with other techniques such as quantization and pruning to further reduce the computational and memory requirements of the models.

In summary, the paper "Decentralized Diffusion Models" (2501.05450) presents a compelling alternative to traditional, monolithic training approaches for diffusion models. By decentralizing the training process, DDMs offer significant advantages in terms of cost, scalability, and resilience. The results demonstrate the potential of DDMs to democratize access to large-scale AI model training and enable new applications in resource-constrained environments.