Enhancing Privacy in ControlNet and Stable Diffusion via Split Learning

Diffusion Models [1] are probabilistic models used to learn a data distribution by gradually denoising a normally distributed variable to generate high-quality images. The stable diffusion model [10] converts images into latent representations with an encoder $\mathcal{E}$ and conducts the diffusion process on the latent domain $Z$. After the sampling process, images are generated through a corresponding decoder $\mathcal{D}$.

Existing diffusion models allow users to guide image generation with text prompts. For example, in a stable diffusion model, we utilize a contrastive language-image pretraining (CLIP) model [11], with parameters $\gamma_{\theta}$, to convert text prompts into features. A cross-attention layer is then employed to combine these features with latent representations.

The image generation process involves a sampling procedure, which is the inverse of the forward process depicted in Fig. 1. In the forward process, we follow a Markov Chain to gradually add Gaussian noise ($\mathcal{N}\sim(0,1)$) to the data, based on a variance schedule $\beta_{1},\ldots,\beta_{T}$, where $t\in[T]$ represents each timestep of noise addition. We denote this Gaussian noise as $\hat{n}$. As an inverse of the diffusion process, during the sampling process, the diffusion model outputs an estimation of noise $n$ at timestep $t$, and we sample the latent $z_{t-1}$ using the equation:

	$\displaystyle z_{t-1}=$	$\displaystyle\sqrt{\alpha_{t-1}}\left(\frac{z_{t}-\sqrt{1-\alpha_{t}}n}{\sqrt{\alpha_{t}}}\right)$		(1)
		$\displaystyle+\sqrt{1-\alpha_{t-1}-\lambda_{t}^{2}}n+\lambda_{t}o(z_{t})$

Here, $\alpha_{t}=1-\beta_{t}$, $\lambda_{t}$ is a noise coefficient, and $o(z_{t})$ is a small value randomly generated from a standard normal distribution. The sampling process begins with randomized Gaussian noise and gradually samples until we obtain $z_{0}$, which corresponds to the latent representation of the image we wish to generate.

During the training process, we uniformly sample timesteps from $[T]$ and train the diffusion model (DM) with given text prompts $y$, aiming to minimize the loss:

$$\min\mathcal{L}_{\rm DM}=\mathbb{E}_{\mathcal{E}(x),y,n\sim\mathcal{N}(0,1),t}\left[\lVert\hat{n}-n_{\theta}(z_{t},t,\gamma_{\theta}(y))\rVert^{2}_{2}\right]$$

(2)

In each training step, we need to follow Eq. 1 to generate a random noise $\hat{n}$ as a label, which serves as the ground truth. The diffusion model’s objective is to learn the parameters $\theta$, which enable it to infer the noise $n$. This inferred noise is the output of the diffusion model, used for denoising the image.

To enable users to control the generated images with more detailed conditions such as scribbles [2], canny lines [12], depth maps [13], HED lines [14], and segmentation maps [15], in addition to the given text prompts, a conditional diffusion model is proposed. An example is ControlNet [2], as shown in Fig. 3. We can generate a stormtrooper with the same skeletons as in the left image of the depth maps.

ControlNet copies the encoders from the backbone diffusion models and replaces the decoders of the backbone diffusion models with convolution layers initialized with zeros (referred to as zero convolution). A control network comprises copied encoders, zero convolutions, and a condition encoder for converting conditions into latent representations. ControlNet consists of this control network and the original stable diffusion model. For the detailed structure of ControlNet, we will explain it later in Section III-B.

Apart from stable diffusion models, ControlNet can also leverage other backbones such as LCM [16] and ControlLoRA [17]. Concurrent works, T2I-Adapter [18] and Composer [19], feature much smaller and much larger control networks, respectively. The control network can have other structures, such as those seen in T2IAdapter [18] and Composer [19]. Our method can also be applied to these works. FreeDoM [20] is a training-free conditional diffusion model. However, generating images with fine-grained conditions, such as using canny edge maps, can be challenging, resulting in poor guidance. Training-required methods still remain the optimal solution for conditional diffusion models.

