Loss Crossovers and Phase Transitions Under Ambient Training in Diffusion Models

Abstract

The distinct behaviours of memorisation and generalisation in generative models has been noted empirically, and linked to sharp phase transitions as model capacity increases. Further, recent work has shown that for the standard diffusion loss objective crossovers between memorising and generalising solutions predict the phase transition location. We investigate how loss crossovers behave in the ambient regime introduced by Shah et al. (2025). Our main result in fig. 2, shows that crossover locations inherit the same bell-shaped dependence on tn demonstrated empirically. All the experiments in our project are run on a rented NVIDIA 4080 Ti with a compute budget of less than £10.

1 Introduction

In the study of generative models memorisation refers to behaviour where samples copy images from the training data. In contrast, when a model produces novel samples from the true distribution it is said to have generalised. Recent study has noted a phase shift between memorising and generalising behaviour Buchanan et al. (2025); Wang et al. (2024); Zhang et al. (2023), where once model-capacity passes a threshold in relation to the dataset the shift from memorisation to generalisation occurs suddenly.

2 Diffusion Modelling and the Theoretical Optimum Behaviour

Given a target probability distribution π⋆ and a dataset (x0(i))i=1N of samples from π⋆, diffusion models attempt to learn to sample from the true distribution. The core mechanism of diffusion models is a forward and a backward process. The forward process gradually removes signal from a sample x0 until it transformed to xt∼𝒩⁢(0,𝐈); whereas, the reverse process does the opposite. Despite originally being presented as a discrete time evolution in Ho et al. (2020) we present the continuous parametrisation following Song et al. (2020) with t∈[0, 1]. The dynamics of the forward process are governed by f,g:[0,1]→ℝ a drift and diffusion coefficient, which control the process through the following SDE

d⁢xt=f⁢(t)⁢xt⁢d⁢t+g⁢(t)⁢d⁢wt. (1)

The Gaussian nature of the forward process allows for a tractable parametrisation of the conditional marginal distributions at time t;

pt⁢(xt∣x0)=𝒩⁢(xt;st⁢x0,st2⁢σt2⁢𝐈),st:=exp⁡(∫0tf⁢(ξ)⁢𝑑ξ),σt:=∫0tg2⁢(ξ)s2⁢(ξ)⁢𝑑ξ.

The reverse process evolves a sample from 𝒩⁢(0,𝐈) to π⋆ by gradually removing the noise added by the forward process. The true reverse process is typically intractable as it depends on the score function of the unconditional marginal distribution pt through the following ODE

d⁢xt=(f⁢(t)⁢xt−12⁢g2⁢(t)⁢∇log⁡pt⁢(xt))⁢d⁢t. (2)

To sample from the true distribution we approximate the score function using a neural network, there are a variety of ways of doing this and it is an active research area in of itself Karras et al. (2022), we present the denoising autoencoder: where a model 𝐱θ⁢(⋅,t) is trained to approximate 𝔼(x0∣xt) by optimising for the objective

ℒden⁢(𝐱θ,λw)=𝔼t[λw⁢(t)⁢𝔼x0∼π⋆𝔼xt∣x0[‖𝐱θ⁢(xt,t)−x‖2]]. (3)

During sampling we leverage Tweedie’s formula to relate 𝔼(x0∣xt) to the score function via st⁢𝔼(x0∣xt)=xt+st2⁢∇log⁡pt⁢(xt). Note the introduction of an arbitrary weighting function λw. Typically, 𝐱θ is parametrised by UNet Ho et al. (2020). To train 𝐱θ on the objective 3 we use the Monte-Carlo approximation of the true distribution

ℒNden⁢(𝐱θ,λw)=𝔼t[λw⁢(t)⁢ℒN,tden⁢(𝐱θ)],ℒN,tden⁢(𝐱θ):=1N⁢∑i=1N𝔼xt∣x0(i)[‖𝐱θ⁢(xt,t)−x0(i)‖2] (4)

with ℒNden⁢(𝐱θ,λw)≈ℒden⁢(𝐱θ,λw). However, the authors of Raphan and Simoncelli (2011); Karras et al. (2022) note that the non-parametric optimum to 4 simply memorises the training dataset by learning a kernel weighted mixture. We denote this model

