Safe Latent Diffusion: Mitigating Inappropriate Degeneration in Diffusion Models

In plain words

Text-to-image models learn from huge collections of images scraped from the web. They therefore also learn what people post there, including nudity, violence and hateful or disturbing material. Asked for an ordinary picture, they sometimes produce such content anyway. This paper does two things. It builds a benchmark of real user prompts that tend to trigger such images, so the problem can be measured. It also proposes a way to steer the image while it is being generated: at each step the model is pushed away from a written description of unwanted content, and still follows the user's prompt.

The method works with the model's own understanding of what is inappropriate, so no separate classifier is added and nothing is retrained. The authors also point out that the same mechanism could be used for censorship, so anyone who uses it should say which content is being suppressed.

TL;DR

• Stable Diffusion, trained on unfiltered web data, produces inappropriate images for many ordinary-looking prompts §5 §6.
• Safe latent diffusion (SLD) adds a safety term to classifier-free guidance during sampling. The term pushes the prediction away from a text description of unwanted content §3 eq. 4 eq. 3.
• No training and no external classifier are needed §1 §3.
• Author-reported: on the I2P benchmark the overall probability of an inappropriate image falls from 0.39 to 0.09 at the strongest setting, with image quality and text alignment close to the baseline in a user study tab. 2 tab. 3.

Problem

Models trained on large unfiltered datasets reproduce and amplify the biases and degenerate content in those datasets §2. Stable Diffusion v1.4 was trained on a subset of LAION-5B and released without removing the annotated inappropriate content §2. The authors measure this on a new benchmark and find that a large share of generated images is inappropriate, in every category they test §6 tab. 2.

Post-hoc filters are the common alternative. Stable Diffusion ships a classifier that withholds images after generation. The authors judge its false-positive rate too high for the nuanced editing they need app. D. Negative prompts are the other obvious tool; the authors find they mitigate less than SLD and offer less control over how far the image moves §4 tab. 2. SLD is meant to act during generation, without retraining the model and without changing the prompt the user writes §1 §3.

Method

Background. Classifier-free guidance combines an unconditional noise prediction with a prompt-conditioned one, scaled by a guidance weight §3 eq. 2. SLD keeps this formula and adds a third prediction, conditioned on a text description of the unwanted concept S. The prediction is moved toward the prompt and away from the concept §3 eq. 3.

Training objective (background, not changed by SLD). The denoiser is trained with a weighted squared error on the clean sample x, conditioned on the prompt c_p §3 eq. 1.

\mathbb{E}_{\mathbf{x},\mathbf{c}_p,\epsilon,t}\left[ w_t \left\| \hat{\mathbf{x}}_\theta(\alpha_t \mathbf{x} + \omega_t \epsilon, \mathbf{c}_p) - \mathbf{x} \right\|_2^2 \right] eq. 1

Classifier-free guidance. The guided noise estimate uses guidance scale s_g §3 eq. 2:

\tilde{\epsilon}_\theta(\mathbf{z}_t, \mathbf{c}_p) = \epsilon_\theta(\mathbf{z}_t) + s_g \left( \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_p) - \epsilon_\theta(\mathbf{z}_t) \right) eq. 2

Step 1: the safety concept. The unwanted concept is a plain-text description. The default concept, given in the paper as a list of words, is: "hate, harassment, violence, suffering, humiliation, harm, suicide, sexual, nudity, bodily fluids, blood, obscene gestures, illegal activity, drug use, theft, vandalism, weapons, child abuse, brutality, cruelty" §4. The text can be replaced for other concepts §4.

Step 2: the safety-guided noise estimate. The prompt estimate is moved away from the concept by a safety term γ §3 eq. 3 eq. 4:

\bar{\epsilon}_\theta(\mathbf{z}_t, \mathbf{c}_p, \mathbf{c}_S) = \epsilon_\theta(\mathbf{z}_t) + s_g \left( \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_p) - \epsilon_\theta(\mathbf{z}_t) - \gamma(\mathbf{z}_t, \mathbf{c}_p, \mathbf{c}_S) \right) eq. 3

\gamma(\mathbf{z}_t, \mathbf{c}_p, \mathbf{c}_S) = \mu(\mathbf{c}_p, \mathbf{c}_S; s_S, \lambda) \left( \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_S) - \epsilon_\theta(\mathbf{z}_t) \right) eq. 4

Step 3: the element-wise scale μ. μ is applied per element. It is zero where the prompt estimate exceeds the concept estimate by at least λ, and otherwise a scale built from φ §3 eq. 5 eq. 6. The paper says the scale is clipped to avoid artefacts §3; see Pitfalls and Paper vs code for how the released code differs.

\mu(\mathbf{c}_p, \mathbf{c}_S; s_S, \lambda) = \begin{cases} \max(1, |\phi|) & \text{if } \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_p) \ominus \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_S) < \lambda \\ 0 & \text{otherwise} \end{cases} eq. 5

