FedPhD: Federated Diffusion Model Training

• The paper introduces a federated learning framework that leverages hierarchical aggregation, homogeneity-aware weighting, and structured pruning to optimize diffusion model training on non-IID data.
• The framework reduces communication by up to 88% and uses only 56% of computation resources, achieving at least 34% improvement in image generation quality over baselines.
• FedPhD employs a three-tier architecture that combines client-edge-global aggregation to effectively mitigate data heterogeneity and resource constraints in edge environments.

Searching arXiv for FedPhD and directly related background papers. FedPhD is a federated learning framework for training diffusion models efficiently and robustly under system and data heterogeneity. It is introduced as “Federated Pruning with Hierarchical Learning of Diffusion Models” (Long et al., 8 Jul 2025) and combines a three-tier hierarchical federated learning architecture, homogeneity-aware aggregation and client-to-edge selection, and structured pruning of diffusion-model U-Nets. The method targets three core difficulties in federated diffusion-model training: excessive communication, non-IID client data, and constrained compute and memory at the edge. In the reported experiments on CIFAR-10 and CelebA, FedPhD achieves strong image-generation quality in terms of Fréchet Inception Distance while reducing communication costs by up to $88\%$, improving FID by at least $34\%$ over baseline methods, and using about $56\%$ of the total computation and communication resources relative to baselines (Long et al., 8 Jul 2025).

1. Problem setting and motivation

Federated learning trains models over distributed clients’ data without sharing raw data. The FedPhD formulation focuses on diffusion models, especially U-Net-based denoisers, whose training in federated environments is difficult for reasons that are also present in federated training of Transformers and convolutional neural networks but are especially acute here (Long et al., 8 Jul 2025).

The first challenge is communication cost. Diffusion models, including the U-Net backbone considered in the paper, are large, with tens of millions of parameters. In standard federated learning, such as FedAvg, repeated transmission of full models over many rounds is expensive. Increasing aggregation frequency can reduce drift under non-IID data, but it correspondingly increases communication overhead (Long et al., 8 Jul 2025).

The second challenge is heterogeneity. Statistical heterogeneity appears as non-IID label distributions across clients, which induces gradient or weight divergence and degrades global convergence. The paper further states that, unlike supervised tasks, diffusion-model training is denoising-based, and drift-correction methods such as SCAFFOLD are less effective in this setting. System heterogeneity appears through variation in client compute, bandwidth, and memory (Long et al., 8 Jul 2025).

The third challenge is local resource limitation. U-Net-based diffusion models are described as memory- and compute-intensive, so edge devices may struggle to store, train, and transmit dense models. This increases training latency and aggravates straggler effects (Long et al., 8 Jul 2025).

FedPhD is presented as filling a gap relative to prior work. Existing federated learning approaches for diffusion models are said either to freeze subsets of parameters, which reduces communication but not heterogeneity, or to rely on data sharing, which is not privacy-preserving. Compression methods for diffusion models such as pruning, quantization, and distillation are characterized as largely assuming centralized, pre-trained teachers. FedPhD is designed to address these limitations jointly through hierarchical aggregation, homogeneity-aware weighting and selection, and coordinated structured pruning (Long et al., 8 Jul 2025).

2. Hierarchical federated learning architecture

The core architectural design of FedPhD is a three-tier hierarchical federated learning system: Clients $\rightarrow$ Edge servers $\rightarrow$ Global server (Long et al., 8 Jul 2025). Clients perform local diffusion-model training and send updates to edge servers. Edge servers aggregate client updates frequently, typically every $r_e = 1$ round, maintain accumulated label distributions, and return edge models to clients. The global server aggregates edge models less frequently, for example every $r_g = 5$ rounds, performs structured pruning after $R_s$ sparse rounds or at initialization in the one-shot setting, and then redistributes the pruned global model (Long et al., 8 Jul 2025).

This hierarchy is intended to separate frequent local stabilization from less frequent central synchronization. Frequent edge aggregation is used to reduce weight divergence caused by non-IID data while periodic global aggregation controls the cost of central communication. A plausible implication is that the architecture treats edge servers as an intermediate statistical smoothing layer, reducing both client drift and backbone transmission volume.

The standard federated objective is written as

$$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$

In the hierarchical formulation, the objective becomes

$$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$

where $34\%$0 is the number of selected clients at edge $34\%$1 (Long et al., 8 Jul 2025).

The reported training loop proceeds by global rounds. Each client first selects an edge server probabilistically. At each edge, the current accumulated distribution $34\%$2 is broadcast, clients perform local diffusion-model training for $34\%$3 epochs, and the edge aggregates client models when the edge-aggregation interval is met. When the global aggregation interval is met, the global server collects edge distributions and models, computes edge-level homogeneity scores, aggregates the edge models, and if $34\%$4, applies structured pruning at ratio $34\%$5 before redistribution (Long et al., 8 Jul 2025).

