FedDM-prox: Federated Diffusion Training

• The paper demonstrates that incorporating a proximal regularizer in the local diffusion loss reduces client drift, yielding up to a 23% FID improvement under moderate data skew.
• The method leverages U-Net-backed diffusion architectures and standard federated averaging without extra communication cost, ensuring seamless integration across clients.
• Empirical results on CIFAR-10 and CelebA show that FedDM-prox maintains image fidelity and sharpness under heterogeneous data conditions compared to baseline approaches.

Federated Diffusion Model Training with Proximal Regularization (FedDM-prox) is a methodological advancement designed to address the unique challenges of training diffusion models in federated settings, particularly under client data heterogeneity. This approach is situated within the FedDM suite of algorithms, which leverage U-Net-backed diffusion architectures and address both communication efficiency and robust model convergence when client data distributions (non-IID) diverge significantly. FedDM-prox incorporates a proximal term in the local client objective to mitigate client-drift, thus enhancing convergence and model fidelity across disparate data partitions (Vora et al., 2024).

1. Federated Diffusion Model Framework

The core federated learning objective is to optimize a global diffusion model $p_\theta(x_0)$ over $K$ participating clients. Each client $k$ holds a private dataset $D_k$ sampled from distribution $q_k(x_0)$. The global optimization is expressed as:

$$\min_\theta F(\theta) \triangleq \sum_{k} w_k F_k(\theta)$$

where $w_k = |D_k| / \sum_j |D_j|$, and the local objective for each client,

$$F_k(\theta) \triangleq \mathbb{E}_{x_0 \sim q_k(x_0)} \left[ \sum_{t=1}^T \mathbb{E}_{q_k(x_t|x_{t-1})}\left[ -\log p_\theta(x_{t-1}|x_t) \right] \right]$$

corresponds to the variational diffusion-model loss. Optimization proceeds in communication rounds, where a subset of clients solve their local subproblems—typically via SGD—before sending updated parameters to a federating server for weighted aggregation.

2. Proximal Objective for Mitigating Client Drift

FedDM-prox directly addresses instability and slow convergence arising from non-IID client data by extending the local diffusion loss with a proximal regularizer centered at the current global model parameter $\theta_r$. The modified per-client optimization in round $r$ is:

$$L_k^\text{prox}(\theta) \triangleq F_k(\theta) + \frac{\mu}{2}\|\theta - \theta_r\|_2^2$$

where $\mu \geq 0$ is a tunable proximal coefficient controlling the strength of regularization. This term constrains local updates, reducing their drift from the global model and consequently improving aggregation stability and convergence—particularly under severe data partition skew.

3. FedDM-prox Algorithmic Workflow

FedDM-prox is instantiated as follows:

1. Initialization: Start with $\theta^0$.
2. Selection & Broadcast: In each round $r$, randomly sample a subset $S_r$ of $k$ clients; broadcast global model $\theta_r$.
3. Local Update (per client $i$):
   • Initialize $\theta_i \leftarrow \theta_r$.
   • For $E$ local epochs:
     • Sample minibatch $B$ from $D_i$.
     • Compute stochastic gradient of local diffusion loss.
     • Apply proximal-regularized SGD:

   $$\theta_i \leftarrow \theta_i - \eta \left[\nabla F_i(\theta_i) + \mu (\theta_i - \theta_r)\right]$$

• Return updated $\theta_i$.

1. Aggregation: Server updates global parameters by weighted averaging:

$$\theta_{r+1} \leftarrow \sum_{i \in S_r} \frac{|D_i|}{\sum_{j \in S_r} |D_j|} \theta_i$$

Repeat for $R$ rounds.

This method retains the same communication cost as FedDM-vanilla, as the proximal regularizer requires no additional communication nor does it increase the model size. Inclusion of the proximal step does increase per-round local computational complexity.

4. Convergence Analysis of FedDM-prox

The convergence guarantee of FedDM-prox extends theoretical results for non-convex federated learning. Under the assumptions:

