# MemControl: Mitigating Memorization in Medical Diffusion Models via
Automated Parameter Selection

Raman Dutt1, Pedro Sanchez1, Ondrej Bohdal1, Sotirios A. Tsaftaris1, Timothy
Hospedales1,2
1The University of Edinburgh 2Samsung AI Center, Cambridge
<EMAIL_ADDRESS>
###### Abstract

Diffusion models show remarkable ability in generating images that closely
mirror the training distribution. However, these models are prone to training
data memorization, leading to significant privacy, ethical, and legal
concerns, particularly in sensitive fields such as medical imaging. We
hypothesize that memorization is driven by the overparameterization of deep
models, suggesting that regularizing model capacity during fine-tuning could
be an effective mitigation strategy. Parameter-efficient fine-tuning (PEFT)
methods offer a promising approach to capacity control by selectively updating
specific parameters. However, finding the optimal subset of learnable
parameters that balances generation quality and memorization remains elusive.
To address this challenge, we propose a bi-level optimization framework that
guides automated parameter selection by utilizing memorization and generation
quality metrics as rewards. Our framework successfully identifies the optimal
parameter set to be updated to satisfy the generation-memorization tradeoff.
We perform our experiments for the specific task of medical image generation
and outperform existing state-of-the-art training-time mitigation strategies
by fine-tuning as few as 0.019% of model parameters. Furthermore, we show that
the strategies learned through our framework are transferable across different
datasets and domains. Our proposed framework is scalable to large datasets and
agnostic to the choice of reward functions. Finally, we show that our
framework can be combined with existing approaches for further memorization
mitigation.

##  1 Introduction

Diffusion models [1, 2, 3, 4] have demonstrated exceptional proficiency in
generating high-quality data across various modalities, including images [5,
6, 7, 8], audio [9, 10] and graphs [11, 12]. They underpin the commercial
success of popular products such as Midjourney, Stable Diffusion and DALL-E,
which make it possible to generate images based on user queries. The appeal of
generative diffusion models lies in their ability to produce high-quality
novel images that capture the underlying intricacies of real data. These
models have diverse practical applications, for example they can be used for
data augmentation that helps train more accurate recognition systems [13].
They can also be helpful for academic research, for example to help with
understanding disease progression [14]. One particularly important and
potentially transformative use case of synthetic data is for surpassing
restrictions on sharing sensitive data.

Generative diffusion models or the generated synthetic data can be shared
instead of the original data, addressesing privacy concerns while making it
possible to drive progress. However, the utility of this approach would be
nullified if diffusion models were to inadvertently disclose part of the
underlying training data by generating images that are their near-replicas.
This phenomenon is known as memorization and has linked diffusion models with
various legal and ethical issues [15]. It has singled them out as potentially
the least private form of generative models [16].

Privacy risks related to memorization are especially significant in the
context of healthcare applications. The risks include exact replication of
images, potentially disclosing information about the patient [17] and making
it possible to conduct re-identification attacks [18]. In re-identification
attacks, models that can identify whether two images belong to the same
patient can allow an attacker with partial information to link a memorized
synthetic image to the associated patient. The growing concern has inspired
the development of methods to mitigate memorization. These methods include
training-time interventions such as de-duplication [19, 20] and noise
regularizers [21], as well as inference-time mitigations such as token-
rewriting [20, 22]. However, a key limitation of the fast and effective token-
rewriting methods is their vulnerability to white-box attacks. An attacker
with access to the model can simply disable the token-rewriting during
inference to extract private data.

In this paper, we introduce a new training-time intervention aimed at
addressing memorization in diffusion models. We hypothesize that the
overcapacity of large neural networks facilitates memorization, and thus we
explore whether regularizing model capacity can mitigate this issue. In
domain-specific diffusion models, such as those used for medical imaging, the
typical workflow involves starting with a general-purpose diffusion model and
fine-tuning it with task-specific data [23]. We posit that there is a trade-
off between the extent of fine-tuning and the potential for the model to
memorize in-domain medical images. If no fine-tuning is performed, the pre-
trained model may not replicate specific in-domain data but will also struggle
to generate high-quality domain-specific images. Conversely, extensive fine-
tuning will push the model to generate high-quality in-domain images but also
increase memorization risk. We investigate whether it is possible to strike an
effective balance, constraining the model enough to produce high-quality
generations while minimizing memorization.

To this end, we leverage parameter-efficient fine-tuning (PEFT) [24, 25, 26,
27, 28, 29], a set of methods that selectively update specific subsets of
parameters during the fine-tuning process. Each PEFT method’s unique approach
to freezing versus updating parameters results in a different balance between
generation quality and memorization. Our study aims to identify which PEFT
strategies can best manage this trade-off, ensuring effective specialization
of pre-trained models with minimal risk of memorization. We introduce a bi-
level optimization framework designed to search for the optimal set of
parameters to update, striking a balance between high-quality generation and
minimal memorization. Our empirical results on the MIMIC medical imaging
dataset [30] demonstrate a significantly improved quality-memorization
tradeoff compared to recent state-of-the-art memorization mitigation
techniques [20, 22]. While our method is fast and robust during inference, it
does incur additional computational cost during training. We show the
potential to mitigate this cost by transferring high-quality updatable
parameter subsets across different learning problems. This suggests our
discovered fine-tuning strategy could be applied off-the-shelf in new
applications, thus providing a versatile and effective solution for generative
tasks in sensitive domains.

The key contributions of our work are summarised as follows: (1) Empirically
demonstrating that reducing model capacity significantly mitigates
memorization; (2) Highlighting the delicate balance between memorization and
generalization and how it is influenced by model capacity (Sec. 4.2); (3)
Proposing a bi-level optimization framework that automates parameter
(capacity) selection to optimize both generative quality and memorization
mitigation (Sec. 4.3); and (4) Showcasing that the parameter subsets
discovered are transferable across different datasets and domains (Sec. 4.4).

##  2 Related Work

###  2.1 Memorization in Diffusion Models

Memorization in deep generative models has been explored in several contexts,
including training data extraction [31, 16], content replication [21] and data
copying [20]. Memorization occurs also in the medical domain and has been
observed for brain MRI and chest X-ray data [32, 33], posing risks for example
via patient re-identification [18]. Given the significant risks associated
with memorization, several strategies have been proposed to mitigate it. Based
on evidence that duplicate samples increase memorization, the most
straightforward approach involves de-duplicating the training set [16, 19,
34]. However, such approach has been shown to be insufficient [20]. In text-
to-image generative models, text conditioning significantly influences
memorization and several mitigation strategies have been proposed based on
this observation [20, 22, 35].

Memorization mitigation strategies can be categorized into training-time and
inference-time methods. Training-time methods include augmenting text input as
shown in [20], and can be implemented for example via (1) using multiple
captions per image, (2) adding Gaussian noise to text embeddings, (3) adding
random words or numbers to text prompts. Among these, adding random words
(RWA) is shown to be the most effective. Another method, threshold mitigation,
is based on differences in model output magnitudes between memorized and non-
memorized prompts, enabling fast detection and mitigation during training
[22]. [35] have introduced a strategy examining cross-attention score
differences, observing higher attention concentration on specific tokens in
memorized prompts. Other studies have explored injecting a deliberate
memorization signal via a watermark and detecting it in generated images with
a trained classifier [36] to prevent unintended data usage during training
[37].

