1 Introduction↩︎

Generative diffusion models have emerged as an effective approach for synthesizing high-quality images that resemble training data and capture underlying concepts. However, this expressiveness also raises significant safety and ethical concerns, as such models may generate inappropriate content, including copyrighted material [1] or harmful content [2]. While policy-makers have made progress in regulating the removal of undesirable data or concepts from trained models [3], [4], it remains an open problem to develop effective technical methods to achieve this goal.

Recently, the classical concept of machine unlearning has been introduced in the context of diffusion models [5]. A standard approach aims to retain the utility of pretrained models, ensuring high-quality image generation, while simultaneously preventing the generation of specific harmful data or concepts—two inherently conflicting objectives. This principle has motivated recent empirical advances, including concept unlearning [6], [7] and data unlearning [8], [9]. Often, these works employ the simple practice of balancing two conflicting objectives by combining them with weighted sums. However, such weight-based methods are inherently ad hoc, as the weights are treated as hyperparameters, making the approach sensitive and poorly generalizable across scenarios.

In this work, we take a principled constrained optimization approach to address the conflicting objectives in unlearning. To do so, we view retaining the utility of trained models as minimizing the distance between the trained model and the pretrained model, while preventing harmful generation is formulated as additional constraints that push the trained model away from specific data or concept distributions. As a result, the tradeoff between these objectives can be systematically characterized by a Lagrangian formulation. To measure the distance between two models, we employ reverse and forward KL divergences, which broadly correspond to concept and data unlearning, respectively. We summarize our key contributions in three aspects below.

1. Problem Formulation. We formulate a distribution optimization problem that minimizes the reverse KL divergence to a pretrained model, while constraining the reverse KL divergences to the distributions associated with the concepts to be unlearned to remain above user-specified thresholds. Formulation 1 corresponds to unlearning at the concept level. We then consider an analogous optimization problem based on the forward KL divergence 2 , which corresponds to unlearning at the data level. Finally, we formulate a distribution optimization problem 3 that constrains the expected likelihood of the concepts to be unlearned to remain below user-specified thresholds. This constrained framework extends the unlearning objectives, thereby revealing their commonalities and differences; see Section 1.1 for exact formulations and discussion.

2. Theoretical Analysis. Although the reverse and forward KL unlearning problems 1 and 2 are nonconvex, we exploit the convexity of the image of a non-atomic vector measure to prove that they are strongly dual, rather than relying on standard convex duality. In contrast, strong duality for the likelihood unlearning problem 3 directly follows from its convexity. In all cases, we explicitly characterize the optimal solutions, which serve as principled unlearning targets; see Figure 1 for visualizations and discussion. A key implication of strong duality is that the unlearning problems can be solved equivalently in the dual domain. This enables our development of primal–dual algorithms instantiated with diffusion models for each formulation.

3. Empirical Results. We demonstrate the effectiveness of our constrained optimization approach for concept and data unlearning in a text-to-image diffusion model. In concept unlearning, our method achieves the same degree of unlearning while deviating less from the pretrained model compared to equal-weights approaches. In data unlearning, assigning the same weight to every sample to be unlearned leads to excessive deviation from the pretrained model, whereas our method learns optimal weights that ensure unlearning while minimizing deviation. Furthermore, we show that our novel likelihood-based unlearning method achieves a better tradeoff between unlearning the target concept and preserving the retained concept. This allows us to perform strong unlearning without deviating significantly from the parts of the distribution we wish to retain.

1.1 Summary of our problem formulations↩︎

Within the optimization framework over the probability measure space \(\Delta\), we outline three principled unlearning targets and defer their exact characterizations to Section 2. Here, \(q\in\Delta\) denotes the data distribution or the distribution induced by the pretrained model to be retained, and \(q_{\mathrm u}^i\in\Delta\) denotes the \(i\)th concept or data distribution to be unlearned.

Reverse KL-Constrained Unlearning. We minimize the reverse KL divergence between the model \(p\) and the reference \(q\) to ensure that the model is close to the reference. To remove the \(m\) undesirable concepts, we impose \(m\) reverse KL divergence constraints, each encouraging the model to stay away from an undesirable concept encoded by \(q_{\text{u}}^i\). To ensure this separation, we enforce a lower bound \(b_i\) on each KL divergence by formulating a reverse KL-constrained distribution optimization problem, \[\label{eq:32unlearning32reverse32KL} \begin{array}{rl} \displaystyle\mathop{\mathrm{minimize}}_{p\,\in\,\Delta} & D_{\text{KL}}(p\,\Vert\, q) \\[0.2cm] \mathop{\mathrm{subject~to}}& D_{\text{KL}}(p\,\Vert\,q_{\text{u}}^i)\; \geq \; b_i \; \text{ for }\; i\in [m]. \end{array}\tag{1}\] Problem 1 aims to minimize deviation from the pretrained model \(q\)—thereby maximizing model utility—while forgetting undesirable concepts. In Section 2.1, we capture a solution to Problem 1 as a ratio of distributions, \[p_{\text{rev}}^\star(\cdot) \;\propto\; \frac{(q(\cdot))^{\alpha_0} }{\prod_{i\,=\,1}^m ( q_{\text{u}}^i(\cdot))^{\alpha_i}}\] where \(\{\alpha_i\}_{i\,=\,0}^m\) is a set of non-negative exponents that balance the relative importance of each model. The subscript \(\text{rev}\) stands for reverse KL. The solution \(p_{\text{rev}}^\star\) reflects unlearning by reducing the probability of sampling in regions where \(q_{\text{u}}^i(\cdot)\) is large. This idea generalizes the concept unlearning target in concept erasing [6] and negation generation [10].

Forward KL-Constrained Unlearning. We employ the forward KL (rather than reverse KL) divergence to measure the closeness between two distributions. To do so, we obtain a variant of Problem 1 in terms of forward KL, \[\label{eq:32unlearning32forward32KL} \begin{array}{rl} \displaystyle\mathop{\mathrm{minimize}}_{p\,\in\,\Delta} & D_{\text{KL}}(q\,\Vert\, p) \\[0.2cm] \mathop{\mathrm{subject~to}}& D_{\text{KL}}(q_{\text{u}}^i\,\Vert\,p)\; \geq \; b_i \; \text{ for }\; i \in [m]. \end{array}\tag{2}\] Since the forward KL divergence is an expectation over the data distribution, it essentially optimizes the log-likelihood of the data under the model—often approximated by the evidence lower bound in diffusion models [11]. Therefore, Problem 2 seeks to maximize the likelihood of the data to be retained while reducing the likelihood of the data to be forgotten. In Section 2.2, we characterize a solution to Problem 2 as a difference of distributions, \[p_{\text{fw}}^\star(\cdot) \;\propto\; q(x) \, - \, \sum_{i\,=\,1}^m \alpha_i \, q_{\text{u}}^i(\cdot)\] where \(\{\alpha_i\}_{i\,=\,1}^m\) is a set of non-negative weights that balance the relative importance of each distribution. The subscript fw stands for forward KL. The log-likelihood reduction of forgetting data is often referred to as data unlearning in the literature [7], [9]. Formulation 2 is distinguished by its weaker assumption: it requires only sample access to the unlearning distributions \(q_{\text{u}}^i\), whereas the reverse KL and likelihood formulations assume access to the associated diffusion models.