With the assistance of ControlNet, users can fine-tune well-trained stable diffusion (SD) models without disrupting the original SD models. However, the conditions and training images involved may contain privacy-sensitive information. One straightforward solution is to train ControlNet entirely on a single device. For inference with a batch size of 1, we need 7.50 GB of GPU memory. However, to train ControlNet, a minimum of 23.82 GB of GPU memory is needed (with a minimal batch size of 2). Such high GPU memory requirements are unfeasible for most client devices, let alone mobile devices. Even if clients possess powerful computing resources and acceptable training times, training ControlNet on the client side remains impossible if the server is unwilling to share the well-trained diffusion model.

Even if a client has enough GPU memory to fine-tune a diffusion model locally, another issue arises when it needs to collect training samples from different users, as it may lack sufficient samples in its local data. To allow users to fine-tune ControlNets without their private data leaving their devices, a common solution is to leverage privacy-preserving decentralized frameworks. One such decentralized training paradigm is federated learning [4, 21]. We follow the standard federated averaging scheme to train the ControlNet with 50 clients, each having 1000 training samples. We train for a total of 100 rounds and aggregate weights after every 250 local iterations. We evaluate the performance on the MS-COCO [22] validation set. As shown in Fig. 3, even under the assumption that clients have powerful computing units and the weights of the diffusion model are available, the ControlNet trained by FedAvg [4] fails to learn the conditions. The generated image does not match the condition at all; for example, the posture of the person in the generated image differs from that in the left condition images. Since federated learning without any privacy-preserving mechanism cannot work, we do not need to further evaluate methods with privacy-preserving mechanisms [21].

There are also data encryption approaches in decentralized systems, such as trusted execution environments, multi-party computation, and homomorphic encryption. However, the overhead is not at the same scale as computing over plaintext data. For example, during inference, the forward time on diffusion models with homomorphic encryption [23] is 79.19 days, compared to 35 seconds with plaintext using NVIDIA A100. Moreover, to the best of our knowledge, there is no encryption method that can be directly applied to the training process of diffusion models.

Considering the challenges involved in training conditional diffusion models either on clients or servers, a suitable solution is to train such models through split learning, involving multiple clients and the server. Unlike federated learning, which pushes the entire model to the edge, split learning employs a neural network spanning both the cloud and the edge. An edge device trains the network up to the partition layer and sends the intermediate features to the server. Upon receiving these features, the server takes over training the remaining layers and completes forward propagation. During backward propagation, the server conducts back-propagation up to the partition point and sends the gradients of the partition layers back to the client. The client then updates the local parameters through back-propagation using the received gradients. A major drawback of split learning is its sequential training manner, resulting in significant resource underutilization and high transmission overhead, leading to longer training times. In each training step, the server and clients exchange features and gradients, with one party waiting while the other computes or transmits data.

Potential threats arise from split learning, as it carries the risk of privacy leakage through data transmission between clients and the server. Literature highlights that an honest-but-curious server could reconstruct private data using the intermediate features sent from clients to the server. Zhang et al. [24] successfully reconstructed private data in a white-box setting. UnSplit[9] further refined this method to conduct a similar attack in a black-box setting. He et al. [25] trained an inverse network using a public dataset, taking intermediate results as inputs to output private data for reconstruction. Conversely, Li et al. [26] demonstrated that private labels on clients also face the risk of leakage. Pasquini et al. [27] proposed an attack to reconstruct private data by manipulating gradients sent back to clients, under the assumption of a dishonest server. Duan et al. [28] introduced a membership inference attack (MIA) tailored specifically for diffusion models, although they acknowledged its limited applicability in real-world scenarios. Direct MIA is excluded from the scope of our paper. Additionally, Carlini et al. [29] utilized leaked text prompts to generate numerous images, subsequently employing MIA to identify which images exist in private datasets.