###  2.2 Parameter-Efficient Fine-Tuning

Models trained on vast online datasets [38] typically lack domain-specific
synthesis capabilities, for example they may not be able to generate medical
chest X-ray data [23, 39]. Consequently, additional fine-tuning following the
pre-train fine-tune paradigm is necessary. While conventional methods update
all parameters via full fine-tuning (FT) [23], recent advancements introduce
parameter-efficient alternatives [25, 26, 28, 29], achieving comparable
performance without full model updates.

PEFT strategies freeze most pre-trained model parameters and fine-tune either
existing parameters (selective PEFT) or introduce and fine-tune new ones
(additive PEFT). Selective PEFT methods include attention tuning [40], bias
tuning [41], and normalization tuning [42], while additive PEFT methods
include LoRA [43], SSF [27], SV-DIFF [28], DiffFit [29] and AdaptFormer [44]
among others, catering to various visual tasks. PEFT matches or exceeds
conventional full fine-tuning [45], including on medical imaging [46]. For
text-to-image generation, strategies such as SV-DIFF [28] and DiffFit [29]
have been proposed, and they excel in both in-domain and out-of-domain
settings [46]. In addition to performance evaluation, bi-level optimization
frameworks based on PEFT have been shown to enhance fairness and reduce bias
in discriminative tasks [47].

##  3 Methodology

###  3.1 Diffusion Model Background

We now summarize the fundamentals of diffusion models. Diffusion models use a
forward and reverse diffusion process. A forward diffusion process involves a
fixed Markov chain spanning $T$ steps, where each step introduces a pre-
determined amount of Gaussian noise. The noise is iteratively added to a data
point $x_{0}$ drawn from the real data distribution $q(x_{0})$ as described
by:

$\displaystyle q\left(x_{t}\mid
x_{t-1}\right)=\mathcal{N}(x_{t};\sqrt{1-\beta_{t}}x_{t-1},\beta_{t}\mathbf{I}),$
(1)

where $\beta_{t}$ represents the scheduled variance at step $t$. Closed-form
expression for this process is:

$x_{t}=\sqrt{\overline{\alpha}_{t}}x_{0}+\sqrt{1-\overline{\alpha_{t}}}\epsilon,$
(2)

where $\overline{\alpha}_{t}=\prod_{i=1}^{t}(1-\beta_{t})$.

In the reverse diffusion process, a Gaussian vector $x_{T}$ sampled from
$\mathcal{N}(0,1)$ undergoes denoising so that it can be mapped onto an image
$x_{0}$ that belongs to the distribution $q(x)$. At each denoising step, a
trained noise predictor $\epsilon_{\theta}$ removes the noise
$\epsilon_{\theta}(x_{t})$ added to $x_{0}$. The formulation for estimating
$x_{t-1}$ is:

$x_{t-1}=\sqrt{\overline{\alpha}_{t-1}}\hat{x}_{0}^{t}+\sqrt{1-\overline{\alpha}_{t-1}}\epsilon_{\theta}(x_{t}),$
(3)

where

$\hat{x}_{0}^{t}=\frac{x_{t}-\sqrt{1-\overline{\alpha}_{t-1}}\epsilon_{\theta}(x_{t})}{\sqrt{\overline{\alpha}_{t}}}.$
(4)

###  3.2 Memorization and Image-Quality Metrics

Our goal is to learn the optimal parameter subset for fine-tuning, given a
PEFT method to limit the amount of memorization while maintaining generated
image quality. There are various metrics used in the community to measure both
memorization [16, 21] and generation quality [48]. In this section, we
introduce the metrics that we use in this paper.

Memorization: Nearest Neighbour Distance The simplest definition of
memorization considers an example $\hat{x}$ drawn from the model to be
memorized if there is a training example $x$ such that $l(x,\hat{x})<\delta$
for some distance metric $l$ and threshold $\delta$. This means a training
example has been memorized if the model produces a near-replica as a sample. A
natural metric for the overall degree of memorization of a model is the
average nearest neighbour distance between each synthetic sample and its
closest example in the training set. For samples $\\{\hat{x}_{i}\\}$ drawn
from the model and a training set $\\{x_{j}\\}$, we follow [49] and quantify
memorization as the average minimum cosine distance (AMD):

$d^{amd}=\frac{1}{N}\sum_{i=1}^{N}\min_{j}l_{\cos}(x_{i},x_{j}).$ (5)

Memorization: Extraction Attack We also evaluate memorization using an image
extraction attack described in [16]. The attack involves two main steps: (1)
Generating numerous image samples for a given prompt, and (2) Conducting
membership inference to identify memorized images. In [16], membership
inference is performed using a heuristic that constructs a graph of similar
samples and identifies sufficiently large cliques—sets of samples that are all
similar to each other. Samples from these cliques are likely to be memorized
examples. Similarity is measured using a modified $\ell_{2}$ distance, defined
as the maximum $\ell_{2}$ distance across corresponding tiles of the two
compared images. A smaller number of extracted images translates to lesser
memorization.

Memorization: Denoising Strength Heuristic A limitation of the two measures
above is that they are comparatively inefficient to compute. An efficient
proxy for memorization based on the magnitude of the diffusion noise
prediction was introduced in [22]. It represents the discrepancy between model
predictions conditioned on a prompt embedding $e_{p}$ and an empty string
embedding $e_{\phi}$ averaged across multiple time steps. The detection metric
can be formulated as follows:

$d^{mem}=\frac{1}{T}\sum_{t=1}^{T}||\epsilon_{\theta}(x_{t},e_{p})-\epsilon_{\theta}(x_{t},e_{\phi})||_{2},$
(6)

for a noise-predictor model $\epsilon_{\theta}$ trained on input samples with
Gaussian noise $x_{t}$ at time-step $t$, incorporating prompt embedding
$e_{p}$, empty string embedding $e_{\phi}$, and $T$ time steps. A higher value
of $d^{mem}$ indicates more substantial memorization [22].

Quality: Fréchet Inception Distance (FID) FID score assesses the image
generation quality by computing a multivariate normal distribution between
real and synthetic images estimated from the features of a pre-trained image
encoder. We use a DenseNet-121 [50] model that is pre-trained on chest X-rays
[51] to compute the features.

Quality: BioViL-T Score: BioViL-T [48] is a state-of-the-art vision-language
model specifically trained for analyzing chest X-Rays and radiology reports.
The model returns a joint score denoting the clinical correlation between an
image and a text caption. A greater similarity score indicates a higher
quality of generated images with better textual guidance.

###  3.3 Parameter-Efficient Fine-Tuning

The key idea of PEFT is to fine-tune a very small subset of model parameters
$\phi\subset\theta$, i.e. $|\phi|\ll|\theta|$. In our framework, the selection
of which parameters are fine-tuned is represented using a binary mask $\omega$
that is applied on the model parameters $\theta$. Each element of $\omega$
determines if a specific parameter at a specific location in the model should
be frozen or fine-tuned. The parameters of the pre-trained model are initially
$\theta_{0}$, and after fine-tuning on data $\mathcal{D}^{train}$ they change
by $\Delta\phi$ compared to their initial values. The fine-tuning process can
then be defined via