Likelihood-Constrained Unlearning. In addition to the nonconvex reverse KL divergence constraints in Problem 1 , we propose an alternative formulation that constrains the likelihood of the unlearning data under the model, \[\label{eq:32unlearning32alignment} \begin{array}{rl} \displaystyle\mathop{\mathrm{minimize}}_{p\,\in\,\Delta} & D_{\text{KL}}(p\,\Vert\, q) \\[0.2cm] \mathop{\mathrm{subject~to}}& \mathbb{E}_{p}[\, q_{\text{u}}^i\,]\; \leq \; \epsilon_i \; \text{ for }\; i \in [m]. \end{array}\tag{3}\] Problem 3 aims to minimize the deviation from the pretrained model \(q\) while directly reducing the likelihood of the unlearning concepts appearing in the generated samples. In Section 2.3 we capture a solution to Problem 3 , \[p_{\text{revl}}^\star(\cdot) \; \propto \; \frac{q(\cdot)}{{\rm e}^{\sum_{i\,=\,1}^m \alpha_i\, q_{\text{u}}^i(\cdot)}}\] where \(\{\alpha_i\}_{i\,=\,1}^m\) is a set of non-negative weights that balance the relative importance of each model. The subscript revl stands for reverse KL with likelihood. The exponential term downweights samples that have high likelihood under the unlearning distributions. Compared with the framework of constrained relative-entropy minimization [12], our formulation chooses the moment functions to be unlearning-distribution likelihoods, thereby turning the exponential tilt into a principled mechanism for pushing the learned distribution away from undesired concepts.

To illustrate, we use a three-Gaussian mixture example to compare the three unlearning solutions, as shown in Figure 1. While the formulations 1 and 3 both reduce the probability of sampling from the unlearning concepts, their resulting solutions differ significantly. In particular, reverse KL-constrained unlearning achieves this by pushing the other modes we wish to retain away from the unlearned concept, whereas likelihood-constrained unlearning does not. This occurs because, for a sample \(x\) where \(q_{\text{u}}^i\) is small, the solution \(p_{\text{rev}}^\star\) is more strongly altered by \(q_{\text{u}}^i\) than \(p_{\text{revl}}^\star\) (proportional vs exponential).

Background on Diffusion Models. A diffusion model consists of two stochastic processes: a forward process and a backward process. The forward process starts from a data point \(x_0\) and gradually adds noise according to \(x_{t} = \sqrt{1-\beta_t} x_{t-1} + \sqrt{\beta_t}w_t\) for \(t\in [T]\), where \(\beta_t\in (0,1)\) is the stepsize and \(w_t\sim\mathcal{N}(0,I)\) is the Gaussian noise. The backward process starts from pure noise \(x_T\sim\mathcal{N}(0,I)\) and iteratively removes noise. For instance, the DDPM sampler updates \(x_{t-1} = \frac{1}{\sqrt{\alpha_t}}(x_t + \eta_t s(x_t, t) + \sigma_t z_t)\), where \(\alpha_t = 1-\beta_t\), \(\eta_t>0\) is the stepsize, \(\sigma_t\) is the noise scaling, and \(z_t\sim\mathcal{N}(0,I)\) is the Gaussian noise. A crucial component of the backward process is the score function \(s(x_t,t)\): \(X\times[T] \to \mathbb{R}^d\); when it is specified accurately, the backward process converges to the data distribution [13], [14]. Throughout the paper, we focus on score-based diffusion models, denoting each probability measure \(p\) by its score function \(s_p\): \(X\times[T] \to \mathbb{R}^d\). We use \(p_0(\cdot; s_p)\) to denote the marginal distribution of the terminal sample \(x_0\), and \(p_{0:T}(\cdot; s_p)\) to denote the path distribution of the trajectory \(x_{0:T}\), both induced by the score function \(s_p\). We use analogous notation for the pretrained models \(q\) and \(q^i\).

2 Unlearning in Distribution Space↩︎

In Sections 2.1–2.3, we formally present, in distribution space, the three formulations presented in Section 1.1. We prove strong duality in all three cases, despite the nonconvexity of the KL-divergence constraints.

2.1 Reverse KL-Constrained Unlearning↩︎

We begin with Problem 1 . Let \(q\): \(\mathcal{B}(X)\to [0, 1]\) be the probability measure induced by a pretrained model, where \(\mathcal{B}(X)\) is a Borel \(\sigma\)-algebra over a bounded domain \(X \subseteq \mathbb{R}^d\). In addition to the pretrained model, we consider a set of \(m\) probability measures \(\{ q_{\text{u}}^i: \mathcal{B}(X)\to [0, 1] \}_{i\,=\,1}^m\), each modeling a concept that we wish the model to forget/unlearn.

To ensure well-posedness, we assume that the probability measures of interest satisfy certain admissibility conditions.

The absolute continuity ensures the boundedness of reverse KL divergence. It is practical to only consider bounded models, since diffusion models are represented by conditional Gaussian distributions. Problem 1 is a nonconvex optimization problem because the superlevel set of reverse KL divergence is not convex. As a result, we cannot invoke the convexity of the feasible set to conclude strong duality from convex duality analysis [15].

Since \(q_{\text{u}}^i\) is focal on a specific concept or class, we can always construct \(p\) such that most of its mass is located on samples with low probability under \(q_{\text{u}}^i\). This makes Assumption 2 practically achievable. Let \(p_{\text{rev}}^\star\) be a solution to Problem 1 and denote \(P_{\text{rev}}^\star = D_{\text{KL}}(p_{\text{rev}}^\star\,\Vert\, q)\). Denote \(\lambda \mathrel{\mathop{:}}=[\,\lambda_1,\ldots,\lambda_m\,]^\top\). Let the Lagrangian for Problem 1  be \[\begin{array}{rcl} L_{\text{rev}}(p, \lambda) & = & \displaystyle D_{\text{KL}}(p\,\Vert\, q) \, + \, \sum_{i\,=\,1}^m \lambda_i \left( b_i - D_{\text{KL}}(p\,\Vert\, q_{\text{u}}^i) \right) \end{array}\] and the associated dual function is given by \(D_{\text{rev}}(\lambda) = \min_{p\,\in\,\Delta} L_{\text{rev}}(p, \lambda)\), where the minimization is achieved at \(p_{\text{rev}}^\star(\cdot;\lambda)\). Let an optimal dual variable be \(\lambda_{\text{rev}}^\star \in \mathop{\mathrm{argmax}}_{\lambda\,\geq\,0} D_{\text{rev}}(\lambda)\), and the optimal value of the dual function be \(D_{\text{rev}}^\star \mathrel{\mathop{:}}=D_{\text{rev}}(\lambda^\star)\).

However, it is unclear when \(p_{\text{rev}}^\star(\cdot;\lambda)\) is a solution to Problem 1 . Due to nonconvexity, \(p_{\text{rev}}^\star(\cdot;\lambda)\) is not necessarily a solution to Problem 1 , even when we fix \(\lambda = \lambda_{\text{rev}}^\star\). To address this issue, we next prove that Problem 1 is strongly dual, despite its nonconvex feasible set.

We defer the proof of Theorem 1 to Appendix 7.1, and sketch our key technical approach below.

A sufficient condition for strong duality to hold for Problem 1 is that its epigraph is convex [16]. To show this, we observe that all probability measures of interest have continuous domains. This allows us to define a non-atomic vector measure in terms of the objective and constraint functions. A key step is a new application of Lyapunov’s convexity theorem [17], which proves the image of a non-atomic vector measure is convex. We stress that it is infeasible to directly apply Lyapunov’s convexity theorem, since the originally-constructed vector measure [17] is not necessarily a probability measure.