\phi = s_S \left( \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_p) - \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_S) \right) eq. 6

Larger λ and larger s_S both shift the sample further from the prompt and away from the concept §3.

Step 4: warm-up. The safety term is not applied until step δ. Before that, the sampler runs ordinary guidance. Larger δ means smaller changes to the composition §3 app. I.

Step 5: momentum. A momentum term ν accumulates the safety term over steps so that consistently guided dimensions move faster. Momentum starts at zero and is built up during warm-up even though the term is not applied then §3 eq. 7 eq. 8.

\gamma_t = \mu(\mathbf{c}_p, \mathbf{c}_S; s_S, \lambda) \left( \epsilon_\theta(\mathbf{z}_t, \mathbf{c}_S) - \epsilon_\theta(\mathbf{z}_t) \right) + s_m \nu_t eq. 7

\nu_{t+1} = \beta_m \nu_t + (1 - \beta_m) \gamma_t, \qquad \nu_0 = \mathbf{0} eq. 8

Here s_m is the momentum scale, in [0, 1], and β_m in [0, 1) sets how slowly momentum changes; larger β_m gives less volatile momentum §3.

Full sampling loop. The paper gives the procedure as Algorithm 1 alg. 1. The pseudocode below follows it. It differs from the paper's listing in naming and in one clip, noted in the comment.

# Safe Latent Diffusion sampler  [alg. 1]
# Inputs: prompt, concept, steps (T), guidance_scale (s_g),
#         concept_scale (s_S), threshold (lambda), warmup (delta),
#         momentum_scale (s_m), momentum_beta (beta_m)

def sld_sample(prompt, concept, steps):
    # Seeded Gaussian latents; momentum starts at zero  [alg. 1] [§3]
    latents = sample_noise(seed)
    momentum = zeros_like(latents)
    for t in range(steps):
        # Three noise estimates: unconditional, prompt, concept  [eq. 3]
        noise_uncond = predict(latents)
        noise_prompt = predict(latents, prompt)
        noise_concept = predict(latents, concept)
        # Element-wise difference between prompt and concept estimates  [eq. 5]
        diff = noise_prompt - noise_concept
        # Scale factor phi  [eq. 6]
        phi = concept_scale * diff
        # Scale mu: max(1, |phi|) where diff < threshold, else 0  [eq. 5]
        # The released code clips to min(1, |phi|) instead; see Pitfalls
        mu = where(diff < threshold, maximum(1, abs(phi)), 0)
        # Safety term with momentum  [eq. 7]
        gamma = mu * (noise_concept - noise_uncond) + momentum_scale * momentum
        # Momentum update uses the same gamma  [eq. 8]
        momentum = momentum_beta * momentum + (1 - momentum_beta) * gamma
        # Warm-up: apply the safety term only from step delta on  [§3] [alg. 1]
        if t >= warmup:
            noise = guidance_scale * (noise_prompt - noise_uncond - gamma)
        else:
            noise = guidance_scale * (noise_prompt - noise_uncond)
        latents = scheduler_step(latents, noise)
    return decode(latents)

Scheduler and model. The experiments use Stable Diffusion v1.4 with a discrete linear multistep (LMS) scheduler, β_start = 8.5e-4 and β_end = 0.012 app. D. The authors state the approach applies to any latent diffusion model that uses classifier-free guidance app. D.

Configurations. Four preset sets of parameters, from weakest to strongest, are given in tab. 1. The text gives recommended ranges for each parameter §4.

Hyperparameters

Results

Author-reported. Setup: Stable Diffusion v1.4 unless noted; I2P benchmark of 4703 prompts; 10 images per prompt; inappropriate = flagged by Q16 or NudeNet §5 §6 app. D. Hardware and model size: not stated. Baseline: unmodified Stable Diffusion, and negative prompting in one column.

The paper notes that the Q16 classifier is conservative and flags some harmless images, and that its false-positive rate is roughly the Max probability tab. 2. The paper reports the expected-max variance for SLD is large, with outliers §6.

Paper vs code

The official repository is ml-research/safe-latent-diffusion, owned by the authors' lab (ml-research), at commit a42923c3de0e4346bee3f61891a510bdcc2aedd2, MIT licensed. The paper's footnote points to the Hugging Face diffusers page for the same method §1; the repository README calls itself the official implementation of the paper. The code is a pipeline extending diffusers' Stable Diffusion pipeline sld_pipeline.py.

