Memorization and generalization in generative diffusion under the manifold hypothesis

Abstract We study the memorization and generalization capabilities of a diffusion model (DM) in the case of structured data defined on a latent manifold. We specifically consider a set of P data points in N dimensions lying on a latent subspace of dimension <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:mi>D</mml:mi> <mml:mo>=</mml:mo> <mml:msub> <mml:mi>α</mml:mi> <mml:mi>D</mml:mi> </mml:msub> <mml:mi>N</mml:mi> </mml:mrow> </mml:math> , according to the hidden manifold model. Our analysis considers a reverse process given by the empirical score function as a proxy of the true one, and then precisely characterizes the process in the high-dimensional limit in which <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:mi>P</mml:mi> <mml:mo>=</mml:mo> <mml:mi>exp</mml:mi> <mml:mo>⁡</mml:mo> <mml:mo stretchy="false">(</mml:mo> <mml:mi>α</mml:mi> <mml:mi>N</mml:mi> <mml:mo stretchy="false">)</mml:mo> </mml:mrow> </mml:math> and N is large, by exploiting a connection with the random energy model (REM). We provide evidence for the existence of an onset time, <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">o</mml:mi> </mml:mrow> </mml:msub> </mml:mrow> </mml:math> , when traps appear in the time-varying potential, although they do not affect typical trajectories. The size of the basins of attraction of such traps is computed at any time. Moreover, we derive the collapse time, <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">c</mml:mi> </mml:mrow> </mml:msub> <mml:mo>&lt;</mml:mo> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">o</mml:mi> </mml:mrow> </mml:msub> </mml:mrow> </mml:math> , at which trajectories fall in the basin of one of the training points, implying memorization. An explicit formula for <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">c</mml:mi> </mml:mrow> </mml:msub> </mml:mrow> </mml:math> is given as a function of P and the ratio α D , proving that the curse of dimensionality issue does not hold for highly structured data, i.e. <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>α</mml:mi> <mml:mi>D</mml:mi> </mml:msub> <mml:mo>≪</mml:mo> <mml:mn>1</mml:mn> </mml:mrow> </mml:math> , regardless of the non-linearity of the manifold surface. We also prove that collapse coincides with the condensation transition in the REM. Finally, the degree of generalization of DMs is formulated in terms of the Kullback–Leibler divergence between the exact distribution and the one obtained at time t of the reverse process. We show the existence of an additional time <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">g</mml:mi> </mml:mrow> </mml:msub> <mml:mo>&lt;</mml:mo> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">c</mml:mi> </mml:mrow> </mml:msub> <mml:mo>&lt;</mml:mo> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">o</mml:mi> </mml:mrow> </mml:msub> </mml:mrow> </mml:math> such that the distance between the reverse distribution and the ground-truth is minimal. Counter-intuitively, the best generalization performance is found within the memorization phase of the model. We conclude that the generalization performance of DMs benefit from highly structured data since <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">g</mml:mi> </mml:mrow> </mml:msub> </mml:mrow> </mml:math> approaches zero faster than <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>t</mml:mi> <mml:mrow> <mml:mi mathvariant="normal">c</mml:mi> </mml:mrow> </mml:msub> </mml:mrow> </mml:math> when <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" overflow="scroll"> <mml:mrow> <mml:msub> <mml:mi>α</mml:mi> <mml:mi>D</mml:mi> </mml:msub> <mml:mo stretchy="false">→</mml:mo> <mml:mn>0</mml:mn> </mml:mrow> </mml:math> .