$\Delta\phi^{*}=\operatorname*{arg\,min}_{\Delta\phi}\mathcal{L}^{base}\left(\mathcal{D}^{train};\theta_{0}+\omega\odot\Delta\phi\right),$

where $\mathcal{L}^{base}$ is the mean squared error in the context of image
generation problems.

We aim to determine the optimal mask $\omega$ that results in fine-tuning with
minimal memorization and the highest image generation quality.

###  3.4 Our Method: Optimizing PEFT for Mitigating Memorization

Our method employs a multi-objective bi-level optimization formulation. It
trains the model in the inner loop assuming a particular PEFT mask that
specifies which diffusion model parameters to freeze or learn. The PEFT mask
is optimized in the outer loop so as to obtain a diffusion model with low
memorisation and high generation quality. This process is formalized using Eq.
7 and Alg. 1.

$\begin{split}\omega^{*}&=\underset{\omega}{\operatorname{argmin}}\text{
}\mathcal{L}^{outer}\left(d^{mem}\left(\mathcal{D}^{mem};\Delta\phi^{*}\right),d^{fid}\left(\mathcal{D}^{val};\Delta\phi^{*}\right)\right)\\\
\text{such
that}\quad\Delta\phi^{*}&=\operatorname*{arg\,min}_{\Delta\phi}\mathcal{L}^{base}\left(\mathcal{D}^{train};\theta_{0}+\omega\odot\Delta\phi\right),\end{split}$
(7)

where $\mathcal{L}^{base}$ is the usual diffusion model loss, and
$\mathcal{L}^{outer}=({d}^{mem},{d}^{fid})$ is the two-element vector of
memorization and quality objectives respectively (Sec. 3.2). Note that because
the outer loop is a multi-objective optimization, $\omega^{*}$ is not a single
model, but the set of mask configurations that constitute the pareto front
dominating others in terms of memorization and quality. Where necessary, we
pick a single model from the pareto front for evaluation by scalarizing the
score vector (for example by arithmetic mean) and picking the optimum mask.

The overall workflow first optimizes the mask using a small proxy dataset (see
Sec. 4.1) for efficient search, and then fixes a specific optimum mask and re-
trains on the full training dataset prior to evaluation (Figure 1).

Figure 1: Overall schematic of our framework. Stage 1: PEFT Mask Search (top):
We use a subset of the training set to search for the mask that decides which
PEFT components of a pre-trained model $\theta$ should be fine-tuned to
optimise both generation quality ($d^{fid}$) and memorization ($d^{mem}$).
Stage 2: Fine-tune with mask (bottom):  The optimal mask from the HPO search
is used for fine-tuning on full dataset and final results are reported on the
test set ($\mathcal{D}^{test}$).  Algorithm 1 Optimizing PEFT for Mitigating
Memorization

1: Input: pre-trained model $\theta_{0}$, $\alpha$: fine-tuning learning rate,
number of trials $T_{N}$

2: Output: fine-tuned model $\theta_{0}+\omega\odot\Delta\phi$ and mask
$\omega$

3: while number of completed trials $<T_{N}$ do

4: Initialize $\Delta\phi=0$ for $\phi\leftarrow\phi_{0}\subset\theta_{0}$

5: Propose mask $\omega\leftarrow HPO$

6: while not converged do

7:
$\Delta\phi\leftarrow\Delta\phi-\alpha\nabla_{\phi}\mathcal{L}^{base}\left(\mathcal{D}^{train};\theta_{0}+\omega\odot\Delta\phi\right)$
// PEFT

8: end while

9: Evaluate $d^{mem}$ on $\mathcal{D}^{mem}$ and $d^{fid}$ on
$\mathcal{D}^{val}$ with model $\theta_{0}+\omega\odot\Delta\phi$ and report
to HPO

10: end while

11: Return the best masks $\omega$ and fine-tuned models
$\theta_{0}+\omega\odot\Delta\phi$

##  4 Experiments

###  4.1 Experimental Setup

Architecture: Our experiments adopted the pre-trained Stable Diffusion (v1.5)
[52] model implemented in the Diffusers package [53]. The stable diffusion
pipeline consists of three components: (1) a modified U-Net [54] with self-
attention and cross-attention layers for textual guidance, (2) a text-encoder
derived from [55], and a variational auto-encoder [56]. For fine-tuning on
medical data we follow previous practice [46, 23], and fine-tune only the
U-Net while keeping the other components frozen during our experiments.

Baselines:  We compare our approach with different fine-tuning strategies: (1)
an off-the-shelf pre-trained Stable Diffusion model that has all of the
components frozen, (2) conventional full fine-tuning where every parameter of
the U-net is updated, (3) the SV-DIFF PEFT method [28], (4) the DiffFit PEFT
method [29], and (4) an Attention tuning PEFT method [40]. Furthermore, we
implement existing state-of-the-art training-time memorization-mitigation
methods: (1) Random Word Addition (RWA) [20] and threshold mitigation [22].
Overall our main baselines are a combination of four fine-tuning methods and
two mitigation strategies.

Search Space and Search Strategy: Our method is based on optimizing the choice
of fine-tuning parameters, which requires defining a search space for
$\omega$. We conduct our experiments on three different parameter search
spaces associated with distinct PEFT methods: (1) SV-DIFF: Identifying which
attention blocks within the U-Net should include SV-DIFF parameters, (2)
DiffFit: Identifying which attention blocks within the U-Net should include
DiffFit parameters, (3) Attention Tuning: Determining which attention layers
within the U-Net should be fine-tuned. The search space for both SV-DIFF and
DiffFit is $\Omega\in\\{0,1\\}^{13}$, while the search space for attention
tuning is $\Omega\in\\{0,1\\}^{16}$. We also consider an approxmiate global
search over all three search spaces by combining the three pareto fronts from
each individual search space. We use the evolutionary search-based NSGA-II
Sampler [57] as implemented in Optuna [58] to search the space of masks (Eq.
7).

Datasets:  We use the MIMIC dataset [30] that consists of chest X-rays and
associated radiology text reports. The dataset was split into train,
validation and test splits ensuring no leakage of subjects, resulting in 110K
training samples, 14K validation and 14K test samples. Note that that several
images in this dataset can be associated with the same text prompt due to
similarity in clinical findings.

Efficient HPO with a Memorization Subset: Conducting mask search on the whole
unaltered training set is infeasible because only a minority of images are
memorised, and calculating the memorization metric (Eq. 6, Sec. 3.2) is
therefore noisy and slow. To alleviate this, we construct a specific training
set $\mathcal{D}^{train}$ for HPO mask search consisting of (1) 1% of the
original training set, and (2) a special memorization set $\mathcal{D}^{mem}$
to drive mask search. Specifically, we select 100 text-image pairs from the
training set and duplicate them 50 times. With this setup the HPO training set
is small, and we know which images are expected to be memorised. HPO mask
search is then driven by the two validation objectives $d^{mem}$ and $d^{fid}$
(Eq. 7, Sec. 3.2) which are evaluated on the memorization set
$\mathcal{D}^{mem}$ and the validation set $\mathcal{D}^{val}$ respectively. A
detailed explanation on memorization subset creation is provided in Appendix
C.1.