In response to potential privacy leakage in split learning, researchers have made efforts to defend against such attacks. Local differential privacy techniques, such as additive noise and randomized response [30], are employed to prevent reconstruction. Additionally, Gaussian noise is utilized to directly add noise to the raw data [8]. Subsequently, many works have adopted methods involving additive noise [31, 32, 33]. DataMix [8] and CutMix [34] leverage the concept of mixing a batch of samples. DataMix is designed for convolutional neural networks, while CutMix is tailored for vision image transformers. PatchShuffling [7, 35, 36] is a method specifically designed for transformer-structured models, where patches are shuffled among a batch of samples. However, all these methods must provide sufficient privacy guarantees at the cost of significant performance decreases.

Xiao et al. [6] utilized adversarial learning to enable clients to generate intermediate results that the server cannot use to reconstruct images. Shredder[33] introduced noise based on mutual information, while DISCO [37] employed a channel obfuscation method to process features before transmitting them to the server. However, all three of these methods are only applied during the inference stage and require training a network to process features.

Differential privacy was originally conceived to ensure that an adversary’s ability to compromise the privacy of any set of users remains unchanged by an individual’s decision to opt into or out of the dataset [38]. This characteristic ensures that an adversary cannot glean additional information about any specific individual, thereby extending its capability to prevent inversion attacks such as reconstructing private data.

We assume we have two sets $X$ and $X^{\prime}$, with only one sample difference between them. We say a privacy-preserving mechanism $\mathcal{M}$ is locally differentially private (LDP) if we cannot differentiate between $X$ and $X^{\prime}$ based on the outputs of these two sets by $\mathcal{M}$. We provide a formal definition here.

We call this $\epsilon$ privacy budget. In common cases, a larger privacy budget implies easier differentiation between the two probabilities, indicating weaker privacy-preserving ability, and vice versa.

We can apply a differential privacy mechanism to either model weights or features. DPDM [39] and DPGM [40] have implemented differential privacy mechanisms over the model weights of a diffusion model. However, since our focus is on split learning, we need to concentrate on mechanisms applied to the inputs or features. To achieve privacy protection against reconstructing private images, we can add noise to the original inputs or intermediate features [41], making it $(\epsilon,\Delta)$-LDP. The noise added can be Gaussian noise [30, 42].

We will later use this formulation of adding noise in the diffusion model to propose an $(\epsilon,\Delta)$-LDP mechanism over features. Randomized response [43] is another typical mechanism to achieve local differential privacy. It encodes each real value in a feature into bits and randomly flips each bit. A drawback of existing LDP mechanisms, whether applied to model weights or features, is that they aim for privacy preservation at the expense of the quality of generated images.

We first introduce the deployment when we want to fine-tune a diffusion model with ControlNet using split learning. Though this is engineering work to combine two existing frameworks, we would still like to provide the essential rationale behind how we decide the cut layers and model deployments over clients and the server. Different from the partition of dense models such as ResNet [44], which has the structure of blocks being sequentially placed, the conditional diffusion model contains two parts: a diffusion model with frozen weights and a trainable control network. The structure of the ControlNet is shown in Fig. 4. The control network will first process the condition images with a particular condition encoder (during condition encoding) trained from scratch and then mix it with the noisy latent representation as the input of the following blocks.

The diffusion model will first use a pre-trained encoder to convert the original image into the latent representation. The diffusion model has encoder blocks and decoder blocks. As they are all frozen, no parameters need to be updated. But the decoder blocks need to calculate gradients of parameters so as to update the control network parameters during backpropagation. Each diffusion model decoder will take the output of the corresponding encoder and the output of the corresponding block in the control network as the inputs. These inputs are fed into the decoder with a jump connection using a similar structure as the UNet [45].

For each encoder and decoder block in the original stable diffusion model and each encoder block in the control network, we will also put the text prompts and timestep as the inputs to realize the text-to-image generation. The condition encoder and the autoencoder $\mathcal{E}$ only take the conditions or original images as the inputs. During the forward process, the clients need to send text prompts as well as timesteps to the server.

Considering hiding the complete model weights of the well-trained diffusion models from the clients and achieving the best tradeoff between privacy and efficiency through choosing different partition points, we cut right after the first diffusion model encoder and the trainable condition encoder. If we cut deeper, the server still needs the output of the previous encoder blocks as the inputs of the decoders. This will not provide a better privacy guarantee but will increase the transmission and computation overhead on the clients.