𝐱mem⁢(xt,t):=∑i=1Nx0(i)⁢softmax⁢(wi⁢(xt,t)),where⁢wi⁢(xt,t)=−12⁢σt2⁢‖xt−st⁢x0(i)‖2. (5)

This result suggests that a diffusion model optimised for 4 should not generalise and instead memorise, as 𝐱mem points the reverse process toward a weighted sum of images in the dataset. This concludes our introduction to diffusion and the theoretical optimal behaviour.

3 The Phase Transition from Memorisation to Generalization

It has been proven in model examples Wang et al. (2024), as well as demonstrated empirically in Gu et al. (2023); Zhang et al. (2023) that the leading determinant of whether a model memorises or generalises is the relationship between model capacity and the dataset. We draw our attention to the paper Buchanan et al. (2025) which provides a framework which allows this phenomenon to be explored analytically. Let us now give a brief exposition of their findings.

We set the true distribution π⋆ as an equally weighted mixture of K isotropic Gaussians in ℝd 1K⁢∑i=1K𝒩⁢(μi⋆,σ⋆2⁢𝐈), which the dataset (x0(i))i=1N is sampled from. We introduce two denoisers: firstly, the parametric optimum 𝐱gen which is parametrised with (μi⋆)i=1N and σ⋆2, and a partially memorising denoiser 𝐱pmem, M which has memorised M out of N of the dataset points similar to 5. The advantage of this environment is that all denoisers are from same parametric family, see Lemma 2.1 of Buchanan et al. (2025) for details. The authors hypothesis that there exists a weight functions such that a phase shift happens if an only if ℒN,td⁢e⁢n⁢(𝐱gen)>ℒN,td⁢e⁢n⁢(𝐱pmem, M); they verify this claim empirically by finding a weight function which accurately predicts the phase transition observed in practice. A pivotal insight is the following decomposition via the orthogonality principle

ℒN,tden⁢(𝐱θ)=1N⁢∑i=1N𝔼xt∣x0(i)[‖𝐱θ⁢(xt,t)−𝐱mem⁢(xt,t)‖2]+ℒN,tden⁢(𝐱mem),

which can be used to show that

𝔼(x0(i))i=1N[ℒN,tden⁢(𝐱pmem, M)−ℒN,tden⁢(𝐱gen)] =1N⁢∑i=1N𝔼xt∣x0(i)[‖𝐱pmem, M⁢(xt,t)−𝐱mem⁢(xt,t)‖2] (6)
−1N⁢∑i=1N𝔼xt∣x0(i)[‖𝐱gen⁢(xt,t)−𝐱mem⁢(xt,t)‖2].

The components on the RHS are approximated individually, leading to the main result of the paper theorem 1 which shows that the excess loss is monotonically decreasing in M.

Theorem 1.

Taken from Buchanan et al. (2025) With ψt:=αt2/σt2 and given that N=polly⁢(d), mini≠j⁡‖μi⋆−μj⋆‖2=𝒪⁢(d), maxi⁡‖μi⋆‖2=𝒪⁢(d), σ⋆2=𝒪⁢(1) it follows that

𝔼(x0(i))i=1N[ℒN⁢(𝐱pmem, M,λw)−ℒN⁢(𝐱gen,λw)]≈𝔼t[λw⁢(t)⁢(C⁢(1−MN)⁢d⁢σ⋆2−d⁢σ⋆2ψt⁢σ⋆2+1)],

and a result, there exists a M⋆∈[0,N] such that the LHS equals zero.

Under the authors hypothesis M⋆ corresponds to the location of phase transition. It can be calculated by setting the RHS of theorem 1 equal to zero and rearranging for M;

M⋆≈N⁢[1−𝔼t[λw⁢(t)/(ψt⁢σ⋆2+1)]C⁢𝔼t[λw⁢(t)]].

We note that despite M⋆ depending on λw it does not influence whether a crossover exists. We now shift our focus to an augmented training regime which was introduced to reduce memorisation.

4 An Algorithm to Reduce Memorisation