Experimental Settings:  All experiments fine-tune a pre-trained Stable
Diffusion model (v1.5) for 20,000 steps with a batch size of 128, using the
AdamW optimizer [59]. We employ no warmup steps and utilize a constant
learning rate scheduler. Hyperparameter optimization (HPO) search was
implemented using Optuna package [58]. The final generation quality, indicated
by the FID score ($d^{fid}$) is reported on a held-out test set, and
memorization, indicated by the AMD score ($d^{amd})$, is measured on the
training set. The BioViL-T score is measured between the original text
captions and generated images from the test set. The extraction attack is
conducted between synthetic images and the training set images.

###  4.2 PEFT Mask Choices Explore the Memorization-Generation Tradeoff

We first explore the influence of model capacity on memorization vs.
generalization quality tradeoff by recording the FID score ($d^{fid}$) and the
memorization metric ($d^{mem}$) for different parameter subsets explored
during the HPO search. Each HPO iteration considers a different subset of
parameters to freeze or update. The results in Figure 2 show the combined
pareto front from searching within the space of SV-Diff [28], DiffFit [29] and
Attention Tuning [40] PEFT options. We also report the performance of the
defacto full fine-tuning [46, 14], as well as the vanilla configurations of
the PEFT methods. From the performance scatter plot, we see that certain
configurations favour generation quality ($d^{fid}\downarrow$), others might
favour reducing memorization ($d^{mem}\downarrow$) while some can perform
poorly on both fronts ($d^{fid}$ $\uparrow$, $d^{mem}\uparrow$). However, our
framework also discovers the Pareto front of configurations that dominate
others ($d^{fid}$ $\downarrow$, $d^{mem}\downarrow$), including configurations
that are near optimal for both (bottom left). Importantly, the pareto-front
excludes configurations that include the defacto full fine-tuning and PEFT
variants.

Figure 2: Plot illustrating how model capacity affects the memorization vs.
generation quality tradeoff. The Hyperparameter Optimization (HPO) search
explored various combinations of parameter subsets during fine-tuning. Each
combination results in a different model capacity, generation quality
($d^{fid}\downarrow$) and memorization ($d^{mem},\downarrow$). The performance
of these subsets is compared to two default PEFT configurations (orange
squares) and full fine-tuning (green square). The pareto front of optimal
parameter subsets, indicated by the lowest $d^{fid}$ and $d^{mem}$, are marked
in red.

We select a specific model from the Pareto front for further evaluation by
taking the mask configuration with the lowest arithmetic mean of FID score
($d^{fid}\downarrow$) and memorisation score ($d^{mem}\downarrow$). Their
scale is relatively comparable and both are minimized, so taking their
arithmetic mean is reasonable.

###  4.3 Quantitative Comparison vs. State of the Art

Given the optimal PEFT mask discovered by our framework in Sec. 4.2, we fine-
tune using this mask on the full training set, and compare this against
competitors. Recall that we consider two types of competitors for memorization
mitigation: purpose designed regularisers [20, 22], and alternative PEFT
methods [28, 29]. These mitigation methods are independent and potentially
complementary, so we also go beyond prior work in reporting all their
combinations as well. From the results in Table 1, we observe that: (1)
Surprisingly, and despite not being designed to tackle memorization, PEFT
methods SV-DIFF [28] and DiffFit [29] outperform state-of-the-art purpose
designed mitigation methods such as [22] on several metrics. (2) Our
MemControl framework achieves the best performance across all metrics. It
outperforms both purpose designed mitigation methods, PEFT methods, and their
combination across all generation quality and memorization metrics.

Table 1: Comparitive evaluation of generation quality (FID $\downarrow$,
BioViL-T Score $\uparrow$) and memorization (AMD $\uparrow$, Extracted Images
$\downarrow$) on the MIMIC dataset. Our MemControl outperforms purpose
designed memorization mitigations, standard PEFT methods, and their
combinations.

FT Method + Mitigation Strategy | FID Score ($\downarrow$) | AMD ($\uparrow$) | BioViL-T Score ($\uparrow$) | Num. Extracted Images ($\downarrow$)
---|---|---|---|---
Freeze | 323.78 | 0.801 | 0.23 |  000
Full FT [23] |  035.79 | 0.029 | 0.60 | 356
Full FT [23] \+ RWA [20] |  067.48 | 0.041 | 0.58 | 320
Full FT [23] \+ Threshold Mitigation [22] |  025.43 | 0.048 | 0.62 | 298
SV-DIFF [28] |  012.77 | 0.062 | 0.72 |  063
SV-DIFF [28] \+ RWA [20] |  041.25 | 0.059 | 0.52 |  058
SV-DIFF [28] \+ Threshold Mitigation [22] |  023.64 | 0.068 | 0.64 |  045
DiffFit [29] |  015.19 | 0.061 | 0.69 |  069
DiffFit [29] \+ RWA [20] |  017.18 | 0.060 | 0.67 |  064
DiffFit [29] \+ Threshold Mitigation [22] |  015.31 | 0.067 | 0.69 |  051
Attn Tuning [40] |  040.28 | 0.031 | 0.59 | 358
Attn Tuning [40] \+ RWA [20] |  049.71 | 0.035 | 0.52 | 324
Attn Tuning [40] \+ Threshold Mitigation [22] |  025.43 | 0.041 | 0.62 | 299
MemControl (ours) |  011.93 | 0.114 | 0.75 |  028

###  4.4 PEFT Strategy is Transferable

One might question if a universal approach to balance quality and memorization
across multiple datasets exists. To investigate this, we assess the
transferability of our PEFT strategy, initially discovered on the MIMIC
medical dataset in Sec. 4.2, to a different dataset, namely Imagenette [60],
which serves as a smaller and more manageable subset of ImageNet.
Specifically, we fine-tune the Stable Diffusion model on Imagenette using the
mask derived from MIMIC.

Table 2: Evaluation of generative performance and memorization on Imagenette dataset using parameter subsets learnt from the MIMIC dataset. Our framework leads to the best overall performance in terms of both FID and AMD, showcasing the transferability of strategies from one domain to another. FT Method + Mitigation Strategy | FID Score ($\downarrow$) | AMD ($\uparrow$)
---|---|---
Freeze | 62.16 | 0.101
Full FT | 11.01 | 0.044
Full FT + RWA |  08.72 | 0.048
Full FT + Threshold Mitigation |  09.11 | 0.052
SV-DIFF |  08.10 | 0.058
SV-DIFF + RWA | 13.53 | 0.061
SV-DIFF + Threshold Mitigation |  09.38 | 0.065
DiffFit |  07.40 | 0.054
DiffFit + RWA | 11.94 | 0.058
DiffFit + Threshold Mitigation | 12.03 | 0.062
Attention Tuning |  08.04 | 0.047
Attention Tuning + RWA |  09.12 | 0.049
Attention Tuning + Threshold Mitigation |  09.11 | 0.051
MemControl (ours) |  06.34 | 0.093