Having established strong duality, we evaluate a solution to Problem 1 given an optimal dual variable \(\lambda_{\text{rev}}^\star\) in Corollary 1. Denote \(\widehat\lambda_i \mathrel{\mathop{:}}=\frac{\lambda_i}{1-\boldsymbol{1}^\top\lambda }\). When \(\boldsymbol{1}^\top\lambda \neq 1\), we define an unlearning target \(q_{\text{u}}^{\dagger}(\cdot;\lambda)\) as a ratio of tilted distributions, \[\label{eq:32negation32neq} q_{\text{\normalfont u}}^{\dagger} (\cdot;\lambda) \; \mathrel{\mathop{:}}=\; \frac{1}{ Z^\dagger_{\text{\normalfont u}}(\lambda)} \frac{(q(\cdot))^{1/(1-\boldsymbol{1}^\top\lambda)}}{\prod_{i\,=\,1}^m(q_{\text{\normalfont u}}^i(\cdot))^{\widehat\lambda_i}}\tag{5}\] where \(Z_{\text{\normalfont u}}^{\dagger}(\lambda) \mathrel{\mathop{:}}=\int {(q(x))^{1/(1-\boldsymbol{1}^\top\lambda)}}/{\prod_{i\,=\,1}^m(q_{\text{\normalfont u}}^i(x))^{\widehat\lambda_i}} dx\). When \(\boldsymbol{1}^\top\lambda = 1\), 5 can be simplified in Appendix 7.2; however, this is a degenerate solution, so we omit it here. Since \(q\) and \(\{q_{\text{\normalfont u}}^i\}_{i\,=\,1}^m\) share the same support, the unlearning target 5 is a proper probability measure. Outside the support, \(\frac{0}{0} = 0\) by convention. Viewing \(q_{\text{\normalfont u}}^i\) as a concept distribution to be unlearned, 5 reweighs the pretrained model \(q\) by the reciprocal of the tilted distribution \(q_{\text{\normalfont u}}^i\). Hence, 5 excludes, or downweights, samples that are likely in \(q_{\text{\normalfont u}}^i\). Similar targets have been used in concept erasure [6] and negation composition [10].

The proof of Corollary 1 evaluates the partial minimization of \(L(p;\lambda)\) over \(p\in\Delta\), and we provide it in Appendix 7.2. When \(0\leq \boldsymbol{1}^\top\lambda_{\text{\normalfont rev}}^\star<1\), the solution ?? captures the effect of the concept unlearning through the distribution ratio 5 . An example of this is the concept erasing [6]. We note that when \(\boldsymbol{1}^\top\lambda_{\text{\normalfont rev}}^\star > 1\), the distribution ratio 5 is reversed; this leads to an infeasible unlearning target, which we therefore exclude. Hence, the reverse KL constraints effectively characterize concept unlearning.

2.2 Forward KL-Constrained Unlearning↩︎

We next move to Problem 2 . The forward KL objective ensures that the model remains likely under the pretrained model, while the \(m\) forward KL constraints reduce the model’s likelihood of generating samples from \(m\) undesirable distributions. To ensure the well-posedness, we assume the following admissibility conditions.

The absolute continuity ensures the boundedness of forward KL divergence. Problem 2 is a non-convex optimization problem because the superlevel set of forward KL divergence is nonconvex. Again, strong duality is not ensured by convex duality analysis [15].

Compared to Assumption 2, we can always find a \(p\) such that few of its mass is located on samples with high probability under \(q_{\text{u}}^i\). Thus, Assumption 2 is also practically achievable. Let \(p_{\text{fw}}^\star\) be a solution to Problem 2 and denote \(P_{\text{fw}}^\star = D_{\text{KL}}(q \,\Vert\, p_{\text{fw}}^\star)\). Denote \(\lambda \mathrel{\mathop{:}}=[\,\lambda_1,\ldots,\lambda_m\,]^\top\). Let the Lagrangian for Problem 2 be \[\begin{array}{rcl} L_{\text{fw}}(p, \lambda) & = & \displaystyle D_{\text{KL}}(q \,\Vert\, p) \, + \, \sum_{i\,=\,1}^m \lambda_i \left( b_i - D_{\text{KL}}( q_{\text{u}}^i \,\Vert\, p) \right) \end{array}\] and the associated dual function is given by \(D_{\text{fw}}(\lambda) = \min_{p\,\in\,\Delta} L_{\text{fw}}(p, \lambda)\), where the minimization is achieved at \(p_{\text{fw}}^\star(\cdot;\lambda)\). Let an optimal dual variable be \(\lambda_{\text{fw}}^\star \in \mathop{\mathrm{argmax}}_{\lambda\,\geq\,0} D_{\text{fw}}(\lambda)\), and the optimal value of the dual function be \(D_{\text{fw}}^\star \mathrel{\mathop{:}}=D_{\text{fw}}(\lambda_{\text{fw}}^\star)\). Similar to Theorem 1, we next show that \(p_{\text{fw}}^\star(\cdot;\lambda)\) is a solution to Problem 2 if \(\lambda = \lambda_{\text{fw}}^\star\). To do so, we prove that Problem 2 is strongly dual.

Theorem 2 has a similar proof as the one for Theorem 1. We explain the proof of Theorem 2 in Appendix 7.3.

With the strong duality in place, Corollary 2 captures an explicit solution to Problem 2 . When \({q(x)} - \sum_{i\,=\,1}^m \lambda_i{q_{\text{u}}^i(x)} \geq 0\) for any \(x\in X\) and \(\lambda\geq0\), we define an unlearning target \(q_{\text{u}}^\triangleleft(\cdot;\lambda)\) as a difference between a pretrained model and a mixture distribution to be unlearned, \[\label{eq:32mixture32minus} q_{\text{u}}^{\triangleleft}(\cdot;\lambda) \; = \; \frac{1}{Z_{\text{u}}^\triangleleft(\lambda)} \left( {q(\cdot)} \, - \, \sum_{i\,=\,1}^m \lambda_i \, {q_{\text{u}}^i(\cdot)} \right)\tag{7}\] where \(Z_{\text{u}}^\triangleleft(\lambda) \mathrel{\mathop{:}}=\int \left( {q(x)} - \sum_{i\,=\,1}^m \lambda_i{q_{\text{u}}^i(x)}\right) dx\). The unlearning target 7 defines a valid distribution by reducing the likelihood of generating samples from a mixture. This target is similar to the data-unlearning [9], [18], which runs gradient descent on the retained dataset and gradient ascent on the unlearning dataset.

We provide the proof of Corollary 2 in Appendix 7.4.

2.3 Likelihood-Constrained Unlearning↩︎

We now consider Problem 3 . A key distinction from Problem 1 is that it constrains the likelihood of generating samples from undesirable concepts, rather than the KL divergences. The likelihood constraints force the model to generate samples from the unlearning concepts less often, with each bounded by \(\epsilon_i\). This constraint is intuitive since it prevents the model from assigning high probability to regions associated with the unlearning distributions. Since the likelihood-based feasible set is convex, we can directly conclude strong duality [15].

