Does Generation Require Memorization? Creative Diffusion Models using Ambient Diffusion

In plain words

Image generators built on diffusion models can reproduce pictures they were trained on, and they do this more when the training set is small. Earlier ways to stop the copying tend to make the images worse.

The idea is to split the generation process by how much noise is in the picture. The paper argues that copying is needed mainly at the low-noise end, where fine detail is decided. At the high-noise end, the model learns from copies of each training image that were noised once, in advance, so it never gets the clean original at those levels. Below a chosen noise level, training is the usual one.

The authors report that this copies less at the same image quality on small datasets, and that the gain shrinks as the dataset grows.

TL;DR

• Diffusion models trained on small datasets copy training images. The paper keeps standard denoising at noise levels below a chosen t_n and trains on one fixed noisy copy of each image above it §3 alg. 1.
• Author-reported: at about the same FID as DDPM, the DINOv2-based memorization measure drops on CIFAR-10, FFHQ and tiny ImageNet with 300 to 3,000 images tab. 1 §5.1.
• The official code routes samples by a sampled noise level, not by whether they come from the clean or noisy set, and it contains no text-conditional training entry point. Its license is non-commercial. Use it to read, not to copy loss.py.

Problem

The optimal solution of the standard denoising objective, for a finite training set, replicates the training points §2.2. Models that generalise are therefore the ones that fail to fit the training loss perfectly §2.2. Known mitigations for memorization, such as corrupting the training images, often lower image quality §1 §5.1. The question here is whether a model can have both low memorization and high quality §1.

The paper's theory argues that memorization matters mainly at low noise levels, where fine detail lives, and that the heavy-tailed frequency structure of data drives the need to memorize §4.1 app. A.4.

Method

The forward process is the variance-preserving corruption X_t = sqrt(1 - sigma_t^2) X_0 + sigma_t Z eq. 1 §2.1. Standard denoising regresses x_0 at every noise level eq. 3 §2.1. For a noisy copy at level t_n, the ambient score-matching loss eq. 6 §2.3 learns the same conditional expectation from noisy data only, and it applies only for t > t_n.

The training recipe is alg. 1 §3:

1. Before training, add noise at level t_n once to every training image. This set S_tn is created once and not redrawn each epoch §3 alg. 1, line 1.
2. Each batch is drawn uniformly from the union of clean and noisy samples alg. 1.
3. For a noisy sample, draw t uniformly on (t_n, T), add extra noise, and apply the ambient loss alg. 1 eq. 6.
4. For a clean sample, draw t uniformly on (0, t_n), add noise, and apply the standard denoising loss alg. 1 eq. 3.
5. Average the per-sample losses over the batch and take a gradient step alg. 1.

# Limited-data training with one noisy copy per image  [alg. 1]
# Input:  network h, clean set S, noise level t_n, schedule sigma(t),
#         batch size B, horizon T
# Output: trained denoiser h

def train(h, S, t_n, sigma, B, T):
    # Noise each clean image once at level t_n, before training  [§3]
    S_tn = [noise_once(x0, sigma(t_n)) for x0 in S]
    while not converged:
        # Draw a batch uniformly from clean and noisy samples  [alg. 1]
        batch = sample(S + S_tn, B)
        loss = 0
        for x in batch:
            if x belongs_to S_tn:
                # Diffusion time for noisy inputs  [alg. 1]
                t = uniform(t_n, T)
                # Extra noise on top of the stored noisy copy  [alg. 1]
                x_t = add_noise_from(x, t_n, t)
                # Ambient score-matching loss  [eq. 6]
                loss += ambient_loss(h, x_t, t, x, t_n)
            else:
                # Diffusion time for clean inputs  [alg. 1]
                t = uniform(0, t_n)
                x_t = noise_clean(x, t)
                # Standard denoising loss  [eq. 3]
                loss += denoising_loss(h, x_t, t, x)
        loss = loss / B
        # Gradient step on the averaged loss  [alg. 1]
        h = update(h, loss)
    return h

Key equations, as printed in the paper:

X_t = \sqrt{1 - \sigma_t^2}\, X_0 + \sigma_t Z, \quad Z \sim \mathcal N(0, I_d) eq. 1

J(\theta) = \mathbb E_{x_0}\, \mathbb E_{(x_t, t) \mid x_0}\left[ \| h_\theta(x_t, t) - x_0 \|^2 \right] eq. 3

J_{\mathrm{ambient}}(\theta) = \mathbb E_{x_{t_n}}\, \mathbb E_{(x_t, t) \mid x_{t_n}} \left[ \left\| \frac{\sigma_t^2 - \sigma_{t_n}^2}{\sigma_t^2 \sqrt{1 - \sigma_{t_n}^2}}\, h_\theta(x_t, t) + \frac{\sigma_{t_n}^2}{\sigma_t^2} \sqrt{\frac{1 - \sigma_t^2}{1 - \sigma_{t_n}^2}}\, x_t - x_{t_n} \right\|^2 \right] eq. 6