The results presented in Table 2 indicate that our MemControl method continues
to outperform competitors in both generation quality and memorization. This
demonstrates the potential for the transferability of fine-tuning masks across
datasets, implying that these masks could be reused for new datasets in the
future, eliminating the need for repeated mask searches.

###  4.5 Additional Ablations

MemControl with existing mitigation mechanisms: Our results in Table 3 show
that combining MemControl with existing mitigation mechanisms leads to
improved performance. We do observe a decrease in generation performance in
the case of RWA, but its combination with MemControl leads to reduced
memorization (AMD $\uparrow$). A detailed analysis of the performance decrease
observed with RWA is provided in Appendix C.2. Combining MemControl with
Threshold Mitigation leads to the best generation and memorization performance
overall.

Table 3: Combining our MemControl with existing mitigation strategies (RWA
[20] and Threshold Mitigation [22]) on the MIMIC dataset. Our framework can be
seamlessly integrated with these approaches and effectively complement
existing methods to enhance memorization mitigation.

Method |  FID Score ($\downarrow$) |  AMD ($\uparrow$) |  BioViL-T Score ($\uparrow$) |  Num. Extracted Images ($\downarrow$)
---|---|---|---|---
MemControl | 11.93 | 0.114 | 0.75 | 28
MemControl \+ RWA | 13.10 | 0.121 | 0.69 | 23
MemControl \+ Threshold Mitigation | 11.31 | 0.142 | 0.75 | 16

MemControl can optimize individual PEFT methods: Our main result in Table 1
applied MemControl across all three PEFT search spaces considered. In Table 4
we demonstrate that when applied to the search space associated with a single
PEFT method, MemControl can substantially improve its performance.

Table 4: Evaluating MemControl in individual parameter search spaces. Our
method leads to the best performance in the respective search space defined by
distinct PEFT methods (SV-DIFF, DiffFit, Attention Tuning) across all
evaluation metrics.

FT Method + Mitigation Strategy |  FID Score ($\downarrow$) |  AMD ($\uparrow$) |  BioViL-T ($\uparrow$) |  Num. Extracted Images ($\downarrow$)
---|---|---|---|---
SV-DIFF | 12.77 | 0.062 | 0.72 |  063
SV-DIFF + RWA | 41.25 | 0.059 | 0.52 |  058
SV-DIFF + Threshold Mitigation | 23.64 | 0.068 | 0.64 |  045
MemControl (SV-DIFF Search Space) | 11.93 | 0.114 | 0.75 |  028
DiffFit | 15.19 | 0.061 | 0.69 |  069
DiffFit + RWA | 17.18 | 0.060 | 0.67 |  064
DiffFit + Threshold Mitigation | 15.31 | 0.067 | 0.69 |  051
MemControl (DiffFit Search Space) | 12.21 | 0.086 | 0.72 |  042
Attn Tuning | 40.28 | 0.031 | 0.59 | 358
Attn Tuning + RWA | 49.71 | 0.035 | 0.52 | 324
Attn Tuning + Threshold Mitigation | 25.43 | 0.041 | 0.62 | 299
MemControl (Attn Tuning Search Space) | 22.36 | 0.044 | 0.063 | 281

###  4.6 Analyzing the Best Solution

Our best solution relied on a subset of SV-DIFF learnable parameters. We
analyze the mask indicating which subset of parameters should be fine-tuned.
It indicates necessity of fine-tuning all of the cross-attention layers in the
Stable Diffusion U-Net, which strongly aligns with the findings in [28].
Additionally, the mask suggests fine-tuning all of the self-attention layers
except for those in down-blocks 2 and 3 and up-blocks 3 and 4. Adopting this
strategy corresponds to fine-tuning 0.019% of the parameters in Stable
Diffusion U-Net.

###  4.7 Discussion and Limitations

A limitation of our method is its higher computational demand compared to
baselines (see Appendix B for details). However, our results indicate that the
discovered mask exhibits transferability, implying that repeated mask learning
for each new dataset may not be necessary. This means the computational cost
of mask search could be incurred once and then amortized over multiple
deployments across different datasets. While our approach significantly
reduces memorization, it does not completely eliminate it, especially in cases
where the PEFT masks are reused across different downstream scenarios.
Therefore more work is necessary to completely solve this problem.

##  5 Conclusion

In this paper, we approached the challenge of mitigating memorization in
diffusion models from a parameter-efficient fine-tuning capacity control
perspective. We illustrated how the selection of fine-tuning parameter space
influences the balance between memorization and generation quality. By
searching for PEFT masks, we identified configurations that offer a promising
trade-off between generation fidelity and memorization. This introduces a
novel and distinctive approach to address the significant memorization issue
in diffusion models, contrasting with current mainstream regularization- and
token rewriting-based methods.