Let \(p_{\text{revl}}^\star\) be a solution to Problem 3 and denote \(P_{\text{revl}}^\star = D_{\text{KL}}(p_{\text{revl}}^\star\,\Vert\, q)\). Denote \(\lambda \mathrel{\mathop{:}}=[\,\lambda_1,\ldots,\lambda_m\,]^\top\). Let the Lagrangian for Problem 3 be \[\begin{array}{rcl} L_{\text{revl}}(p, \lambda) & = & \displaystyle D_{\text{KL}}(p\,\Vert\, q) \, + \, \sum_{i\,=\,1}^m \lambda_i \left(\, \mathbb{E}_{p} [ \,q_{\text{u}}^i \,] - \epsilon_i \,\right) \end{array}\] and the associated dual function is given by \(D_{\text{revl}}(\lambda) = \min_{p\,\in\,\Delta} L_{\text{revl}}(p, \lambda)\), where the minimization is achieved at \[\label{eq:32likelihood32exp} p_{\text{revl}}^\star(\cdot;\lambda) \;=\; \frac{1}{Z_{\text{u}}^{\circ}(\lambda)}\, q(\cdot)\, {\rm e}^{-\sum_{i\,=\,1}^m \lambda_i \, q_{\text{u}}^i(\cdot)}\tag{8}\] where \({Z_{\text{u}}^{\circ}(\lambda)} \mathrel{\mathop{:}}=\int q(x) {\rm e}^{-\sum_{i\,=\,1}^m \lambda_i q_{\text{u}}^i(x)} dx\). Let an optimal dual variable be \(\lambda_{\text{revl}}^\star \in \mathop{\mathrm{argmax}}_{\lambda\,\geq\,0} D_{\text{revl}}(\lambda)\), and the optimal value of the dual function be \(D_{\text{revl}}^\star \mathrel{\mathop{:}}=D_{\text{revl}}(\lambda_{\text{revl}}^\star)\).

The proof of Theorem 3 is based on the convex duality analysis [15] in Appendix 7.5. Given an optimal dual variable \(\lambda_{\text{revl}}^\star\geq 0\), the optimal model \(p_{\text{revl}}^\star(\cdot;\lambda_{\text{revl}}^\star)\) that we describe in 8 characterizes the effect of the concept unlearning through the negative exponent \(-\sum_{i\,=\,1}^m \lambda_i q_{\text{u}}^i\). Regions with high likelihood under the unlearning distributions are downweighted exponentially.

3 Unlearning for Diffusion Models↩︎

We instantiate the three formulations in Section 1.1 by letting all distributions be induced by diffusion models.

3.1 Reverse KL-Constrained Unlearning↩︎

We specify the probability measures \(p\), \(q\), and \(q_{\text{u}}^i\) in Problem 1 as the point distributions \(p_{0}(\cdot; s_p)\), \(q_{0}(\cdot; s_q)\), and \(q_{0}^i(\cdot; s_q^i)\), respectively, where \(s_p\), \(s_q\), and \(s_q^i\) denote the corresponding score functions. We cast Problem 1 into a constrained optimization with functional constraints, \[\begin{align} \label{eq:32unlearning32reverse32KL32diffusion32models95point-wise95score} \displaystyle \mathop{\mathrm{minimize}}_{s_p\,\in\,\mathcal{S}} \;\; & D_{\text{KL}}(p_0(\cdot; s_p)\,\Vert\, q_0(\cdot; s_q)) \\ \mathop{\mathrm{subject~to}}\;\; & D_{\text{KL}}(p_0(\cdot; s_p)\,\Vert\, q_0^i(\cdot; s_q^i)) \geq b_i \, \text{ for }\, i\in [m].\nonumber \end{align}\tag{10}\] Problem 10 is a nonconvex optimization problem because the mapping from the score function to the KL divergence is nonlinear. Nevertheless, when viewed as a function of the underlying probability measures, it has the same structure as Problem 1 . To exploit this structure, we assume that the score function class is sufficiently expressive to represent any target point distribution within the class of interest. This is mild, as diffusion models are implemented as overparameterized deep neural networks.

Assumption 6 is reasonable, because score-based backward processes converge to a data distribution under very mild conditions [13], [14]. Assumption 7 resembles the rationale behind Assumption 2. Let \(s_{\text{rev}}^\star\) be a solution to Problem 10 and denote \(\widehat P_{\text{rev}}^\star = D_{\text{KL}}(p_{0}(\cdot; s_{\text{rev}}^\star)\,\Vert\, q_{0}(\cdot; s_q))\). The Lagrangian for Problem 10 is given by \(\widehat L_{\text{rev}}(s_p, \lambda) = L_{\text{rev}}(p_{0}(\cdot; s_p), \lambda)\), and its dual function is given by \(\widehat D_{\text{rev}}(\lambda) \mathrel{\mathop{:}}=\min_{s_p\,\in\,\mathcal{S}} L_{\text{rev}}(p_{0}(\cdot; s_p), \lambda)\). Let an optimal dual variable be \(\widehat\lambda_{\text{rev}}^\star \in\mathop{\mathrm{argmax}}_{\lambda\,\geq\,0} \widehat D_{\text{rev}}(\lambda)\), and the optimal value of the dual function be \(\widehat D_{\text{rev}}^\star \mathrel{\mathop{:}}=\widehat D_{\text{rev}}(\widehat\lambda_{\text{rev}}^\star)\).

We prove Theorem 4 in Appendix 8.1. Despite the nonlinearity of Problem 10 , we exploit the convexity of a non-atomic vector measure—now it is defined on the path space—to establish strong duality. Since the score function class is sufficiently expressive, we then carry the result over to the score function space using the saddle-point characterization of strong duality. We note that the convexity of a non-atomic vector measure is unaffected by the nonlinearity introduced in 10 when passing to point distributions.

Exploiting strong duality, we introduce a primal-dual algorithm for solving Problem 10 , alternating between minimizing the Lagrangian via gradient descent, and maximizing the dual function via dual subgradient ascent:

• Primal step is to solve the Lagrangian problem: \(s^+ \in \mathop{\mathrm{argmin}}_{s}\widehat{L}_{\text{rev}}(s,\lambda)\), where the Lagrangian \(\widehat{L}_{\text{rev}}(s,\lambda)\) evaluates the KL divergence to the target 5 ;

• Dual step updates the dual variable via \(\lambda^+ = \lambda + \eta \partial\widehat{D}_{\text{rev}}(\lambda)\), where the subgradient \(\partial\widehat{D}_{\text{rev}}(\lambda)\) is evaluated using the current model \(s\), and \(\eta\) is the stepsize.

See Appendix 9 for the full description.

3.2 Forward KL-Constrained Unlearning↩︎

Instead of the point distributions, we specify the probability measures \(p\), \(q\), \(q_{\text{u}}^i\) in Problem 2 as the path distributions \(p_{0:T}(\cdot; s_p)\), \(q_{0:T}(\cdot; s_q)\), and \(q_{0:T}^i(\cdot; s_q^i)\). Problem 2 can be cast as a constrained optimization problem, \[\begin{align} \label{eq:32unlearning32forward32KL32diffusion32models95path-wise95score} \displaystyle \mathop{\mathrm{minimize}}_{s_p\,\in\,\mathcal{S}} \;\; & D_{\text{KL}}(q_{0:T}(\cdot; s_q)\,\Vert\, p_{0:T}(\cdot; s_p)) \\ \mathop{\mathrm{subject~to}} \;\; & D_{\text{KL}}(q^i_{0:T}(\cdot; s^i_q)\,\Vert\, p_{0:T}(\cdot; s_p)) \geq b_i\, \text{ for }\, i\in [m]. \nonumber \end{align}\tag{11}\] Problem 11 is a nonconvex optimization problem for the same reason as Problem 10 . Similarly, we exploit the strong duality for Problem 2 to establish strong duality.