3. Diffusion-model formulation

FedPhD trains DDPM/DDIM-style diffusion models with a U-Net noise predictor $34\%$6 (Long et al., 8 Jul 2025). The forward diffusion process is given by

$34\%$7

with closed form

$34\%$8

where $34\%$9 and $56\%$0 (Long et al., 8 Jul 2025). The reverse denoising parameterization is

$56\%$1

Training uses the DDPM-like denoising objective

$56\%$2

where

$56\%$3

The paper also states a common DDPM sampling parameterization,

$56\%$4

with $56\%$5, and notes that DDIM uses a deterministic reverse process with fewer steps:

$56\%$6

FedPhD primarily trains DDIM with $56\%$7 timesteps for efficiency (Long et al., 8 Jul 2025).

The backbone is a U-Net “as in Ho et al. (2020)” with $56\%$8 million parameters in dense form (Long et al., 8 Jul 2025, Ho et al., 2020). The paper states that there are no major architectural deviations beyond pruning. This places the innovation of FedPhD in the federated optimization, aggregation, and compression scheme rather than in a new denoiser architecture.

4. Homogeneity-aware aggregation and client-to-edge selection

A defining component of FedPhD is its use of Statistical Homogeneity (SH) to guide both aggregation weights and client-to-edge assignment (Long et al., 8 Jul 2025). For client $56\%$9, let $\rightarrow$0 denote the empirical label distribution over labels $\rightarrow$1, and let $\rightarrow$2 denote the target distribution, often uniform. The client-level SH score is defined as

$\rightarrow$3

Edge servers maintain accumulated distributions. If $\rightarrow$4 denotes the clients attached to edge $\rightarrow$5, then the updated edge distribution is

$\rightarrow$6

where $\rightarrow$7 is the sample count contributed by client $\rightarrow$8 during the current period and $\rightarrow$9 is the edge’s current attached sample count (Long et al., 8 Jul 2025). The edge-level SH score is then

$\rightarrow$0

Global aggregation weights combine sample size and homogeneity:

$\rightarrow$1

with

$\rightarrow$2

Edge aggregation of client models uses the analogous form

$\rightarrow$3

where

$\rightarrow$4

The coefficients $\rightarrow$5 and $\rightarrow$6 regulate the trade-off between SH and sample size (Long et al., 8 Jul 2025).

Client-to-edge selection is also SH-aware. Client $\rightarrow$7 selects edge $\rightarrow$8 with probability

$\rightarrow$9

where $r_e = 1$0 and $r_e = 1$1 are the SH score and attached sample count after hypothetically adding client $r_e = 1$2’s data (Long et al., 8 Jul 2025). This explicitly favors edges that are closer to the target distribution and less loaded.

The paper reports that homogeneity-aware edge selection increases SH at edges and balances client loads, with lower variance than random selection, thereby improving convergence and stability (Long et al., 8 Jul 2025). Formal convergence bounds are not derived. The stated theoretical position is therefore empirical and mechanism-based rather than theorem-driven: frequent hierarchical aggregation reduces weight divergence, while SH-aware weighting and selection bias the effective training distribution toward target labels (Long et al., 8 Jul 2025).

5. Federated structured pruning

FedPhD integrates structured pruning into the federated hierarchy rather than applying compression only as a post hoc centralized step (Long et al., 8 Jul 2025). The pruning is structured at the level of channels, groups, or blocks rather than unstructured sparsity, and is intended to preserve hardware efficiency and reduce payload size. A dependency graph, denoted DepGraph, constructs parameter groups $r_e = 1$3 across the U-Net for coordinated pruning (Long et al., 8 Jul 2025).

Two pruning strategies are described. The first is one-shot (OS) pruning before training, intended for resource-limited clients. It uses magnitude-based criteria such as $r_e = 1$4 or group norms to choose groups to prune at a fixed ratio $r_e = 1$5. The second performs pruning after sparse training rounds, during which a group-norm regularizer shapes groups so that they can later be pruned according to group magnitudes (Long et al., 8 Jul 2025).

For client $r_e = 1$6, the sparse training objective is

$r_e = 1$7

with group-norm regularization

$r_e = 1$8

Layer-aware weighting is used to favor pruning in the mid U-Net layers, which are said to exhibit more redundancy. The group score is

$r_e = 1$9

and the regularization coefficient is set as

$r_g = 5$0

with $r_g = 5$1 tuned by grid search (Long et al., 8 Jul 2025).

Pruning is coordinated at the global server either after $r_g = 5$2 sparse rounds or at initialization for OS pruning. The global server then distributes the pruned mask or pruned model to all edges and clients, ensuring a consistent sparse structure across the hierarchy (Long et al., 8 Jul 2025). This design is central to FedPhD’s claim that pruning reduces transmitted parameters per update as well as local multiply–accumulate counts after pruning.