## References

  * [1] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256–2265, Lille, France, 07–09 Jul 2015. PMLR.
  * [2] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  In International Conference on Learning Representations, 2021.
  * [3] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 6840–6851. Curran Associates, Inc., 2020.
  * [4] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  In International Conference on Learning Representations, 2021.
  * [5] Prafulla Dhariwal and Alexander Quinn Nichol.  Diffusion models beat GANs on image synthesis.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
  * [6] Tero Karras, Samuli Laine, and Timo Aila.  A style-based generator architecture for generative adversarial networks.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2019.
  * [7] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen.  Hierarchical text-conditional image generation with clip latents.  arXiv:2204.06125, 2022.
  * [8] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684–10695, 2022.
  * [9] Heeseung Kim, Sungwon Kim, and Sungroh Yoon.  Guided-tts: A diffusion model for text-to-speech via classifier guidance.  In International Conference on Machine Learning, pages 11119–11133. PMLR, 2022.
  * [10] Rongjie Huang, Jiawei Huang, Dongchao Yang, Yi Ren, Luping Liu, Mingze Li, Zhenhui Ye, Jinglin Liu, Xiang Yin, and Zhou Zhao.  Make-an-audio: Text-to-audio generation with prompt-enhanced diffusion models.  In International Conference on Machine Learning, pages 13916–13932. PMLR, 2023.
  * [11] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang.  Geodiff: A geometric diffusion model for molecular conformation generation.  In International Conference on Learning Representations, 2022.
  * [12] Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard.  Digress: Discrete denoising diffusion for graph generation.  In The Eleventh International Conference on Learning Representations, 2023.
  * [13] Brandon Trabucco, Kyle Doherty, Max A Gurinas, and Ruslan Salakhutdinov.  Effective data augmentation with diffusion models.  In The Twelfth International Conference on Learning Representations, 2024.
  * [14] Sourav Kumar, Shell Xu Hu, Timothy Hospedales, Praveer Singh, and Jayashree Kalpathy-cramer.  Monitoring disease progression with stable diffusion- a visual approach.  In Medical Imaging with Deep Learning, 2024.
  * [15] Matthew Butterick.  Stable diffusion litigation · joseph saveri law firm & matthew butterick, Jan 2023.
  * [16] Nicolas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito, and Eric Wallace.  Extracting training data from diffusion models.  In 32nd USENIX Security Symposium (USENIX Security 23), 2023.
  * [17] Jinsung Yoon, Lydia N Drumright, and Mihaela Van Der Schaar.  Anonymization through data synthesis using generative adversarial networks (ads-gan).  IEEE journal of biomedical and health informatics, 24(8):2378–2388, 2020.
  * [18] Virginia Fernandez, Pedro Sanchez, Walter Hugo Lopez Pinaya, Grzegorz Jacenków, Sotirios A Tsaftaris, and Jorge Cardoso.  Privacy distillation: reducing re-identification risk of multimodal diffusion models.  arXiv:2306.01322, 2023.
  * [19] Ryan Webster, Julien Rabin, Loic Simon, and Frederic Jurie.  On the de-duplication of laion-2b.  arXiv:2303.12733, 2023.
  * [20] Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein.  Understanding and mitigating copying in diffusion models.  Advances in Neural Information Processing Systems, 36, 2023.
  * [21] Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein.  Diffusion art or digital forgery? investigating data replication in diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
  * [22] Yuxin Wen, Yuchen Liu, Chen Chen, and Lingjuan Lyu.  Detecting, explaining, and mitigating memorization in diffusion models.  In The Twelfth International Conference on Learning Representations, 2024.
  * [23] Pierre Chambon, Christian Bluethgen, Jean-Benoit Delbrouck, Rogier Van der Sluijs, Małgorzata Połacin, Juan Manuel Zambrano Chaves, Tanishq Mathew Abraham, Shivanshu Purohit, Curtis P Langlotz, and Akshay Chaudhari.  Roentgen: Vision-language foundation model for chest x-ray generation.  arXiv preprint arXiv:2211.12737, 2022.
  * [24] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  LoRA: Low-rank adaptation of large language models.  In ICLR, 2022.
  * [25] Dawid Jan Kopiczko, Tijmen Blankevoort, and Yuki M Asano.  VeRA: Vector-based random matrix adaptation.  In The Twelfth International Conference on Learning Representations, 2024.
  * [26] Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang-Ting Cheng, and Min-Hung Chen.  Dora: Weight-decomposed low-rank adaptation.  arXiv:2402.09353, 2024.
  * [27] Dongze Lian, Daquan Zhou, Jiashi Feng, and Xinchao Wang.  Scaling & shifting your features: A new baseline for efficient model tuning.  Advances in Neural Information Processing Systems, 35:109–123, 2022.
  * [28] Ligong Han, Yinxiao Li, Han Zhang, Peyman Milanfar, Dimitris Metaxas, and Feng Yang.  Svdiff: Compact parameter space for diffusion fine-tuning.  arXiv:2303.11305, 2023.
  * [29] Enze Xie, Lewei Yao, Han Shi, Zhili Liu, Daquan Zhou, Zhaoqiang Liu, Jiawei Li, and Zhenguo Li.  Difffit: Unlocking transferability of large diffusion models via simple parameter-efficient fine-tuning.  arXiv:2304.06648, 2023.
  * [30] Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark.  Mimic-iii, a freely accessible critical care database.  Scientific data, 2016.
  * [31] Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ulfar Erlingsson, et al.  Extracting training data from large language models.  In 30th USENIX Security Symposium (USENIX Security 21), 2021.
  * [32] Muhammad Usman Akbar, Wuhao Wang, and Anders Eklund.  Beware of diffusion models for synthesizing medical images-a comparison with gans in terms of memorizing brain mri and chest x-ray images.  Available at SSRN 4611613, 2023.
  * [33] Salman Ul Hassan Dar, Arman Ghanaat, Jannik Kahmann, Isabelle Ayx, Theano Papavassiliu, Stefan O Schoenberg, and Sandy Engelhardt.  Investigating data memorization in 3d latent diffusion models for medical image synthesis.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, 2023.
  * [34] Ali Naseh, Jaechul Roh, and Amir Houmansadr.  Memory triggers: Unveiling memorization in text-to-image generative models through word-level duplication.  arXiv:2312.03692, 2023.
  * [35] Jie Ren, Yaxin Li, Shenglai Zen, Han Xu, Lingjuan Lyu, Yue Xing, and Jiliang Tang.  Unveiling and mitigating memorization in text-to-image diffusion models through cross attention.  arXiv:2403.11052, 2024.
  * [36] Zhenting Wang, Chen Chen, Lingjuan Lyu, Dimitris N Metaxas, and Shiqing Ma.  Diagnosis: Detecting unauthorized data usages in text-to-image diffusion models.  In The Twelfth International Conference on Learning Representations, 2023.
  * [37] Alexandre Sablayrolles, Matthijs Douze, Cordelia Schmid, and Hervé Jégou.  Radioactive data: tracing through training.  In International Conference on Machine Learning, pages 8326–8335. PMLR, 2020.
  * [38] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al.  Laion-5b: An open large-scale dataset for training next generation image-text models.  Advances in Neural Information Processing Systems, 2022.
  * [39] Raman Dutt, Dylan Mendonca, Huai Ming Phen, Samuel Broida, Marzyeh Ghassemi, Judy Gichoya, Imon Banerjee, Tim Yoon, and Hari Trivedi.  Automatic localization and brand detection of cervical spine hardware on radiographs using weakly supervised machine learning.  Radiology: Artificial Intelligence, 2022.
  * [40] Hugo Touvron, Matthieu Cord, Alaaeldin El-Nouby, Jakob Verbeek, and Hervé Jégou.  Three things everyone should know about vision transformers.  In ECCV, 2022.
  * [41] Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel.  BitFit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  In ACL, May 2022.
  * [42] Samyadeep Basu, Daniela Massiceti, Shell Xu Hu, and Soheil Feizi.  Strong baselines for parameter efficient few-shot fine-tuning.  arXiv, 2023.
  * [43] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  Lora: Low-rank adaptation of large language models.  arXiv:2106.09685, 2021.
  * [44] Shoufa Chen, Chongjian Ge, Zhan Tong, Jiangliu Wang, Yibing Song, Jue Wang, and Ping Luo.  Adaptformer: Adapting vision transformers for scalable visual recognition.  arXiv, 2022.
  * [45] Xiaohua Zhai, Joan Puigcerver, Alexander Kolesnikov, Pierre Ruyssen, Carlos Riquelme, Mario Lucic, Josip Djolonga, Andre Susano Pinto, Maxim Neumann, Alexey Dosovitskiy, et al.  A large-scale study of representation learning with the visual task adaptation benchmark.  arXiv:1910.04867, 2019.
  * [46] Raman Dutt, Linus Ericsson, Pedro Sanchez, Sotirios A Tsaftaris, and Timothy Hospedales.  Parameter-efficient fine-tuning for medical image analysis: The missed opportunity.  arXiv, 2023.
  * [47] Raman Dutt, Ondrej Bohdal, Sotirios A. Tsaftaris, and Timothy Hospedales.  Fairtune: Optimizing parameter efficient fine tuning for fairness in medical image analysis.  In The Twelfth International Conference on Learning Representations, 2024.
  * [48] Shruthi Bannur, Stephanie Hyland, Qianchu Liu, Fernando Perez-Garcia, Maximilian Ilse, Daniel C Castro, Benedikt Boecking, Harshita Sharma, Kenza Bouzid, Anja Thieme, et al.  Learning to exploit temporal structure for biomedical vision-language processing.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
  * [49] Ching-Yuan Bai, Hsuan-Tien Lin, Colin Raffel, and Wendy Chi-wen Kan.  On training sample memorization: Lessons from benchmarking generative modeling with a large-scale competition.  In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, 2021.
  * [50] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, 2017.
  * [51] Joseph Paul Cohen, Mohammad Hashir, Rupert Brooks, and Hadrien Bertrand.  On the limits of cross-domain generalization in automated x-ray prediction.  In Medical Imaging with Deep Learning, 2020.
  * [52] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10684–10695, June 2022.
  * [53] Patrick von Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj, Dhruv Nair, Sayak Paul, William Berman, Yiyi Xu, Steven Liu, and Thomas Wolf.  Diffusers: State-of-the-art diffusion models.  https://github.com/huggingface/diffusers, 2022.
  * [54] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18. Springer, 2015.
  * [55] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In International conference on machine learning. PMLR, 2021.
  * [56] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv:1312.6114, 2013.
  * [57] Kalyanmoy Deb, Amrit Pratap, Sameer Agarwal, and TAMT Meyarivan.  A fast and elitist multiobjective genetic algorithm: Nsga-ii.  IEEE transactions on evolutionary computation, 6(2):182–197, 2002.
  * [58] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama.  Optuna: A next-generation hyperparameter optimization framework.  In KDD, 2019.
  * [59] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In ICLR, 2018.
  * [60] Jeremy Howard.  Imagenette: A smaller subset of 10 easily classified classes from imagenet, March 2019.
  * [61] Chendong Xiang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu.  A closer look at parameter-efficient tuning in diffusion models.  arXiv:2303.18181, 2023.
  * [62] Jae-hun Shim, Kyeongbo Kong, and Suk-Ju Kang.  Core-set sampling for efficient neural architecture search.  arXiv, 2021.
  * [63] Savan Visalpara, Krishnateja Killamsetty, and Rishabh Iyer.  A data subset selection framework for efficient hyper-parameter tuning and automatic machine learning.  In ICML Workshops, 2021.