Assumption 8 is more restrictive than Assumption 6, as it uses joint distributions instead of marginal ones. Let \(s_\text{fw}^\star\) be a solution to Problem 11 and denote \(\widehat P_{\text{fw}}^\star = D_{\text{KL}}( q_{0:T}(\cdot; s_q)\,\Vert\,p_{0:T}(\cdot; s_{\text{fw}}^\star))\). The Lagrangian for Problem 11 is given by \(\widehat L_\text{fw}(s_p, \lambda) = L_\text{fw}(p_{0:T}(\cdot; s_p), \lambda)\), and its dual function is given by \(\widehat D_\text{fw}(\lambda) \mathrel{\mathop{:}}=\min_{s_p\,\in\,\mathcal{S}} L_\text{fw}(p_{0:T}(\cdot; s_p), \lambda)\). Let an optimal dual variable be \(\widehat\lambda_\text{fw}^\star \in\mathop{\mathrm{argmax}}_{\lambda\,\geq\,0} \widehat D_\text{fw}(\lambda)\), and the optimal value of the dual function be \(\widehat D_\text{fw}^\star \mathrel{\mathop{:}}=\widehat D_\text{fw}(\widehat\lambda_\text{fw}^\star)\).

The proof of Theorem 5 in Appendix 8.2 is similar to that of Theorem 4, except that the KL divergence is defined over a path distribution. The convexity of a non-atomic vector measure continues to hold. Leveraging strong duality, we introduce a primal-dual algorithm for solving Problem 11 :

• Primal step is to solve the Lagrangian problem: \(s^+ \in \mathop{\mathrm{argmin}}_{s}\widehat{L}_{\text{fw}}(s,\lambda)\), where the Lagrangian \(\widehat{L}_{\text{fw}}(s,\lambda)\) has the form of standard score-matching or noise prediction losses;

• Dual step updates the dual variable via \(\lambda^+ = \lambda + \eta \partial\widehat{D}_{\text{fw}}(\lambda)\), where the subgradient \(\partial\widehat{D}_{\text{fw}}(\lambda)\) is evaluated using the current model \(s\), and \(\eta\) is the stepsize.

See Appendix 9 for the full description.

3.3 Likelihood-Constrained Unlearning↩︎

As done in Section 3.1, we introduce the point distributions \(q_0^i(\cdot;s_q^i)\) to Problem 3 to rewrite it as \[\begin{align} \label{eq:32unlearning32alignment32diffusion32models} \displaystyle\mathop{\mathrm{minimize}}_{s_p\,\in\,\mathcal{S}} \;\;& D_{\text{KL}}(p_{0:T}(\cdot; s_p)\,\Vert\, q_{0:T}(\cdot; s_q)) \\ \mathop{\mathrm{subject~to}}\;\;& \mathbb{E}_{x_0 \,\sim\, p_0}\left[\, q_0^i (x_0; s^i_q)\,\right] \; \leq \; \epsilon_i \; \text{ for } \; i \in [m].\nonumber \end{align}\tag{12}\] In contrast to Problem 3 , Problem 12 is nonconvex due to the nonlinear mapping from the score function to the KL divergence and the likelihood. We note that Problems 12 and 10 are similar, differing only in the form of the constraints. Critically, the convexity of a non-atomic vector measure continues to hold despite this change in how the constraints depend on the probability measures.

Let \(s_{\text{revl}}^\star\) be a solution to Problem 12 and denote \(\widehat P_{\text{revl}}^\star = D_{\text{KL}}(p_{0:T}(\cdot; s_{\text{revl}}^\star)\,\Vert\, q_{0:T}(\cdot; s_q))\). The Lagrangian for Problem 12 is given by \(\widehat L_{\text{revl}}(s_p, \lambda) = L_{\text{revl}}(p_{0:T}(\cdot; s_p), \lambda)\), and its dual function is given by \(\widehat D_{\text{revl}}(\lambda) \mathrel{\mathop{:}}=\min_{s_p\,\in\,\mathcal{S}} L_{\text{revl}}(p_{0:T}(\cdot; s_p), \lambda)\). Let an optimal dual variable be \(\widehat\lambda_{\text{revl}}^\star \in\mathop{\mathrm{argmax}}_{\lambda\,\geq\,0} \widehat D_{\text{revl}}(\lambda)\), and the optimal value of the dual function be \(\widehat D_{\text{revl}}^\star \mathrel{\mathop{:}}=\widehat D_{\text{revl}}(\widehat\lambda_{\text{revl}}^\star)\).

We prove Theorem 6 in Appendix 8.3. The proof is similar to that of Theorem 4, except for the likelihood constraints. Since the likelihood constraints can be lifted to path-space constraints, we can still apply the convexity of a non-atomic vector measure. Because of strong duality, we similarly use a primal-dual algorithm for solving Problem 12 :

• Primal step is to solve the Lagrangian problem: \(s^+ \in \mathop{\mathrm{argmin}}_{s}\widehat{L}_{\text{revl}}(s,\lambda)\), where the Lagrangian \(\widehat{L}_{\text{revl}}(s,\lambda)\) involves the computation of the KL divergence and the point probability. To estimate the probability \(q_0^i(x_0)\) for a sample \(x_0\), we apply the information-theoretic diffusion results [19] to compute \(\log q_0^i(x_0)\) by viewing the pretrained diffusion model as a Gaussian noise channel (see Appendix 8.4);

• Dual step updates the dual variable via \(\lambda^+ = \lambda + \eta \partial\widehat{D}_{\text{revl}}(\lambda)\), where the subgradient \(\partial\widehat{D}_{\text{revl}}(\lambda)\) is evaluated using the current model \(s\), and \(\eta\) is the stepsize.

See Appendix 9 for the full description.

4 Computational Experiments↩︎

We demonstrate our constrained unlearning approach through a series of computational experiments using Stable Diffusion v1.4 [20] as the base model. Sections 4.1, 4.2, and 4.3 cover likelihood-constrained concept unlearning, Forward KL-based removal of memorized samples, and reverse KL-based multi-concept unlearning, respectively.

4.1 Likelihood-Constrained Unlearning↩︎

We begin with the likelihood-constrained problem 12 by considering a simple setting in which a diffusion model is pretrained on a three-Gaussian mixture, and one component is to be unlearned. As a baseline, we consider concept erasing [6], which targets the unlearning distribution \({(q(\cdot))^{1+\eta}}/{(q_{\text{\normalfont u}}(\cdot))^{\eta}}\) for some \(\eta>0\). This is similar to the unlearning target 5 in the reverse KL case. We visualize the distributions of the pretrained and trained models for each approach in the three leftmost plots of Figure 4. At the same unlearning-likelihood level, our likelihood-constrained method achieves a smaller KL divergence to the retained modes, as shown in the right plot of Figure 4, where each point represents a constraint threshold or regularization weight for our method or the baseline.

We further demonstrate this improvement in a pretrained text-to-image model. To unlearn attributes from a concept, we set \(q = p_{\text{pre}}(\cdot\,|\,c)\), and \(q^i_{\text{u}} = p_{\text{pre}}(\cdot\,|\,c_i)\), where \(c\) is a retained concept (e.g., cowboy) that we aim to stay close to, and \(c_i\) denotes other concepts or biases associated with the concept \(c\) that we wish to remove (e.g., horse). We show the retention–unlearning tradeoff in the two rightmost plots in Figure 5, where retention performance is measured by the Kernel Inception Distance (KID) [21] (Lower KID means better retention performance) and the KL divergence to the pretrained model, respectively. Compared to the baseline, our likelihood-constrained approach achieves strong unlearning (low likelihood) while more effectively limiting deviation from the retained model.