We send the output of the model and gradients back to the clients In such a way, the images are generated on the clients. The output is the $n$ going to be used in Eq. 1 for denoising. It is an inferred random noise following the standard normal distribution. The server will not be able to generate images with only knowing $n$. Regarding the diffusion model, if we place the partition point before the first encoder block, the server can subtract the estimated noise $n$ from received $z_{t}$ in Eq. 1 to recover the $z_{0}$ and retrieve the private images.

In summary, practical applications of split learning face four threats. The first is a potential attack using gradient descents in a white-box scenario, particularly if clients utilize pre-trained weights. The second threat is an UnSplit attack, while the third involves training inverse networks to infer private data without prior knowledge of the data. The fourth threat is the leakage of text prompts.

We summarize potential split learning attacks in Table III and their effectiveness. The remaining effective method is inverse network-based attacks for reconstructing condition images. Another valid threat is the leakage of text prompts.

We have empirically shown that ControlNet itself is quite effective in defending against several attacks in split learning. If we look at its structure carefully, we will find that the forward process in Fig. 1 has already contained the process of adding noise over the latent representation. We will next show that such a mechanism is $(\epsilon,\Delta)-$LDP using the Definition III.2. Based on this property, we propose a new sampling scheme over timesteps during the diffusion process, preserving privacy.

With a given latent representation $z_{0}$, we will generate the noisy latent representation according to the timestep $t$, scheduling parameter $\beta_{t}$ and a randomly generated noise $\hat{n_{t}}\sim\mathcal{N}(0,1)$:

		$\displaystyle z_{t}=\sqrt{1-\beta_{t}}z_{0}+\sqrt{\beta_{t}}\hat{n_{t}}$		(5)
		$\displaystyle\frac{z_{t}}{\sqrt{1-\beta_{t}}}=z_{0}+\sqrt{\frac{\beta_{t}}{1-\beta_{t}}}\hat{n_{t}}$

According to Fig. 4, $z_{t}$ is the input to the first encoder block of stable diffusion model. Because $\beta_{t}$ is usually a small number, we approximate Eq. 5 as,

$$z_{t}\approx z_{0}+\sqrt{\frac{\beta_{t}}{1-\beta_{t}}}\hat{n_{t}}$$

(6)

We can view this equation as adding a noise following distribution $\mathcal{N}(0,\frac{\beta_{t}}{1-\beta_{t}})$ over $z_{0}$ to get $z_{t}$. We then substitute the variance in Definition III.2. For convenience, we notate $H$ as hyper-parameter $2\ln\frac{1.25}{\Delta}\alpha^{2}$.

		$\displaystyle\frac{\beta_{t}}{1-\beta_{t}}=H\cdot\frac{1}{\epsilon^{2}}$		(7)
		$\displaystyle\epsilon=\sqrt{H\cdot\frac{1-\beta_{t}}{\beta_{t}}}$

In the diffusion model, we employ the linear scheduling as the default method which is $\beta_{t}=k\cdot t+\beta_{0}$, where $k$, $\beta_{0}$ are scheduling parameters. Hence, we can derive the following relationship between privacy budget $\epsilon$ and $t$, $k$, and $\beta_{0}$:

$$\epsilon(t,k,\beta_{0})=\sqrt{H\cdot\left(\frac{1}{kt+\beta_{0}}-1\right)}$$

(8)

From this equation, we can see that the privacy budget is related to the timestep. However, to protect different types of images, we need to set different levels of privacy budgets. For example, during the inversion attack for the original image, even if the $t=0$, privacy can still be protected. However, for condition images which are simpler than detailed natural images, the budget at $t=0$ is not enough to protect privacy. In such a case, we need to set bigger privacy budgets. Fortunately, based on Eq. 8, we can set proper privacy budget by setting different $t,k,\beta_{0}$ in fine-tuning ControlNet.