In Li and Chen (2024) the authors demonstrate that output diversity is primarily controlled by the high-noise regime. Inspired by this, Shah et al. (2025) leverage Ambient score matching which is introduced in Daras et al. (2024) to create an algorithm which encourages generalisation in the high noise regime. Ambient score matching allows training without access to images from the true distribution by training on the objective

ℒtnAmb⁢(𝐱θ,λw)=𝔼t>tn[λw⁢(t)⁢𝔼xtn𝔼xt∣xtn[‖σt2−σtn2σt2⁢1−σtn2⁢𝐱θ⁢(xtst,t)+σtn2σt2⁢xtst−xtnstn‖2]]. (7)

In practice the expectation over xtn is replaced with the dataset Stn of noisy labels. The algorithm proposed by Shah et al. (2025) further augments training by simultaneously optimising for ℒN,tden and ℒtnAmb. This is achieved by fixing the time parameter tn∈[0, 1], and then training using ambient score matching when t>tn and denoising score matching when t≤tn. See algorithm 1 for an explicit presentation using the variance preserving drift and diffusion coefficients αt,σt which we use in all of our experiments.

Algorithm 1 Algorithm for reducing memorisation (taken from Shah et al. (2025)).
  1. Input: untrained network \(\mathbf{x}_{\theta}\), set of samples \(S\), noise level \(t_{\mathrm{n}}\), batch size \(B\), diffusion time \(T\)
  2. \(S_{t_{\mathrm{n}}}\leftarrow\bigl\{\sqrt{1-\sigma_{t_{\mathrm{n}}}^{2}}\,x_{0}^{(i)}+\sigma_{t_{\mathrm{n}}}\epsilon^{(i)}\mid x_{0}^{(i)}\in S,\;\epsilon^{(i)}\sim\mathcal{N}(0,\mathbf{I})\bigr\}\) Noise the training set at level \(t_{\mathrm{n}}\).
  3. while not converged do
  4. Form a batch \(\mathcal{B}\) of size \(B\) uniformly sampled from \(S\cup S_{t_{\mathrm{n}}}\)
  5. \(\mathrm{loss}\leftarrow 0\) Initialize loss.
  6. for each sample \(x\in\mathcal{B}\) do
  7. \(\epsilon\sim\mathcal{N}(0,\mathbf{I})\) Sample noise.
  8. if \(x\in S_{t_{\mathrm{n}}}\) then
  9. \(x_{t_{\mathrm{n}}}\leftarrow x,\; t\sim\mathcal{U}(t_{\mathrm{n}},1)\) Noisy sample.
  10. \(x_{t}\leftarrow\sqrt{\frac{1-\sigma_{t}^{2}}{1-\sigma_{t_{\mathrm{n}}}^{2}}}\,x_{t_{\mathrm{n}}}+\sqrt{\frac{\sigma_{t}^{2}-\sigma_{t_{\mathrm{n}}}^{2}}{1-\sigma_{t_{\mathrm{n}}}^{2}}}\,\epsilon\) Prepare for ambient score matching.
  11. \(\mathrm{loss}\leftarrow\mathrm{loss}+\Bigl\|\frac{\sigma_{t}^{2}-\sigma_{t_{\mathrm{n}}}^{2}}{\sigma_{t}^{2}\sqrt{1-\sigma_{t_{\mathrm{n}}}^{2}}}\mathbf{x}_{\theta}(x_{t},t)+\frac{\sigma_{t_{\mathrm{n}}}^{2}}{\sigma_{t}^{2}}\sqrt{\frac{1-\sigma_{t}^{2}}{1-\sigma_{t_{\mathrm{n}}}^{2}}}\,x_{t}-x_{t_{\mathrm{n}}}\Bigr\|^{2}\)
  12. else
  13. \(x_{0}\leftarrow x,\; t\sim\mathcal{U}(0,t_{\mathrm{n}})\) Clean sample.
  14. \(x_{t}\leftarrow\sqrt{1-\sigma_{t}^{2}}\,x_{0}+\sigma_{t}\epsilon\)
  15. \(\mathrm{loss}\leftarrow\mathrm{loss}+\|\mathbf{x}_{\theta}(x_{t},t)-x_{0}\|^{2}\) Regular denoising score matching.
  16. end if
  17. end for
  18. \(\mathrm{loss}\leftarrow\frac{\mathrm{loss}}{B},\;\theta\leftarrow\theta-\eta\nabla_{\theta}\) Average loss and backpropagate.
  19. end while