• Scale μ. Paper: max(1, |phi|) where the difference is below λ eq. 5 alg. 1. Code: clamp(|diff| * sld_guidance_scale, max=1), i.e. min(1, |phi|), so the scale is bounded above by 1 sld_pipeline.py. The paper's text says the scale is clipped §3, which matches the code, not the equation. The equation and the code disagree; the released pipeline ships the clipped form.
• Momentum update. Algorithm 1 as printed writes beta_m * nu_t (1 - beta_m) * gamma_t with no + between the two terms alg. 1. Eq. 8 has a + eq. 8, and the code adds the two terms sld_pipeline.py.
• Momentum in γ. Eq. 7 and the code both include s_m ν_t in the γ used for the momentum update sld_pipeline.py.
• Default parameters. The code defaults are δ=10, s_S=1000, λ=0.01, s_m=0.3, β=0.4 sld_pipeline.py. These match the Medium set in tab. 1. The README says the defaults are the medium configuration README.md.
• Guidance scale and steps. Code defaults: guidance 7.5 and 50 steps sld_pipeline.py. The paper does not state s_g for the I2P or quality runs (the caption that gave 7 is commented out in the source) §6.
• Concept text. Identical to the paper's list sld_pipeline.py §4.
• Stable Diffusion's own safety checker. The code runs it whenever one is loaded and returns per-image flags sld_pipeline.py. The paper says it did not use the built-in checker for its experiments app. D.
• Negative prompts. If a negative prompt is passed, the code uses it in place of the empty unconditional branch, so SLD's baseline changes sld_pipeline.py sld_pipeline.py.
• Scheduler. The code accepts DDIM, LMS and PNDM sld_pipeline.py. The paper's experiments use LMS app. D. The code does not pin the scheduler.
• Evaluation. The I2P classifier pipeline, Q16 and NudeNet scoring, and the user-study code are not in this repository. The paper links an evaluation repository ml-research/i2p §5. It was not cloned or read, so no claim is made about it.

No other discrepancies found.

Pitfalls

• Clip versus max. Eq. 5 gives a scale of at least 1, which can be large when s_S is in the thousands. The text says the scale is clipped §3, and the code clips it to at most 1 sld_pipeline.py. Pick one knowingly; the released pipeline ships the clipped form.
• Missing plus sign. The momentum line in Algorithm 1 has no + alg. 1. Implement eq. 8.
• Undefined operator. Eq. 5 uses an operator ⊖ that the paper does not define. The algorithm and the code use a plain element-wise subtraction of the prompt and concept estimates alg. 1 sld_pipeline.py.
• Threshold range. The text recommends λ in [0.0, 0.03] §4, but the Max preset uses λ = 1.0 tab. 1. λ is compared against raw noise-prediction differences; a value tuned for one model or guidance scale does not transfer automatically (the paper gives no such check).
• Concept scale. The text says s_S can be made arbitrarily high because μ is clipped §4, while the recommended range stops at 3000 §4 and the algorithm allows 5000 alg. 1. Larger values only raise the number of elements that are set to the clip, per the paper §4.
• Momentum. The paper reports artefacts above s_m = 0.5 app. I, and the text recommends at most 0.5 §4.
• Guidance scale is not stated for the main experiments §6. The ablation figures use s_g = 15 and s_S = 100 app. I. Expect to tune s_g yourself.
• Cost. The code runs three UNet forward passes per step (unconditional, prompt, concept) instead of two for plain guidance sld_pipeline.py sld_pipeline.py. This is read from the code; the paper reports no timing.
• Negative prompts interact with SLD. A user negative prompt replaces the unconditional branch that SLD subtracts from sld_pipeline.py. Results are not those of the paper's protocol if one is passed.
• Built-in checker. Loading the Stable Diffusion safety checker changes the output path sld_pipeline.py. The paper evaluated without it app. D.
• Long concept text. The concept is tokenised and truncated silently to the tokenizer length sld_pipeline.py. Keep the concept short enough to fit.
• Warm-up with other step counts. With fewer steps the momentum window shrinks §4. Keep δ at most half the steps §4.
• Measurement depends on two classifiers. Inappropriateness is decided by Q16 or NudeNet app. D. Q16 is described as conservative tab. 2. NudeNet only flags exposed genitalia app. D, so "sexual" results do not cover suggestive content.
• Bias is reduced, not removed. After SLD a correlation of r = 0.52 remains between the country-level nudity rates with and without SLD §6. Dataset-level fixes are still needed §6.
• Variance. The expected-max results show high variance under SLD, with outliers §6.
• Count mismatch. The paper gives 4703 prompts §5 and, in the toxicity paragraph, 4702 §5. Use 4703 as the dataset size.
• Disclosure. The authors ask that any deployment state which content is suppressed, since the same mechanism can be used for censorship §7.

Lineage

• Builds on: classifier-free guidance §3 (cited there as ho2022classifier), which SLD extends with a third noise estimate and a warm-up and momentum scheme §3.
• Compared against: Stable Diffusion v1.4 without safety measures tab. 2; negative prompting §4 tab. 2; Stable Diffusion v2.0 with training-data filtering app. E.
• Related text-based image editing methods the paper cites: Prompt-to-Prompt, Text2LIVE, Imagic, UniTune §2. The paper argues these need fine-tuning or extra components, which SLD does not §2.
• Superseded by: not stated.