##  Appendix A Appendix

###  A.1 Qualitative Examples

We present several qualitative examples in this section. In Figure 3, we
illustrate the comparison in generation quality and memorization between a
conventional fine-tuning method (Full FT), adopting a mitigation strategy with
conventional fine-tuning (Full FT + RWA) and our method. Upon a close
inspection, it can be observed that full fine-tuning leads to near-duplicates
of the original image (col. 1 and 2). Using a mitigation mechanism such as RWA
reduces the number of duplicates but decreases the image generation quality
(col 3 and 4). Our solution, MemControl, preserves both image quality and
prevents the generation of duplicate samples.

Figure 3: Plot illustrating a qualitative comparison between full fine-tuning
(Full FT), Full FT + RWA, and MemControl. Full FT (col. 1 & 2) generates
images that are near-replicas of the original training image. Combining Full
FT with RWA mitigation strategy (col. 3 & 4) diversifies the generated images
but deteriorates the quality. MemControl (col. 5 & 6) preserves image quality
and prevents generating identical images.

Figure 4 illustrates some of the memorized samples obtained through the image
extraction attack [16]. For each original image, multiple generations across
different seeds have been shown.

Figure 4: Plot illustrating that conventional full fine-tuning leads to the
generation of memorized samples for the same prompt across multiple different
seeds.

##  Appendix B Analysis of the Compute Cost

In this section, we analyze the compute cost of our approach and the different
baselines. We use hyperparameter optimization (HPO) for each of them. Each HPO
trial was conducted on 1% of the original MIMIC dataset (1100 samples)
combined with 100 image-text pairs from the memorization subset, repeated 50
times (5000 samples), resulting in a total of 6100 samples. Each trial
consisted of 3000 optimization steps with a batch size of 128, requiring 2 GPU
hours (GPUh) per iteration on an NVIDIA V100 GPU. In total, we ran 30 trials
for MemControl, amounting to 60 GPU hours for the entire process. We compare
our approach with tuned baselines. For example, full FT required 10 similar
HPO trials to determine the optimal learning rate, totalling 20 GPU hours.
Table 5 summarizes the compute requirement for each experiment and baseline.

Table 5: An analysis of the compute cost in GPU hours for the different fine-
tuning and memorization mitigation strategies. Tuned Baseline refers to the
baseline obtained after searching for the optimal learning rate and other
hyperparameters.

Method |  | Compute Cost
---
| Full FT/ SV-DIFF/ DiffFit/ Attention (Tuned Baseline)
---
20 hours
| Full FT/ SV-DIFF/ DiffFit/ Attention + RWA(Tuned Baseline)
---
27 hours
| Full FT/ SV-DIFF/ DiffFit/ Attention + Threshold Mitigation (Tuned Baseline)
---
35 hours
MemControl (ours) | 60 hours

##  Appendix C Further Discussion on Experimental Details and Results

###  C.1 Creating the Memorization Subset

Figure 5: Count for the most frequently occurring tokens in the MIMIC dataset.

The memorization subset was created with the objective of forcing the model to
memorize specific image-text pairs within the dataset during training and
measuring memorization on these specific samples for reporting to the HPO.
First, we identified the most frequently occurring tokens within all text
captions in the MIMIC dataset (Figure 5). As a next step, we selected 100 text
captions having the highest count of these particular tokens. These 100 image-
text pairs constituted our memorization subset. During training, this subset
was duplicated 50 times resulting in 5000 additional samples.

###  C.2 Results with Random Word Addition (RWA)

Random Word Addition (RWA) has been proposed as a successful mitigation
strategy in [20]. The method involves replacing tokens/words in the caption
with a random word. This method turned out to be the most successful
mitigation strategy in the analysis performed in [20]. However, our results
indicate that RWA often leads to a substantial decrease in generative
performance. A simple explanation for this observation is that medical
vocabulary, and hence prompts in a medical dataset, are very different from
those in a dataset containing natural images. Since the stable diffusion
pipeline uses a tokenizer trained on natural prompts, it often lacks
significant knowledge about the specific medical terminology used in medical
datasets. RWA works by generating a random integer (between 0 and 49,400) and
decoding that integer into the corresponding text token using the tokenizer.
With a non-medical tokenizer, the random word that is added often decreases
the quality of the caption which in turn significantly impacts the textual
guidance in generating new images. Although RWA did decrease memorization in
certain cases in our experiments, it came at the cost of lower generation
quality.

##  Appendix D Elucidating the Parameter Search Space

Our framework performs a hyperparameter search to identify the optimal
locations for fine-tuning or model components for adding specific PEFT
parameters. This process utilizes a binary mask where each element’s position
corresponds to a specific component in the Stable Diffusion U-Net. The
element’s value (0 or 1) indicates whether the parameter at that location
should be frozen or fine-tuned. A similar investigation was conducted in [61],
which found that modulating the placement of PEFT parameters around cross-
attention blocks in the U-Net yields the best performance. Building on this
insight, we constrain our search space to include both self-attention and
cross-attention blocks in the U-Net for each PEFT method.