This improved tradeoff can also be qualitatively observed. While both the concept erasure approach and our likelihood constrained approach successfully unlearn the undesired conept, in this case ‘horse’, (Figure 5 Left), our approach keeps the parts of the distribution that we wish to retain, closer to the pre-trained model (Figure 5 Middle left)

4.2 Forward KL-Constrained Unlearning↩︎

To demonstrate the effectiveness of the forward KL constraints in 11 , we work within the data unlearning paradigm [9], which seeks to prevent the model from generating specific samples that it may have memorized (e.g., for copyright reasons). We use the unconditional DDPM trained on the CelebA-HQ dataset as our pretrained model [22], and randomly select three images from the dataset to unlearn. As shown in Figure 2, our constrained unlearning approach effectively removes the given samples in the left plot, while retaining the capability of generating other concepts, as shown in the right plot. This significantly improves upon the unconstrained baseline [9], which essentially solves the Lagrangian problem for Problem 11 with equal weights assigned to all of the samples we wish to remove. We also show the retention–unlearning tradeoff in Figure 3, where the unlearning performance is measured by max SSCD [23] (Higher max SSCD indicates greater similarity to a forgotten sample and therefore weaker unlearning), the retention performance is measured by KID. We see that at the same level of unlearning, measured by max SSCD, our constrained unlearning approach deviates less from the pretrained model, as measured by KID. This is achieved by dynamically assigning large weights to the samples that are hard to unlearn and small weights to the easier samples.

4.3 Reverse KL-Constrained Unlearning↩︎

As in Section 4.1, we consider concept unlearning by setting \(q = p_{\text{pre}}(\cdot\,|\,c)\), and \(q^i_{\text{u}} = p_{\text{pre}}(\cdot\,|\,c_i)\) in Problem 10 , where \(c\) is a retained concept, and \(c_i\) denotes other concepts or biases associated with the concept \(c\) that we wish to remove. In Figure 6, we show the tradeoff between the worst-case unlearning performance (i.e., minimum KL divergences to unlearning concepts) and the retention performance (i.e., the KL divergence to the pretrained model), with each point representing a constraint threshold or regularization weight. We see that our constrained retention-unlearning tradeoff curve lies below the unconstrained baseline, achieving the same unlearning performance while deviating less from the pretrained model (Figure 6, Left). To confirm this, we also use the maximum CLIP text-to-image similarity score [24] among the unlearning concepts, observing a similar result (Figure 6, Right).

5 Conclusion↩︎

We have developed a constrained optimization framework that unifies unlearning as minimizing deviation from a pretrained model, subject to explicit separation constraints from the unlearning distributions. Despite the nonconvexity of these constraints, we prove that constrained unlearning has zero duality gap, which allows us to characterize optimal unlearning targets and develop primal–dual algorithms to approximate them. Empirically, we demonstrate our constrained unlearning approach on image generation tasks, showing a significantly-improved retention–unlearning tradeoff.

Limitations. Despite our theoretical characterizations and promising empirical results, additional experiments are needed to evaluate our method’s effectiveness on unlearning tasks beyond image generation. Further theoretical study is required to establish convergence and sample-complexity guarantees for the proposed primal-dual algorithms.

Impact Statement↩︎

This paper aims to improve the responsible behavior of deep generative models by ensuring compliance with unlearning requirements. By establishing a unified optimization framework, our work provides practical guidance for developing more reliable training algorithms, with potential impact across unlearning applications such as language, image, speech, and audio generation.

6 Related Work↩︎

Unlearning in diffusion models. Our constrained unlearning framework is related to prior work on unlearning in diffusion models. In this setting, given a pretrained diffusion model, unlearning aims to preserve the ability to generate diverse samples while avoiding the generation of samples that resemble specified samples or concepts to be forgotten. To handle these two conflicting objectives, existing approaches primarily adopt a classical regularization strategy, optimizing a weighted sum of retention and forgetting losses [8], [9], [25]–[31]. However, the weighted-sum objective has three key drawbacks: (i) the tradeoff between model utility and unlearning performance is tuned heuristically; (ii) the degree or likelihood of forgetting specific data or concepts is difficult to evaluate; and (iii) multiple unlearning objectives cannot be naturally incorporated within a single regularization term. In contrast, we formulate unlearning from the perspective of constrained distribution optimization. We treat the preservation of pretrained model utility as minimizing the distance between the trained and pretrained models, while expressing unlearning objectives as additional constraints that push the trained model away from specific data or concept distributions. This provides a more principled alternative to existing ad hoc approaches in [6], [9]. Our unlearning framework (i) offers a theoretical guarantee of an optimal tradeoff between model utility retention and unlearning performance, and (ii) allows for the direct imposition of multiple unlearning constraints. We also remark that our constrained unlearning approach is more tractable, both theoretically and practically, than the bilevel optimization approaches in [18], [32], [33], which often suffer from high computational costs and limited theoretical guarantees.

Diffusion models under constraints. Our work is also related to recent research on imposing constraints in diffusion models. Two types of constraints have been explored to control the samples generated by diffusion models: (i) hard constraints on individual samples [34]–[40] and (ii) soft constraints imposed in expectation [41]–[43]. Our constrained unlearning approach falls into the second category. In contrast to prior work [41], [43], our constrained unlearning problem is nonconvex even in the distribution space, which constitutes the main challenge we address. A closely related work is the constrained unlearning formulation of [7], assuming a particular class of unlearning distributions. In contrast, our constrained unlearning framework is more general and does not impose such additional restrictions.

7 Proofs in Section 2↩︎

7.1 Proof of Theorem 1↩︎

7.1.1 Strong duality for Problem 13↩︎

The proof is an application of Lyapunov’s convexity theorem in measure theory. We note that weak duality always holds, i.e., \(\widehat{D}_{\text{rev}}^\star \leq P_{\text{rev}}^\star\). Thus, it is sufficient to establish \(\widehat{D}_{\text{rev}}^\star \geq P_{\text{rev}}^\star\). We define an epigraph for Problem 13 as \[\Gamma \; \mathrel{\mathop{:}}=\; \big\{ (\gamma, t) \,\vert\, \exists\, p \in \mathcal{P}, \mathbb{E}_\mu \left[\, f(p(x)) \,\right] \leq \gamma_0, b_i - \mathbb{E}_\mu \left[\, f_i(p(x)) \,\right] \leq \gamma_i \text{ for } i=1,\ldots,m, \text{ and } \mathbb{E}_{\mu} \left[\, p(x) \,\right] - 1= \tau \big\}\] where \(\gamma \mathrel{\mathop{:}}=[\,\gamma_0, \gamma_1,\ldots, \gamma_m\,]^\top\).

