SPIRE: Conditional Personalization for Federated Diffusion Generative Models

TL;DR

SPIRE introduces a parameter-efficient federated diffusion model personalization framework that separates a global backbone from lightweight client embeddings, enabling effective adaptation and reducing forgetting.

Contribution

It presents a novel framework for federated diffusion models using shared backbones and client embeddings, with theoretical analysis linking conditional diffusion to maximum likelihood estimation.

Findings

SPIRE matches or surpasses strong baselines in collaborative pretraining.

It significantly outperforms baselines in adapting to unseen clients.

It reduces Kernel Inception Distance while updating only hundreds of parameters.

Abstract

Recent advances in diffusion models have revolutionized generative AI, but their sheer size makes on device personalization, and thus effective federated learning (FL), infeasible. We propose Shared Backbone Personal Identity Representation Embeddings (SPIRE), a framework that casts per client diffusion based generation as conditional generation in FL. SPIRE factorizes the network into (i) a high capacity global backbone that learns a population level score function and (ii) lightweight, learnable client embeddings that encode local data statistics. This separation enables parameter efficient finetuning that touches $\leq 0.01\%$ of weights. We provide the first theoretical bridge between conditional diffusion training and maximum likelihood estimation in Gaussian mixture models. For a two component mixture we prove that gradient descent on the DDPM with respect to mixing weights loss recovers the optimal mixing weights and enjoys dimension free error bounds. Our analysis also hints at how client embeddings act as biases that steer a shared score network toward personalized distributions. Empirically, SPIRE matches or surpasses strong baselines during collaborative pretraining, and vastly outperforms them when adapting to unseen clients, reducing Kernel Inception Distance while updating only hundreds of parameters. SPIRE further mitigates catastrophic forgetting and remains robust across finetuning learning rate and epoch choices.