The reported ablation indicates that pruning up to approximately $r_g = 5$3 causes minimal degradation in FID and IS on CIFAR-10, whereas more aggressive pruning substantially harms generation quality, with FID rising to $r_g = 5$4 at $r_g = 5$5 (Long et al., 8 Jul 2025). This suggests a bounded operating region in which structural redundancy can be removed without severe denoising degradation.

6. Empirical evaluation, efficiency, and limitations

The experiments use CIFAR-10 and CelebA. CIFAR-10 comprises $r_g = 5$6k images at $r_g = 5$7, and CelebA comprises $r_g = 5$8k images at $r_g = 5$9 (Long et al., 8 Jul 2025). The model is DDIM with $R_s$0 steps and the $R_s$1M-parameter U-Net. The hierarchy uses $R_s$2 clients and $R_s$3 edges, with $R_s$4 and $R_s$5. Non-IID partitioning is severe: CIFAR-10 clients each hold $R_s$6 classes, while CelebA clients each hold one of four attribute classes, defined as young/old $R_s$7 male/female (Long et al., 8 Jul 2025).

Optimization uses Adam, with batch size $R_s$8 for CIFAR-10 and $R_s$9 for CelebA, and learning rates $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$0 and $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$1, respectively. The distributed baselines do not use EMA; EMA appears only in centralized comparisons (Long et al., 8 Jul 2025). Baselines are FedAvg, FedProx, FedDiffuse, MOON, and SCAFFOLD. The evaluation metrics are FID and Inception Score for quality, and parameter count, MACs, and communication volume per central aggregation for efficiency (Long et al., 8 Jul 2025). FID is computed between real and generated InceptionV3 features over $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$2k generated samples with batch size $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$3, according to

$$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$4

The main reported outcomes are summarized below.

Setting	Reported result
CIFAR-10, non-IID	FedPhD achieves FID $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$5, IS $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$6
CelebA, non-IID	FedPhD achieves FID $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$7–$$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$8, IS $$\min_\theta F(\theta) = \sum_{n=1}^N \rho_n F_n(\theta), \qquad \rho_n = \frac{D_n}{D}, \quad D=\sum_n D_n.$$9–$$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$0; OS is best at $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$1
FedAvg comparison	FedAvg reports $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$2 on CIFAR-10 and $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$3 on CelebA
Overall gains	At least $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$4 improvement in FID versus baselines under comparable budgets
Resource efficiency	Uses about $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$5 of total computation and communication resources relative to baselines
Communication reduction	Up to $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$6 reduction in communication cost

On CIFAR-10 at $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$7, parameters are reduced from $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$8M to $$\min_\theta F(\theta) = \sum_{e=1}^{N_e} \sum_{n=1}^{M_e} \rho_{en} F_{en}(\theta),$$9M, MACs from $34\%$00G to $34\%$01G, and model size from about $34\%$02 MB to about $34\%$03 MB (Long et al., 8 Jul 2025). Communication volume per central aggregation falls from $34\%$04 GB for FedAvg to $34\%$05 GB for FedPhD (Long et al., 8 Jul 2025). The paper also gives a generic communication model,

$34\%$06

as well as link-cost expressions following ShapeFL assumptions:

$34\%$07

with $34\%$08 and $34\%$09 denoting transmitted volume (Long et al., 8 Jul 2025).

Scalability experiments with $34\%$10 and $34\%$11 clients under fixed data size show only modest FID increase for FedPhD, exemplified by approximately $34\%$12, while baselines degrade much more, exemplified by approximately $34\%$13 (Long et al., 8 Jul 2025). The paper interprets this as robustness to increasing heterogeneity.

Several limitations are stated explicitly. FedPhD does not provide formal convergence analysis for hierarchical federated learning with SH-aware weighting. Differential privacy and secure aggregation are proposed but not implemented. EMA synchronization across the hierarchy is not implemented. Selection parameters $34\%$14 and $34\%$15 require dataset-specific tuning. Aggressive pruning above $34\%$16 degrades generation quality. The method currently focuses on unconditional diffusion models (Long et al., 8 Jul 2025).

These limitations define the immediate extension space identified by the paper: adaptive pruning schedules and regrowth, personalized heads or adapters on top of a shared sparse backbone, secure aggregation and differentially private SH computation, drift correction specialized for denoising objectives, and EMA synchronization or momentum aggregation in hierarchical federated learning (Long et al., 8 Jul 2025). A plausible implication is that FedPhD should be understood not as a closed solution to federated diffusion modeling, but as a specific systems-and-optimization template for jointly managing heterogeneity, communication, and resource constraints.