---
title: 'FedPhD: Federated Diffusion Model Training'
url: https://www.emergentmind.com/topics/fedphd
type: topic
---

# FedPhD: Federated Diffusion Model Training

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” [2507.06449] 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 [2507.06449].

## 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 [2507.06449].

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 [2507.06449].

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 [2507.06449].

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 [2507.06449].

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 [2507.06449].

## 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 [2507.06449]. 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 [2507.06449].

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 \(M_e\) is the number of selected clients at edge \(e\) [2507.06449].

The reported training loop proceeds by global rounds. Each client first selects an edge server probabilistically. At each edge, the current accumulated distribution \(q_e(y)\) is broadcast, clients perform local diffusion-model training for \(E\) 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 \(r = R_s\), applies structured pruning at ratio \(s_p\) before redistribution [2507.06449].

## 3. Diffusion-model formulation

FedPhD trains DDPM/DDIM-style diffusion models with a U-Net noise predictor \(\epsilon_\theta\) [2507.06449]. The forward diffusion process is given by
$$
q(x_t \mid x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}\,x_{t-1}, \beta_t I),
$$
with closed form
$$
q(x_t \mid x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t}\,x_0, (1-\bar{\alpha}_t)I),
$$
where \(\alpha_t = 1-\beta_t\) and \(\bar{\alpha}_t = \prod_{s=1}^t \alpha_s\) [2507.06449]. The reverse denoising parameterization is
$$
p_\theta(x_{t-1} \mid x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t,t), \Sigma_\theta(x_t,t)).
$$

Training uses the DDPM-like denoising objective
$$
L(\theta) = \mathbb{E}_{t,x_0,\epsilon}\big[\|\epsilon - \epsilon_\theta(x_t,t)\|_2^2\big],
$$
where
$$
x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,I).
$$
The paper also states a common DDPM sampling parameterization,
$$
x_{t-1} = \frac{1}{\sqrt{\alpha_t}}
\left[x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t,t)\right] + \sigma_t z,
$$
with \(z \sim \mathcal{N}(0,I)\), and notes that DDIM uses a deterministic reverse process with fewer steps:
$$
x_{t-1} = \sqrt{\bar{\alpha}_{t-1}}\,x_0 + \sqrt{1-\bar{\alpha}_{t-1}}\,\epsilon_\theta(x_t).
$$
FedPhD primarily trains DDIM with \(100\) timesteps for efficiency [2507.06449].

The backbone is a U-Net “as in Ho et al. (2020)” with \(35.7\) million parameters in dense form [2507.06449; 2006.11239]. 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 [2507.06449]. For client \(n\), let \(q_n(y)\) denote the empirical label distribution over labels \(\mathcal{Y}\), and let \(q_u(y)\) denote the target distribution, often uniform. The client-level SH score is defined as
$$
\mu_n = 2 - \sqrt{\sum_{y \in \mathcal{Y}} |q_n(y) - q_u(y)|^2 }.
$$

Edge servers maintain accumulated distributions. If \(\mathcal{M}_e\) denotes the clients attached to edge \(e\), then the updated edge distribution is
$$
q'_e(y) = \frac{q_e(y) + \sum_{n \in \mathcal{M}_e} q_n(y)\cdot n_n}{n_e + \sum_{n \in \mathcal{M}_e} n_n},
$$
where \(n_n\) is the sample count contributed by client \(n\) during the current period and \(n_e\) is the edge’s current attached sample count [2507.06449]. The edge-level SH score is then
$$
\mu_e = 2 - \sqrt{\sum_{y \in \mathcal{Y}} |q_e(y) - q_u(y)|^2 }.
$$

