CollaFuse: Collaborative Diffusion Models

TL;DR

This work introduces a novel approach for distributed collaborative diffusion models inspired by split learning, which facilitates collaborative training of diffusion models while alleviating client computational burdens during image synthesis.

Abstract

In the landscape of generative artificial intelligence, diffusion-based models have emerged as a promising method for generating synthetic images. However, the application of diffusion models poses numerous challenges, particularly concerning data availability, computational requirements, and privacy. Traditional approaches to address these shortcomings, like federated learning, often impose significant computational burdens on individual clients, especially those with constrained resources. In response to these challenges, we introduce a novel approach for distributed collaborative diffusion models inspired by split learning. Our approach facilitates collaborative training of diffusion models while alleviating client computational burdens during image synthesis. This reduced computational burden is achieved by retaining data and computationally inexpensive processes locally at each client while outsourcing the computationally expensive processes to shared, more efficient server resources. Through experiments on the common CelebA dataset, our approach demonstrates enhanced privacy by reducing the necessity for sharing raw data. These capabilities hold significant potential across various application areas, including the design of edge computing solutions. Thus, our work advances distributed machine learning by contributing to the evolution of collaborative diffusion models.

Figures (7)

• Figure 1: Overview of CollaFuse for collaborative image synthesis by splitting the denoising process between the server and clients. Based on client-specific conditioning $y_i$, the first $T-t_\zeta$ denoising steps are run on a trusted server model, while the remaining $t_\zeta$ denoising steps are run locally by the $i$-th client. Thereby, external computing resources can be utilized while keeping the clients' raw data private.
• Figure 2: Comparison between random samples from the training set (top row) and images generated with our collaborative diffusion models trained with cut point $t_\zeta=100$ (bottom row). Images were not cherry-picked and generated starting with the same initial noise. The results demonstrate that collaboratively training diffusion models can achieve high image quality and attribute fidelity.
• Figure 3: Fidelity evaluation of clients and server using FID$\downarrow$ and FCD$\downarrow$: We assess 10,000 / 10,000 / 3,730 real $x_0$ and collaboratively generated images ($\hat{x}_0^c \circ \hat{x}^s_{t_\zeta}(\epsilon)$) from the CIFAR-10 / CelebA / Awa2 dataset across clients. Our findings show that cut points with $t_\zeta \leq 200$ outperform the baseline of independent client models (ICMs) ($t_\zeta = 1000$). Notably, smaller cut points even surpass the global model (GM) ($t_\zeta = 0$), which is trained on all client datasets. In particular, training with smaller images or latent sizes benefits from CollaFuse. As expected, information disclosure almost consistently decreases as the cut point ($t_\zeta$) increases across all models.
• Figure 4: Samples generated by CollaFuse trained on CelebA with different cut points $t_\zeta$. The top row depicts images produced by the server, which are then sent to the client. The bottom row shows the samples after the final denoising performed by the client. For $t_\zeta=0$, the server performs the full denoising process, for $t_\zeta=1000$, each client trains a separate diffusion model without a server component.
• Figure 5: The figure shows generated images exemplarily for different scenarios across different cut points. Column I describes server-only generated images, while column II describes client-only generated images. Column III shows images at the cut point, while column IV shows images generated collaboratively.