However, as we can see from the Eq. 5, if we directly send the encoded condition mixed with the noisy latent representation to the server, as the server knows the timestep $t$ and the label $\hat{n}$, it can directly subtract the added noise from Eq. 5. As a result, based on observation of the success defense by the diffusion model part, we find that leveraging some functions, especially non-linear, after the noisy latent representation and before sending to the server will confuse the attacker from stealing privacy. Hence, we propose to add a noise-confounding activation layer before sending features to the server. To design an activation function keeping privacy-preserving property while maintaining the image generation performance, we pass the sum of the encoded condition and the noisy latent representation through such a function:

$$y=\lvert x\rvert\cdot\left(\frac{2}{1+e^{-x}}\right)+\delta$$

(9)

The $\delta$ is a randomized noise following distribution $\mathcal{N}\sim(0,1)$. The noise $\delta$ is randomized at the beginning of the training and fixed during the training. The server or the attacker has no access to $\delta$. The function graph of Eq. 9 is shown in Fig. 10. The functionality of this function is to prevent the attacker from inferring the sum of the latent representation and encoder condition. In order to maintain image generation performance, we adopt a symmetric design with an SiLU-like shape. The SiLU function [50] is a widely used activation function, which can help improve the performance of neural networks. As $\delta$ is fixed during the training, the quality of the summation will not be degraded.

Another solution is to put the SD Encoder Block 1 of the control network on the clients. However, such moving will increase the computation overhead on the clients and break the structure of gradient back-sending free. On the other hand, adding this activation function will help preserve privacy without adding any overhead. Besides, moving the place of SD encoder block is specific for ControlNet while adding activation function can also apply to other conditional diffusion model such as T2I-Adapter [18].

Apart from keeping the privacy of images, text prompts can also contain private information. In the original ControlNet, for each encoder block and decoder block in the diffusion model and control network, text prompts will be input into the blocks and attention modules will be applied in these blocks to let the generate images learn the text prompts. As a result, the clients have to upload the text prompts or the text features to the server. The server can get the raw text information and privacy will be leaked. Uploading the output features of text encoder may be able to let the server not get texts, but the server can still use the same method proposed by Carlini et al. [29] to extract training dataset using text features as inputs.

As a result, to hide prompts from the server, we propose to not send the text prompts to the server directly. During the training, only the SD Encoder Block 1 in Fig. 4 on the clients, will take in text prompts as the input. The text prompts will not be uploaded to the server. Therefore, other encoder and decoder blocks in ControlNet, situated on the server, won’t utilize text prompts as inputs. Removing text prompts will not affect the performance of condition and image encoders as they are irrelevant to prompts. However, the input distribution of server-side encoders and decoders has changed. To maintain high-quality image generation, we introduce the following prompts-hiding training methods.

As the diffusion model is frozen and able to always keep the generation performance of a well-trained diffusion model while the control network needs further training, we use different policies for the encoder and decoder blocks in the control network and diffusion model. For blocks in control networks, no text features will be input and the attention modules will be replaced by self-attention modules for the condition features. Since we still need to further fine-tune the control network, the distribution drift caused by removal of text input will be mitigated during training.

To maintain the high image generation performance of the frozen diffusion model, we will keep the text attention modules but input a zero text feature. The zero text feature has the same feature dimension (by default 768) as the original text feature but with a length of one and always has the weights of zero. We need to keep the distribution of text input the same for these blocks as they will not be further trained. Otherwise distribution drift will affect image generation performance.

We first introduce details about experimental settings and then evaluate the effectiveness of our methods and state-of-the-art mechanisms defending against successful inversion attacks. We conduct all our experiments on Plato[51], an open-source research framework for deploying decentralized training on multiple devices. Plato can support large-scale decentralized training and help deploy the server and the clients on separate devices conveniently. We use the same setting as previous where we have 50 clients in total and each has 1000 training samples. The number of clients will effect efficiency and scalability but will not effect image generation performance or privacy-preserving ability. Since the main focus of this paper is on the latter two aspects, we do not particularly study other settings.