Intuition the authors give §3: the ambient loss does not point the score at training points, because it predicts the noisy inputs, and noisy versions are harder to memorize than clean ones. Below t_n the learned score matches standard training. Limits: the paper claims no privacy guarantee and no end-to-end optimality result §6.

Theory, in brief §4. For the ambient-trained model, the distribution reached at t = t_n under the optimal score is the uniform distribution over the noisy set S_tn [Lemma 4.1]. Under the Feldman-style subpopulation model with heavy-tailed frequencies, the generalisation error is at least the unseen-population error plus a term scaled by tau_1 times the loss on subpopulations seen once [Theorem 4.2 informal] §4.1. tau_1 is defined as a ratio of expectations over the mixing weights §4.1. Heavy-tailedness is the condition that a constant mass of mixing weights lies in [1/(2n), 1/n] eq. 7 §4.1.3. Adding noise merges nearby components, which lightens these tails §4.1.3 app. A.4. The concatenation claim (ambient trajectory up to t_n, memorized DDPM below it, need not memorize) is stated as a conjecture §4.

Hyperparameters

Values come from the paper unless marked "code". Code values are read from the official repository at code_commit. Items the paper does not state are marked not stated.

Results

Author-reported throughout. Setup: unconditional models trained from scratch on CIFAR-10, FFHQ and tiny ImageNet with 300, 1,000 and 3,000 images; FID over 50,000 samples; memorization by DINOv2 similarity §5.1. The paper does not state the unit of the "S>" rows; they are reproduced as printed.

Table 1 (DDPM vs ours at matched FID, tab. 1):

Bold marks the better value as printed. The 300-image tiny-ImageNet column is absent by design app. D.

Table 2 (mitigation comparison, FFHQ; tab. 2). "Masking" is linear corruption (random inpainting) and "Noise" is training on noisy data only, as the paper describes them §5.1:

Table 3 (text-conditional, Stable Diffusion v2 fine-tuned on 10k LAION pairs; tab. 3 app. D). Sim is the mean nearest-neighbour similarity, 95% the 95th percentile, CLIP the image-text alignment:

Figures: fig. 1 shows FID-memorization pairs for a sweep of sigma_tn on 300 FFHQ images. fig. 2 shows qualitative text-conditional samples. fig. 3 shows denoised outputs at several noise levels for different training conditions. Figures are not reproduced here.

Paper vs code

Official repository kulinshah98/memorization_noisy_data at code_commit. The repository has no LICENSE file; file headers carry NVIDIA's CC BY-NC-SA 4.0 notice and the README says the code is adapted from the EDM repository. Cite it, do not copy from it.