The authors argue that the images in the noised dataset are harder to memorise as they have been corrupted by Gaussian noise which is incompressible. Further, distinct samples at the zero noise level merge together during the forward process further reducing the incentive for memorisation. Studying the algorithm 1, it is clear that as tn approaches one we recover the typical diffusion training objective, and as tn approaches zero the benefit of the image corruption is reduced. Thus, we would expect any measure of incentive-to-memorise plotted against tn to have an approximate bell shape – we refer to this as the bell-shaped property.

A noteable drawback of algorithm 1 is that is reduces image quality. It is shown in Figure 1 of Shah et al. (2025) that there is trade-off between decreasing memorisation and image quality. However, algorithm 1 currently produces state of the art results on this trade off. We now begin the contributions that as far as we know are unexplored in the literature. We begin by formulating the loss objective of algorithm 1 as follows

ℒtAlg1⁢(𝐱θ)=𝟏t<tn⁢[ℒtDen⁢(𝐱θ)]+𝟏t>tn⁢[ℒtnAmb⁢(𝐱θ)], (8)

and thus by the law of total expectation

ℒtnAlg1⁢(𝐱θ,λw)=tn⋅𝔼t<tn[λw⁢(t)⁢ℒtDen⁢(𝐱θ)]+(1−tn)⋅𝔼t>tn[λw⁢(t)⁢ℒtnAmb⁢(𝐱θ)]. (9)

We notice that the loss is split into an affine combination of two loss functions and that in the mixture of Gaussian setup of the previous section we can apply theorem 1 to the first component. This motivates a study of the excess loss between 𝐱pmem, M and 𝐱gen to explain the phase transition in the ambient regime.

Proposition 1.

Our own analogue to theorem 1. Under the assumptions of theorem 1 we have that

where the inputs to 𝐱mem and 𝐱pmem, M are removed for brevity. This result follows from applying the orthogonality principle to 8, and then applying the approximation in theorem 1.

The crossover dynamics in proposition 1 has gained an additional component compared to theorem 1, thus analytical study of crossover is reduced to creating an approximation to 10 for both 𝐱gen and 𝐱pmem, M. In fig. 2 we provide evidence that 𝒞new⁢(𝐱pmem, M)−𝒞new⁢(𝐱mem) is approximately linear in M/N.

𝒞new⁢(𝐱):=𝔼(x0(i))i=1N𝔼t>tn𝔼xt∣xtn⁡‖σt2−σtn2σt2⁢1−σtn2⁢𝐱−𝐱mem‖2 (10)

5 The Properties of Loss Crossovers in Ambient Training

Our main finding is that the crossover location as a function of tn exhibits the same bell-shaped structure that we observe empirically; suggesting that crossovers remain predictive in the regime governed by algorithm 1. We begin with an empirical study of the affect that algorithm 1 has on memorisation as tn varies. Figure 1 shows that there is an inverted bell-curve when memorisation ratio is plotted against tn as when tn is close to 0 or 1 the trained model exhibits large amounts of memorisation, but for intermediate tn memorisation is reduced. This suggests that the location of the phase transition from memorisation to generalisation with respect to model capacity would be delayed for intermediate values of tn.

Refer to caption
Figure 1: Memorisation ratio of a diffusion model trained on 64 images from the CIFAR-10 dataset as a function of the noise threshold σtn which is parameterised by tn. Memorisation is reduced for intermediate values of σtn, exhibiting an inverted bell-shaped dependence consistent with the behaviour observed under algorithm 1. Memorisation ratio is defined in Gu et al. (2023) The model is a ddpmpp U-Net Ho et al. (2020) with 128 channels and channel multipliers [2, 2, 2], trained using the Adam optimiser with learning rate 12.5⋅10−6 and batch size 64.