Stable Diffusion U-Net Design:  The Stable diffusion U-Net is a modified
version of the original model [54] with additional self-attention and cross-
attention blocks for textual guidance. Architecturally, the model is divided
into 4 down-blocks, 1 mid-block, and 4 up-blocks. The self-attention and
cross-attention blocks are contained in the mid-block, more specifically in
the first and the last three down and up blocks respectively.

###  D.1 Search Space for SV-DIFF

SV-DIFF [28] starts by conducting a Singular Value Decomposition (SVD) on the
weight matrices of a pre-trained diffusion model as a one-time calculation.
During fine-tuning, the original weight matrices are kept frozen, and only the
spectral shift parameters are updated, resulting in a very small number of
parameters being adjusted (0.02% of the total parameters in the U-Net). In its
original setup, SV-DIFF is applied to the weight matrices of all U-Net layers
(including Conv1D, Conv2D, Linear, Embedding, LayerNorm, and GroupNorm).
Building on the brief investigation in [28] regarding the fine-tuning of
spectral shifts within specific parameter subsets of the U-Net, we define our
search space to include all self-attention and cross-attention blocks within
the U-Net.

Given that there are 6 attention blocks (3 self-attention and 3 cross-
attention) in both down-blocks and up-blocks, and 1 attention block in the
mid-block, our search space for the binary mask is $\Omega\in\\{0,1\\}^{13}$.
For a binary mask of length 13, the first 6 elements represent the down-
blocks, the next 6 elements denote the up-blocks, and the other element
corresponds to the mid-block.

###  D.2 Search Space for DiffFit

DiffFit PEFT strategy [29] introduces learnable scale parameters $\gamma$ into
the diffusion model blocks. During fine-tuning, the majority of the model is
frozen and only the bias, normalization, class-condition module and the
learnable parameters $\gamma$ are updated.

Similar to SV-DIFF, we explore the various positions for adding DiffFit
parameters in the attention blocks of the U-Net. Given that there are 6
attention blocks (3 self-attention and 3 cross-attention) in both down-blocks
and up-blocks, and 1 attention block in the mid-block, our search space for
the binary mask is $\Omega\in\\{0,1\\}^{13}$. For a binary mask of length 13,
the first 6 elements represent the down-blocks, the next 6 elements denote the
up-blocks, and the other element corresponds to the mid-block.

##  Appendix E Experimental Settings for Hyperparameter Optimization (HPO)

Parameter Sampler: We use the NSGA-II algorithm [57] to sample hyperparameter
values. Initially, an initial population of random solutions is generated. The
objective functions are then evaluated on each solution in the population.
Finally, the candidate solutions are sorted according to Pareto dominance. The
algorithm maintains diversity by giving preference to solutions in less
crowded regions.

Bi-level Optimization Using NSGA-II: In each trial, the NSGA-II sampler
selects values for the binary mask to be used for fine-tuning. After fine-
tuning, we obtain a value for the objective metric associated with the sampled
mask values. Over numerous trials, the sampler identifies which values better
optimize the objective metric and gives these values higher preference in the
sampling process. In the case of single-objective optimization, pruning
techniques such as Successive Halving (SH) can be employed for early
termination of unpromising trials.

Scalability to Large Datasets:  Employing our framework on the MIMIC dataset
consisting of 110K training samples enabled us to test the scalability of our
framework to large datasets. We randomly sampled 1% (1,100 samples) from the
dataset for the HPO search, and the mask obtained on this subset was finally
used for fine-tuning on the entire dataset. Previously, such sampling
strategies for conducting HPO search on large datasets have been shown to be
effective [47, 62, 63].

###  E.1 Analysis of the HPO

In this section, we present an analysis of the HPO search over different PEFT
search spaces.

Convergence of the HPO:  Figure 6 illustrates the outer loop convergence in
the bi-level optimization for the specific case of the SV-DIFF parameter
search space. The plot demonstrates that the memorization metric $d^{mem}$
successfully converges for both parameter search spaces (SV-DIFF and DiffFit)
within the given number of outer loop iterations. This indicates that the HPO
search ran for a sufficient number of trials to reach the optimal objective
value.

Figure 6: Figure depicting the optimization history of the memorization metric
during the HPO search for SV-DIFF method.

## NeurIPS Paper Checklist

  1. 1.

Claims

  2. Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

  3. Answer: [Yes]

  4. Justification: Yes. The abstract claims to propose a bi-level optimization framework that selects optimal parameter subsets to balance both generation quality and mitigate memorization in text-to-image diffusion models. These claims have been justified and experimentally shown throughout the paper.

  5. 2.

Limitations

  6. Question: Does the paper discuss the limitations of the work performed by the authors?

  7. Answer: [Yes]

  8. Justification: Yes, the limitations have been discussed in Section 4.7.

  9. 3.

Theory Assumptions and Proofs

  10. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

  11. Answer: [N/A]

  12. Justification: The paper does not contain theoretical results. However, sufficient equations have been provided at appropriate sections.

  13. 4.

Experimental Result Reproducibility

  14. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

  15. Answer: [Yes]

  16. Justification: All experimental details have been provided in section 4. Additionally, the source code would be open-sourced for better reproducibility.

  17. 5.

Open access to data and code

  18. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

  19. Answer: [Yes]

  20. Justification: All the experimental details have been provided in section 4. Additionally, the source code would be open-sourced for better reproducibility.

  21. 6.

Experimental Setting/Details

  22. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

  23. Answer: [Yes]

  24. Justification: All the experimental details have been provided in section 4. Additionally, the source code would be open-sourced for better reproducibility.

  25. 7.

Experiment Statistical Significance

  26. Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

  27. Answer: [N/A]

  28. Justification: The experiments have been run once due to computational costs, showing large improvements compared to the baselines.

  29. 8.

Experiments Compute Resources

  30. Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

  31. Answer: [Yes]

  32. Justification: A detailed analysis on the compute cost incurred in each experiment has been discussed and presented in Appendix B.

  33. 9.

Code Of Ethics

  34. Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?

  35. Answer: [Yes]

  36. Justification: Yes, the research conducted in this work abides by the NeurIPS Code of Ethics.

  37. 10.

Broader Impacts

  38. Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

  39. Answer: [Yes]

  40. Justification: The broader impacts have been discussed in Section 4.7.

  41. 11.

Safeguards

  42. Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

  43. Answer: [N/A]

  44. Justification: Addition of safeguards are not applicable to the research presented in this work.

  45. 12.

Licenses for existing assets

  46. Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

  47. Answer: [Yes]

  48. Justification: All the original authors and creators of the models, algorithms, and datasets used in the paper have been appropriately cited in the paper.

  49. 13.

New Assets

  50. Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

  51. Answer: [N/A]

  52. Justification: The paper does not release any new assets.

  53. 14.

Crowdsourcing and Research with Human Subjects

  54. Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

  55. Answer: [N/A]

  56. Justification: No crowdsourcing was conducted for this work.

  57. 15.

Institutional Review Board (IRB) Approvals or Equivalent for Research with
Human Subjects

  58. Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

  59. Answer: [N/A]

  60. Justification: IRB approval was not required in this study.