1. Routing of samples. The paper builds each batch from S ∪ S_tn and applies the ambient loss to noisy members and the denoising loss to clean members alg. 1. The code gives every image a fixed nature noise, then picks the branch per example from a sampled noise level: ambient when sigma is at or above the nature-noise level, EDM denoising otherwise loss.py loss.py training_loop.py. The same underlying image can enter either branch.
2. Time sampling. The paper samples t uniformly on each interval alg. 1. The code samples sigma log-normally with P_mean -1.2 and P_std 1.2 loss.py loss.py. The paper does not state the time distribution.
3. Weighting. The paper's ambient loss eq. 6 is an unweighted squared error. The code multiplies the squared error of the converted target by the EDM weight (sigma^2 + sigma_data^2) / (sigma sigma_data)^2 with sigma_data 0.5 loss.py loss.py.
4. Parameterisation. The paper uses the variance-preserving process eq. 1. The code adds additive noise with std sigma loss.py and converts the denoiser output with from_x0_pred_to_xnature_pred_ve_to_ve [code: torch_utils/ambient_utils/init.py#L74-L76]. The two forms are not shown to be equivalent in the paper.
5. Nature noise. The paper creates one noisy version per image alg. 1, line 1. The code draws Gaussian noise of std sigma per image, seeded by image index when nature_noise_fixed is set dataset_utils.py dataset_utils.py. It also computes a "dust" corruption mask that is returned but never applied to the noise dataset_utils.py dataset_utils.py. The per-image and per-pixel corruption probabilities are set to 1.0 train.py, so the mask has no effect in the shipped configuration.
6. Training length. The paper says 30,000 iterations at batch 256 for CIFAR-10 and FFHQ app. D. The script sets --duration=30 train.sh train.sh; train.py reads duration in Mimg and sets total_kimg = duration * 1000 train.py train.py. That is 30 million images, about 117,188 iterations at batch 256 (derived by us from these values).
7. Learning rate default. The paper uses 1e-4 app. D. The script passes --lr=1e-4 train.sh. The train.py default is 1e-3 train.py.
8. Noise level in the shipped script. The script uses sigma 0.5 train.sh and the loss-type switch selects the ambient branch train.sh. The paper's sweep is [0.4, 4] fig. 1. The train.py default is 0.0 train.py.
9. Augmentation, dropout, EMA, warm-up. The paper does not state these. The script uses augmentation 0.15 and dropout 0.05 train.sh. The defaults are 0.12 and 0.13 train.py. EMA half-life defaults to 0.5 Mimg train.py. Snapshots store the EMA weights training_loop.py. Learning-rate warm-up is linear over 10,000 kimg training_loop.py training_loop.py. train.py does not override the warm-up, so with the script's 30,000 kimg it covers one third of training (derived).
10. Text-conditional experiments. The paper uses Stable Diffusion v2 with LAION-10k, batch 16, 100,000 iterations app. D. The repository's train.py implements only the EDM-style unconditional and class-conditional loss types train.py. diffusers_utils.py holds text-encoding and sampling helpers diffusers_utils.py diffusers_utils.py, but no text-conditional training entry point was found. The text-conditional recipe is therefore not reproducible from this repository. The paper says these experiments build on Somepalli et al.'s implementation app. D.
11. Memorization metric. The paper says DINOv2 nearest-neighbour similarity §5.1. The code loads dinov2_vitg14 find_dataset_neighbors.py, normalises features by default find_dataset_neighbors.py find_dataset_neighbors.py, takes the inner product with all training features find_dataset_neighbors.py, and keeps the maximum find_dataset_neighbors.py. The script caps generated samples at 5,000 memorization_metrics.sh. The paper does not state the sample count for the metric.
12. Script syntax. The trailing comments in scripts/train.sh come after a backslash followed by spaces, so the backslash does not continue the line train.sh. Under the usual shell rules the command would split. This was not run here.

Repository state at the cloned commit: unconditional training, nearest-neighbour search, FID and generation scripts are present. The text-conditional pipeline is absent.

Pitfalls

• The schedule of t_n is unspecified. The paper says reasonable values form an interval, with too low reverting to DDPM and too high bringing memorization back §5.1. The per-setting t_n is not reported, so reproduce by matching the DDPM FID first, then sweep.
• Noise is drawn once. Redrawing the noisy copies each epoch changes the method. The paper states that the one-time construction is what reduces memorization §3.
• Sample routing differs from the paper. The code routes by sampled noise level rather than dataset membership (Paper vs code, item 1). With the script's sigma 0.5 and the default time distribution, the ambient branch takes about a third of samples (derived from the log-normal defaults). The paper's uniform batch from S ∪ S_tn would give about half noisy samples if |S_tn| = |S| (derived from alg. 1). A faithful re-implementation of the paper needs the routing from alg. 1, not from the code.
• Iteration count. The paper's 30,000 iterations and the script's 30 Mimg differ by a factor of about four (Paper vs code, item 6). Check which one your target setting means before comparing numbers.
• Abstract claim. The text says "the same or better FID". In Table 1 every FID cell is better except FFHQ at 3,000 images, where ours is 6.46 against 6.42 for DDPM tab. 1 §5.1.
• Benefit shrinks with dataset size. The authors say the gain in both FID and memorization shrinks as data grows §5.1.
• Metric dependence. The memorization numbers depend on the DINOv2 variant and on thresholds tuned separately per dataset, and neither the variant (in the paper) nor the threshold values are stated §5.1. Compare only within this paper's tables.
• Text conditioning. Without a text-memorization fix the method does not reduce memorization much, because the model relies on the text §5.2 tab. 3. The reported gains come only in combination with [S23] or [W24].
• Privacy. The method has no privacy guarantee §6.
• Theory assumptions. The optimal-score lemma assumes the score is learned perfectly §4. The Gaussian-mixture constants used in the appendix are chosen for the argument, not measured app. A.4.
• Licensing of the code. Non-commercial terms apply (see code_license). Do not copy snippets into commercial work.

Lineage

• Builds on: ambient score matching and corrupted-data training from Daras et al., cited in the paper as §1 and §2.3; the EDM codebase, on which the code is based app. D; the subpopulation and heavy-tail model of Feldman (2020) §4.1.
• Compared against: the text-memorization mitigations of Somepalli et al. (S23) and Wen et al. (W24) §5.2; linear-corruption ("Masking") and noise-only ("Noise") training §5.1 tab. 2.
• Concurrent work cited: Wu et al. and Chen et al. (2025), theoretical treatments of memorization trade-offs in diffusion models §1.