Global aggregation weights combine sample size and homogeneity:
$$
\theta = \sum_{e=1}^{N_e} \rho_e \cdot \theta_e,
$$
with
$$
\rho_e = \frac{\operatorname{ReLU}(n_e + a\mu_e + b)}
{\sum_{e'} \operatorname{ReLU}(n_{e'} + a\mu_{e'} + b)}.
$$
Edge aggregation of client models uses the analogous form
$$
\theta_e = \sum_{n \in \mathcal{M}_e} \rho_{en} \cdot \theta_n,
$$
where
$$
\rho_{en} = \frac{\operatorname{ReLU}(n_n + a\mu_n + b)}
{\sum_{m \in \mathcal{M}_e} \operatorname{ReLU}(n_m + a\mu_m + b)}.
$$
The coefficients \(a\) and \(b\) regulate the trade-off between SH and sample size [2507.06449].

Client-to-edge selection is also SH-aware. Client \(n\) selects edge \(e\) with probability
$$
P_n(e) = \frac{\operatorname{ReLU}(a \cdot \mu_e^{n'} - n_e^{n'} + b)}
{\sum_{j=1}^{N_e} \operatorname{ReLU}(a \cdot \mu_j^{n'} - n_j^{n'} + b)},
$$
where \(\mu_e^{n'}\) and \(n_e^{n'}\) are the SH score and attached sample count after hypothetically adding client \(n\)’s data [2507.06449]. 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 [2507.06449]. 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 [2507.06449].

## 5. Federated structured pruning

FedPhD integrates structured pruning into the federated hierarchy rather than applying compression only as a post hoc centralized step [2507.06449]. 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 \(\mathcal{G}=\{\theta^1,\dots,\theta^{|G|}\}\) across the U-Net for coordinated pruning [2507.06449].

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 \(L_2\) or group norms to choose groups to prune at a fixed ratio \(s_p\). 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 [2507.06449].

For client \(n\), the sparse training objective is
$$
F_n(\theta) = \frac{1}{D_n}\sum_{x \in \mathcal{D}_n}\mathcal{L}(\theta;x) + \Omega(\mathcal{G},k),
$$
with group-norm regularization
$$
\Omega(\mathcal{G},k) = \sum_{k=1}^K \sum_{\theta^g \in \mathcal{G}} \lambda_g \|\theta^g[k]\|_2^2.
$$
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
$$
Q(\theta^g) = \frac{1}{L}\sum_{l \in \theta^g} |l - l_{\text{mid}}|,
$$
and the regularization coefficient is set as
$$
\lambda_g = \frac{\lambda_0}{Q(\theta^g)},
$$
with \(\lambda_0\) tuned by grid search [2507.06449].

Pruning is coordinated at the global server either after \(R_s\) 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 [2507.06449]. 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 \(s_p \approx 44\%\) causes minimal degradation in FID and IS on CIFAR-10, whereas more aggressive pruning substantially harms generation quality, with FID rising to \(22.34\) at \(s_p = 74\%\) [2507.06449]. 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 \(50\)k images at \(32\times 32\), and CelebA comprises \(163\)k images at \(64\times 64\) [2507.06449]. The model is DDIM with \(T=100\) steps and the \(35.7\)M-parameter U-Net. The hierarchy uses \(N=20\) clients and \(N_e=2\) edges, with \(r_e=1\) and \(r_g=5\). Non-IID partitioning is severe: CIFAR-10 clients each hold \(2\) classes, while CelebA clients each hold one of four attribute classes, defined as young/old \(\times\) male/female [2507.06449].

Optimization uses Adam, with batch size \(128\) for CIFAR-10 and \(64\) for CelebA, and learning rates \(2\times 10^{-4}\) and \(2\times 10^{-5}\), respectively. The distributed baselines do not use EMA; EMA appears only in centralized comparisons [2507.06449]. 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 [2507.06449]. FID is computed between real and generated InceptionV3 features over \(30\)k generated samples with batch size \(256\), according to
$$
\mathrm{FID} = \|\mu_r - \mu_g\|_2^2 + \operatorname{Tr}\left(\Sigma_r + \Sigma_g - 2(\Sigma_r \Sigma_g)^{1/2}\right).
$$

The main reported outcomes are summarized below.

| Setting | Reported result |
|---|---|
| CIFAR-10, non-IID | FedPhD achieves FID \(16.74\), IS \(4.24\) |
| CelebA, non-IID | FedPhD achieves FID \(8.32\)–\(7.48\), IS \(2.76\)–\(2.88\); OS is best at \(7.48/2.88\) |
| FedAvg comparison | FedAvg reports \(21.63/4.16\) on CIFAR-10 and \(11.29/2.81\) on CelebA |
| Overall gains | At least \(34\%\) improvement in FID versus baselines under comparable budgets |
| Resource efficiency | Uses about \(56\%\) of total computation and communication resources relative to baselines |
| Communication reduction | Up to \(88\%\) reduction in communication cost |

On CIFAR-10 at \(s_p=44\%\), parameters are reduced from \(35.7\)M to \(20.3\)M, MACs from \(6.06\)G to \(3.42\)G, and model size from about \(136.5\) MB to about \(77.9\) MB [2507.06449]. Communication volume per central aggregation falls from \(109.22\) GB for FedAvg to \(65.46\) GB for FedPhD [2507.06449]. The paper also gives a generic communication model,
$$
C_{\text{comm}} \approx R \cdot (1-s_p)\cdot |\theta| \cdot b,
$$
as well as link-cost expressions following ShapeFL assumptions:
$$
C_{ne} = 0.002 \cdot d_e \cdot V, \qquad C_{ce} = 0.02 \cdot d_c \cdot V,
$$
with \(d_c \approx 10 \cdot d_e\) and \(V\) denoting transmitted volume [2507.06449].

Scalability experiments with \(N=50\) and \(N=100\) clients under fixed data size show only modest FID increase for FedPhD, exemplified by approximately \(+0.4\), while baselines degrade much more, exemplified by approximately \(+3.5\) [2507.06449]. 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 \(a\) and \(b\) require dataset-specific tuning. Aggressive pruning above \(44\%\) degrades generation quality. The method currently focuses on unconditional diffusion models [2507.06449].

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 [2507.06449]. 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.

Source: https://www.emergentmind.com/topics/fedphd