For the pre-trained models, we used stable diffusion V-1.5 and ControlNet V-1.1. The autoencoder is from pre-trained CLIP model with ViT-Large-Patch14 [11]. The resolution of the input and generate images is $512\times 512$. We used the MS-COCO [22] as the training dataset for fine-tuning diffusion models to generate high quality images with given conditions. The MS-COCO dataset contains over 120K wild images with proper prompts. It is a common dataset used in fine-tuning large diffusion models and text-to-image generation tasks. The model is fine-tuned over MS-COCO for 25000 iterations with a batch size of 4. The rest training settings is the same as default implementation of ControlNet [2] where we use AdamW optimizer with learning rate of $1\times 10^{-5}$. The noise coefficient $\lambda_{t}$ is 0.

Evaluation Metrics. For comparing the performance, we need to verify that the privacy-preserving method will not harm image generation performance and that an adversary will not be able to reconstruct private images. For the first objective, we use Fréchet Inception Distance [52] (FID) to evaluate quality of generated images and the CLIP score [11] to evaluate whether the prompts and generated images are matched (in range of $[0,100]$). We use the MS-COCO validation set with over 5000 images to evaluate the quality of generated images. Lower FID indicates better quality of generated images. A higher CLIP score indicates that the text prompts and the generated images match to each other better.

For the second prospective, we use PSNR and SSIM as mentioned in previous sections. The images are generated with the same random seed. The settings for evaluating privacy against reconstructing private data is the same as in Section V. Lower PSNR and SSIM indicate the reconstructed images are less similar to the private data, meaning better privacy-preserving effectiveness.

For inverse-network based attacks, if we use the original split learning structure, the attack is under a black-box setting. If we use our designed split learning structure, as the weights of models on clients only involve weights of pre-trained CLIP models downloaded from [53], the attack is under a white-box setting. Within the realm of conditional image generation, various tasks involve different conditions. We assess three types of conditions: canny lines, scribbles, and segmentation maps. These conditions represent a range from detailed to coarse-grained, with the lines drawn in the condition images varying accordingly.

During the training process of ControlNet, the timestep is sampled among the range of $[t_{s},t_{\max}]$. With the default $k$ and $\beta_{0}$, we say $\epsilon_{s}=\epsilon(t_{s},k,\beta_{0})$. So, according to Eq. 8, during the training process, the privacy budget is equal to or larger than $\epsilon_{s}$. Because we need to ensure that the privacy of every image is preserved, when we evaluate the effectiveness of privacy-preserving methods, in terms of both numerical data and visualization, we consider the worst case of least noise added and sample the timestep as $t_{s}$ and send the intermediate features to the server.

For our and other privacy-preserving methods, we implement them with our designed gradient sending-back free structure. For $(\epsilon,\Delta)$-LDP mechanism, $\Delta=1\times 10^{-4}$ and we calculate that $\alpha\approx 0.16$. The training latency remains the same after adding our privacy-preserving methods.

We present qualitative results in Table IV. The conclusion is that from numerical data and visualization, we can see that our method is the only method that can protect privacy without loss of image generation quality. The methods that can generate images correctly cannot preserve privacy well. The methods can preserve privacy well are not able to generate good images. Though more advanced methods such as Mixup and PS can provide good privacy protection in some cases, they fail to correctly generate images of good quality conforming to the conditions.

One of our new insights is that all previous privacy-preserving methods try to propose a general method for split learning, overlooking the variance between different use cases. We can easily extend methods like DataMix from image classification to different tasks. However, they cannot achieve satisfactory performance when we really verify them on the task of image generation. Our privacy-preserving method is tailored for diffusion models, considering the specialty of the overall model structure of diffusion model to how prompts are processed. Let’s take a look at the detailed analysis.

In Theorem VI.1, we can set different privacy budgets with proper $k$ and $\beta_{0}$. In Fig. 11, we change privacy budgets by setting different scheduling parameters $k$ and $\beta_{0}$ respectively. In the default setting of our method, the privacy budget is $8$. We try the other two cases of setting privacy budgets as $0.3$ and $2$. As shown in Fig. 11, our method can still generate images of good quality. However, for the baseline methods, they fail to generate good images and protect privacy when they use the same privacy budgets of $0.3$ and $2$.