• Each $F_k(\cdot)$ is $L$-smooth,
• Stochastic-gradient noise is bounded with variance $\sigma^2$,
• Local learning rate $0 < \eta \leq 1/L$ and $\mu > 0$,

then after $R$ communication rounds, the expected squared gradient norm satisfies:

$$\mathbb{E}\left[\|\nabla F(\theta)\|^2\right] = O\left(\frac{F(\theta^0) - F^*}{\eta R} + \frac{\sigma^2}{\mu R} + \mu\right)$$

This yields $O(1/R)$ convergence to a stationary point. The proximal term effectively regulates the deviation between local and global models, particularly critical when data heterogeneity is high (Vora et al., 2024).

5. Diffusion Model Backbone and Federated Specifics

FedDM algorithms, including FedDM-prox, utilize U-Net architectures as denoisers ($\epsilon_\theta(\cdot, t)$), as introduced in Ho et al. 2020:

• Encoder–decoder path with skip connections;
• GroupNorm and Swish activations;
• Self-attention at intermediate resolutions.

In DDPMs, the diffusion process operates directly in image space over $T = 1000$ timesteps ($\beta_t$ linear from $1 \times 10^{-4}$ to $2 \times 10^{-2}$). In Latent Diffusion Models (LDMs), an encoder maps $x_0$ to latent $z_0$, diffusion operates in latent space, then $z_0$ is decoded to pixels. Only model weights and batch statistics are client-local during federated training; the protocol for communication and aggregation is otherwise identical.

6. Hyperparameter Choices and Empirical Findings

Typical hyperparameter settings for FedDM-prox in experiments are:

Parameter	Value
Total clients ($K$)	10
Per-round clients ($k$)	6
Global rounds ($R$)	20
Local epochs ($E$)	15
Batch size	128
Learning rate ($\eta$)	$1 \times 10^{-4}$
Proximal coeff. ($\mu$)	$0.1$
Diffusion steps ($T$)	$1000$, $\beta$ as above

Empirical evaluation on CIFAR-10 under varying levels of label skew demonstrates robust improvements over FedDM-vanilla:

Skew Level	FedDM-vanilla FID	FedDM-prox FID
Low	7.24	6.10
Moderate	10.8	8.3
High	15.6	11.1

Under moderate skew, FedDM-prox reduces FID by approximately 23%. Communication remains approximately $65$ MB per round. Reference comparison: FedDM-quant achieves $4\times$ reduction (down to $\sim16$ MB/round) with a marginal FID increase.

7. Trade-Offs, Recommendations, and Qualitative Results

FedDM-prox achieves substantial enhancement in non-IID robustness without additional communication cost, at the price of increased local computational load due to the proximal term. Improved FID (up to $40\%$ lower than FedDM-vanilla) is observed particularly under severe heterogeneity. In contrast, FedDM-quant enables $4\times$ bandwidth reduction with only minor FID degradation, but necessitates a per-client calibration routine.

Comparison summary:

Algorithm	Comm. Cost/round	FID (CIFAR-10, moderate skew)	Notes
FedDM-vanilla	$\sim65$ MB	$10.8$	Standard federated averaging
FedDM-prox	$\sim65$ MB	$8.3$	Stronger on non-IID, no extra comm.
FedDM-quant	$\sim16$ MB	$7.06$	8-bit quantization, needs calibration

Qualitative results (e.g., on CelebA) show that FedDM-prox preserves image sharpness and detailed features under non-IID conditions, in contrast to the degradation observed in unregularized federated training. FID versus communication rounds demonstrates the steadier and ultimately superior convergence of FedDM-prox under high skew, with lower plateau values compared to baseline.

FedDM-prox is recommended for federated diffusion model scenarios where client data heterogeneity is the primary concern and bandwidth constraints are moderate; for cases where bandwidth is the primary concern, FedDM-quant offers a more appropriate trade-off (Vora et al., 2024).