Now, motivated by the relationship between crossovers and phase transitions in Buchanan et al. (2025) we investigate the properties of the crossovers in the regime of algorithm 1. In fig. 2 we demonstrate that the crossovers exhibit the same bell-shaped dependency on tn that we observed empirically. We take this as strong evidence that crossovers could be used to accurately predict the affect that tn has on phase transitions. An exciting prospect for future work is to investigate whether there exists a weight function such that the crossovers predict the location of phase transitions as well as the properties.

(a)Refer to caption

(b)Refer to caption

(c)Refer to caption

Figure 2: The crossover location varies non-monotonically with tn exhibiting a bell-shaped structure consistent with empirical findings. (left) We plot crossover point against tn (centre) We plot the losses of 𝐱mem,𝐱pmem, M,𝐱gen against tn. (right) We plot the difference in loss between 𝐱pmem, M and 𝐱gen as t and tn vary. All plots are in the environment N=200,M=100,K=12,σ⋆2=1, taken from Buchanan et al. (2025), with (μi⋆)i=1N sampled from 𝒩⁢(0, 30⁢𝐈).

We finish our study under the assumption that loss crossovers predict the locations of phase transition. A notable property of the original crossover theorem 1 is that M⋆ is linear in N. In fig. 3(left) we plot the crossover point as a function of N and tn in the ambient regime and notice limited variation in N. This suggests that similar to theorem 1 a crossover occurs when M/N is equal to a fixed constant. Together, all of these results suggest that the loss crossovers retain predictive power over the phase shift even in modified training regimes.

(a)Refer to caption

(b)Refer to caption

Figure 3: The crossover ratio is largely invariant with respect to N. (left) The crossover point is plotted with respect to N and tn Invariance is not observed for the variation of K; seen in the (right) plot. We calculate crossovers for N∈{50, 100, 150, 300, 600, 1000} and K∈{10, 20, 30, 40, 50, 60, 70, 80, 90, 100} with all other parameters as in 2.

References

  • S. Buchanan, D. Pai, Y. Ma, and V. De Bortoli (2025) On the edge of memorization in diffusion models. arXiv preprint arXiv:2508.17689. Cited by: §1, §3, §3, Figure 2, §5, Theorem 1.
  • G. Daras, A. G. Dimakis, and C. Daskalakis (2024) Consistent diffusion meets tweedie: training exact ambient diffusion models with noisy data. arXiv preprint arXiv:2404.10177. Cited by: §4.
  • X. Gu, C. Du, T. Pang, C. Li, M. Lin, and Y. Wang (2023) On memorization in diffusion models. arXiv preprint arXiv:2310.02664. Cited by: §3, Figure 1.
  • J. Ho, A. Jain, and P. Abbeel (2020) Denoising diffusion probabilistic models. Advances in neural information processing systems 33, pp. 6840–6851. Cited by: §2, Figure 1.
  • T. Karras, M. Aittala, T. Aila, and S. Laine (2022) Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems 35, pp. 26565–26577. Cited by: §2.
  • M. Li and S. Chen (2024) Critical windows: non-asymptotic theory for feature emergence in diffusion models. arXiv preprint arXiv:2403.01633. Cited by: §4.
  • M. Raphan and E. P. Simoncelli (2011) Least squares estimation without priors or supervision. Neural computation 23 (2), pp. 374–420. Cited by: §2.
  • K. Shah, A. Kalavasis, A. R. Klivans, and G. Daras (2025) Does generation require memorization? creative diffusion models using ambient diffusion. arXiv preprint arXiv:2502.21278. Cited by: Loss Crossovers and Phase Transitions Under Ambient Training in Diffusion Models, §4, §4, Algorithm 1.
  • Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2020) Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. Cited by: §2.
  • P. Wang, H. Zhang, Z. Zhang, S. Chen, Y. Ma, and Q. Qu (2024) Diffusion models learn low-dimensional distributions via subspace clustering. arXiv preprint arXiv:2409.02426. Cited by: §1, §3.
  • H. Zhang, J. Zhou, Y. Lu, M. Guo, P. Wang, L. Shen, and Q. Qu (2023) The emergence of reproducibility and generalizability in diffusion models. arXiv preprint arXiv:2310.05264. Cited by: §1, §3.