With Lemma 1, we next show strong duality for Problem 13 . First of all, it is easy to check that \((P_{\text{rev}}^\star, 0, 0)\) is not an interior point of \(\Gamma\); otherwise, there always exists \(\epsilon>0\) such that \((P_{\text{rev}}^\star-\epsilon, 0, 0) \in \Gamma\), which contradicts the optimality of \(P_{\text{rev}}^\star\). By the supporting hyperplane theorem, there exists a non-zero vector \((y, z) \in \mathbb{R}^{m+2}\) such that \[\label{eq:32supporting32hyperplane} \sum_{i\,=\,0}^{m}y_i \gamma_i + z \tau \; \geq \; y_0 P_{\text{rev}}^\star \; \text{ for } (\gamma, \tau) \in \Gamma.\tag{19}\] We note that \(\Gamma\) is unbounded above, i.e., for any \((\gamma, \tau) \in \Gamma\), there exists \((\gamma', \tau) \in \Gamma\) such that \(\gamma'\geq \gamma\) elementwise. To ensure 19 , we must require \(y_i\geq0\) for all \(i\); otherwise, for some \(y_i<0\), we can set the left-hand side of the inequality 19 to be \(-\infty\) by increasing \(\gamma_i\). Furthermore, we can show that \(y_0 > 0\); otherwise, \(y_0=0\) simplifies 19 to be \(\sum_{i\,=\,1}^m y_i \gamma_i + z \tau \geq 0\) for \((\gamma, \tau) \in\Gamma\). This contradicts the strict feasibility assumption, since setting \((\gamma, \tau)\) to be the strict feasible point \(\bar{p}\) concludes that \(\gamma_i<0\) for all \(i\), and \(t=0\). Hence, \(y_0 > 0\), and thus, \[\gamma_0 + \sum_{i\,=\,1}^m \bar{y}_i \gamma_i + \bar{z} \tau\; \geq \; P_{\text{rev}}^\star \; \text{ for } (\gamma, \tau) \in \Gamma\] where \(\bar{y}_i \mathrel{\mathop{:}}=\frac{y_i}{y_0}\) and \(\bar{z} \mathrel{\mathop{:}}=\frac{z}{y_0}\). It further implies that \[\mathbb{E}_\mu \left[\, f(p(x)) \,\right] + \sum_{i\,=\,1}^m \bar{y}_i \left( b_i - \mathbb{E}_\mu \left[\, f_i(p(x)) \,\right]\right) + \bar{z} \left( \mathbb{E}_\mu \left[\,p(x) \,\right] -1 \right) \; \geq \; P_{\text{rev}}^\star\] or, equivalently, \(\widehat{L}_{\text{rev}}(p; \bar{y}, \bar{z}) \geq P_{\text{rev}}^\star\) for \(p\in\mathcal{P}\) and \(\bar{y}\geq 0\). Thus, minimization of \(\widehat{L}_{\text{rev}}(p; \bar{y}, \bar{z})\) over \(p\in\mathcal{P}\) leads to \(\widehat{D}_{\text{rev}}(\bar{y},\bar{z}) \geq P_{\text{rev}}^\star\), and thus \(\widehat{D}_{\text{rev}}^\star \geq P_{\text{rev}}^\star\).

7.2 Proof of Corollary 1↩︎

7.3 Proof of Theorem 2↩︎

7.4 Proof of Corollary 2↩︎

7.5 Proof of Theorem 3↩︎

8 Proofs in Section 3↩︎

8.1 Proof of Theorem 4↩︎

8.2 Proof of Theorem 5↩︎

8.3 Proof of Theorem 6↩︎

8.4 Likelihood Estimation↩︎

9 Algorithms↩︎

Here we detail the primal-dual algorithms discussed in Section 3 of the main paper.

9.1 Reverse KL-Constrained Unlearning Algorithm↩︎

As discussed in the main paper the algorithm alternates between minimizing the Lagrangian (primal) function, and maximizing the dual function.

Primal step: In the primal step we minimize the Lagrangian: \[\label{eq:32reverse32kl32alg32lagrangian} \mathop{\mathrm{minimize}}_{s_p}\;\widehat{L}_{\text{rev}}(s_p,\lambda) \; \mathrel{\mathop{:}}=\; D_{\text{KL}}(p_{0}(\cdot; s_p)\,\Vert\, q_{0}(\cdot; s_q)) -\sum_i \lambda_i \left( D_{\text{KL}}(p_{0}(\cdot; s_p)\,\Vert\, q^i_{0}(\cdot; s^i_q)) \; - b_i \right)\tag{27}\]

We can estimate the point-wise KL divergences in 27 by sampling trajectories from the model and using the following result (Lemma 2 from [43]):

We can then minimize the weighted sum of the KL divergences using standard gradient-based optimization techniques.

Dual step: In the dual step we update the dual multipliers via \(\lambda^+ = \lambda + \eta \partial\widehat{D}_{\text{rev}}(\lambda)\), using the fact that the constraint violations give us a subgradient \(\partial\widehat{D}_{\text{rev}}(\lambda)\) of the dual function. \[\label{eq:32dual32update32reverse32kl32123} \lambda^+_i = \left(\lambda_i - \eta_d \left( D_{\text{KL}}(p_{0}(\cdot; s_p)\,\Vert\, q^i_{0}(\cdot; s^i_q)) \; - b_i \right) \right)_+.\tag{29}\] A simple observation regarding the dual update is that if a constraint is not satisfied, we increase the corresponding multiplier.

9.2 Forward KL-Constrained Unlearning Algorithm↩︎

Similar to the reverse KL algorithm, we alternate between minimizing the Lagrangian (primal) function, and maximizing the dual function.

Primal step: In the primal step we minimize the Lagrangian: \[\mathop{\mathrm{minimize}}_{s_p}\; \widehat{L}_{\text{fw}}(s_p,\lambda) \; \mathrel{\mathop{:}}=\; D_{\text{KL}}(q_{0:T}(\cdot; s_q)\,\Vert\, p_{0:T}(\cdot; s_p)) -\sum_i \lambda_i \left( D_{\text{KL}}(q^i_{0:T}(\cdot; s^i_q)\,\Vert\, p_{0:T}(\cdot; s_p))\; - b_i \right).\]

The forward KL divergence \(D_{\text{KL}}(q_{0:T}(\cdot; s^i_q)\,\Vert\, p_{0:T}(\cdot; s_p))\) is the standard diffusion model training objective which we can estimate using samples from the distribution \(q\):

\[\label{eq:32score32matching} D_{\text{KL}}(q_{0:T}(\cdot; s_q)\,\Vert\, p_{0:T}(\cdot; s_p)) \; = \; \mathbb{E}_{x_0\sim q} \mathbb{E}_{q(x_t|x_0)} \left[ \Vert s_p(x_t, t) - \nabla \log q(x_t, t) \Vert^2_2 \right] + \text{constant}.\tag{30}\]

The first term on the RHS of 30 is the standard score matching objective with samples from the distribution \(q\). The Lagrangian then becomes: \[\widehat{L}_{\text{fw}}(s,\lambda) = \mathbb{E}_{x_0\sim q} \mathbb{E}_{q(x_t|x_0)} \left[ \Vert s_p(x_t, t) - \nabla \log q(x_t, t) \Vert^2_2 \right] -\sum_i \lambda_i \left( \mathbb{E}_{x_0\sim q_i} \mathbb{E}_{q(x_t|x_0)} \left[ \Vert s_p(x_t, t) - \nabla \log q^i(x_t, t) \Vert^2_2 \right]\; - b_i \right)\] which we can minimize using standard gradient-based optimization techniques using samples from \(q\) and \(q^i\).

Dual step: In the dual step, similar to the reverse KL algorithm we simply update the dual multipliers dependent on the constraint violations: \[\lambda^+_i = \left(\lambda_i - \eta_d \left( D_{\text{KL}}(q^i_{0:T}\,\Vert\, p_{0:T}(\cdot; s_p)) \; - b_i \right) \right)_+.\]

9.3 Likelihood-Constrained Unlearning Algorithm↩︎

The algorithm similarly alternates between primal and dual steps.

Primal step: For the primal step we need to minimize the Lagrangian:

\[\label{eq:32lagrangian32revl} \mathop{\mathrm{minimize}}_{s_p} \;\widehat{L}_{\text{revl}}(s_p,\lambda) \; \mathrel{\mathop{:}}=\; D_{\text{KL}}(p_{0:T}(\cdot; s_p)\,\Vert\, q_{0:T}(\cdot; s_q)) +\sum_i \lambda_i \left( \mathbb{E}_{p} \bigl[ q^i(x_0) \bigr]\; - \epsilon_i \right).\tag{31}\]

We can treat the likelihood of the unlearning distribution \(q^i(x_0)\) as a reward function that can be evaluated for any given sample using a diffusion model that samples from the distribution \(q^i(\cdot)\). We discuss this in detail in Lemma 2 the result of which we summarize here: \[\log q^i_{0}(x_0) - \log q_{0}(x_0) \; \approx \; \frac{1}{2} \sum_{t\,=\,0}^T \left( \mathbb{E}_{q(x_t|x_0)} \left[ \left\|{\nabla \log q(x_t) - s_q(x_t, t)}\right\|_2^2\right] - \mathbb{E}_{q(x_t|x_0)} \left[ \left\|{\nabla \log q^i(x_t) - s_{q^i}(x_t,t)}\right\|_2^2\right] \right)\] Note we have further approximated the integral in Lemma 2 with a sum over discrete diffusion time steps.

Now that we have a way to estimate the likelihood, the primal minimization becomes the well-known problem of aligning a diffusion model with a penalty function with the reward being the unlearning likelihood \(q(x_0)\) in this case. We can use the following result from [45] to compute the gradient of the expected likelihood: \[\nabla_\theta \mathbb{E}_{p_\theta(x_0)} \bigl[ q(x_0) \bigr] \;=\; \mathbb{E}_{p_\theta(x_{0:T})} \left[ q(x_0) \sum_{t=1}^{T} \nabla_\theta \log p_\theta(x_{t-1} \mid x_t) \right]\] which allows us to minimize the Lagrangian in 31 using gradient-based optimization.

Dual step: For the dual steps we update the dual variables via: \[\lambda_i \;=\; \left(\lambda_i + \eta_d \left( \mathbb{E}_{p} \bigl[ q^i(x_0) \bigr] - \epsilon_i \right)\right)_+.\]

10 Additional Implementation Details↩︎

Here we provide some additional implementation details regarding the experiments discussed in Section 4. All experiments were run on an NVIDIA B200 GPU.

10.1 Likelihood-Constrained Unlearning↩︎

10.1.1 Gaussian Mixture Unlearning↩︎

For the Gaussian mixture unlearning experiments we trained diffusion models with ResNet [46] networks as denoisers. We used a DDPM [22] scheduler with \(T = 100\) diffusion time steps. For the reverse KL baseline, we set the dual update learning rate to \(0\) to essentially minimize the Lagrangian with a fixed multiplier \(\lambda\), and we ran for different values of \(\lambda\) ranging from \(0.01\) to \(1.0\). For the Likelihood constrained approach we ran for different constraint thresholds ranging from \(0.05\) to \(1.0\). The likelihood is normalized such that the expected likelihood of the unlearn concept is approximately \(1.0\) when evaluated over samples of the initial pretrained model.

The constraint thresholds are chosen by multiplying the expected likelihood of the unlearning evaluated on samples from the pretrained model, by a factor.

We also note that since Lemma 2 allows the estimation of ratio of likelihoods, we consider the likelihood of the distribution that is proportional to \(q^i(x)/q_{pre}(x)\), where \(q_{pre}(x)\) is the pretrained model.

The full list of important hyperparameters is listed in Table 1.

10.1.2 Concept Unlearning with Stable Diffusion↩︎

We use a distilled version of Stable Diffusion v1.4 for these experiments. We fine-tune the parameters of LoRA adapters [47] added to the cross-attention layers in the diffusion U-Net following the esd-x approach from [6].

For computing KID to the retain distribution, we construct a retain set by sampling images from the pretrained model and keeping those whose likelihood/CLIP score for the unlearn concept is lower than a certain threshold i.e., samples in which the unlearn concept does not appear. The threshold is the expected likelihood/CLIP over samples from the pretrained model multiplied by a factor.

Similarly the likelihood constraint thresholds were chosen by multiplying the expected likelihood of the unlearn concept for the pretrained model (\(\bar \ell\)) by factors smaller than \(1\). We highlight that since Lemma 2 allows the estimation of ratio of likelihoods, we consider the likelihood of the distribution that is proportional to \(q_{pre}(x, c_i)/q_{pre}(x)\), where \(q_{pre}(x)\) is the unconditional pretrained distribution of Stable Diffusion and \(q_{pre}(x, c_i)\) is conditioned on the unlearning concept \(c_i\).

We also note that, unlike the Gaussian case where the likelihood constraint is reliably satisfied regardless of threshold level, for Stable Diffusion, optimizing the expected likelihood is challenging as the loss sometimes diverges. In Figure 5 we report only the convergent runs.

The full list of important hyperparameters is listed in Table 2.

10.2 Forward KL-Constrained Unlearning↩︎

For forward KL constrained unlearning we unlearn specific samples from a model pretrained on the CelebA-HQ dataset. Our implementation is based on modifying that of [9], namely, SISS without importance sampling, to incorporate dual updates.

For the unconstrained baseline, we use the same fixed multiplier \(\lambda\) for all the unlearning samples. We sweep \(\lambda\) over a range from \(0.05\) to \(0.5\) to observe SSCD values from \(1.0\) (almost no unlearning) to around \(0.3\) which by visual inspection represents complete unlearning of the memorized sample.

We compute the SSCD score between two images: the original image sample we wish to unlearn from the model, and the image obtained by adding noise to the unlearning sample up to time step \(t = 200\) and then denoising using the model to see how closely it can recover the original sample. This comparison is visually represented in Figure 2, Left.

The full list of important hyperparameters is listed in Table 3.

10.3 Reverse KL-Constrained Unlearning↩︎

Dual-Only Algorithm. For the reverse KL experiments we modify the algorithm slightly. We can skip the primal step in the algorithm discussed in Appendix 9.1, by observing the fact that we already know from section 2.1 that for a given \(\lambda\) the optimal primal minimizer is given by \[q_{\text{\normalfont u}}^{\dagger} (\cdot;\lambda) \; = \; \frac{1}{ Z^\dagger_{\text{\normalfont u}}(\lambda)} \frac{(q(\cdot))^{1/(1-\boldsymbol{1}^\top\lambda)}}{\prod_{i\,=\,1}^m(q_{\text{\normalfont u}}^i(\cdot))^{\widehat\lambda_i}}.\] Since we have access to the score functions of \(q(\cdot)\), and the scores of the individual unlearning distributions \(q^i(\cdot)\), we can bypass the primal step and directly do the dual updates in 29 .

This dual-only algorithm allows us to find the optimal multipliers \(\lambda^\star\) that satisfy the unlearning reverse KL constraints. If needed, the model weights can then be fine-tuned so that the distribution of generated samples fits \(q_{\text{\normalfont u}}^{\dagger} (\cdot;\lambda^\star)\). [43] propose a similar dual-only algorithm to find the optimal weights for diffusion model composition.

To choose the constraint thresholds, we first estimate the reverse KL between the pretrained model, and each of the individual unlearning concepts. Then we set the threshold as the average of these KL divergences, multiplied by a factor.

For the fixed equal weight, unconstrained baseline, we choose all dual multipliers to be equal, and across runs we sweep the multiplier values so that the sum of the dual multipliers, \(\lambda^\top 1\), goes from \(0\) to \(0.95\). Recall from Section 2.1 that to avoid degenerate solutions we need \(\lambda^\top 1< 1\).

References↩︎