Speed-accuracy relations for diffusion models: Wisdom from nonequilibrium thermodynamics and optimal transport
Authors: Kotaro Ikeda, Tomoya Uda, Daisuke Okanohara, Sosuke Ito
Organizations: Department of Mathematical Engineering and Information Physics, School of Engineering, The University of Tokyo, 7-3-1 Hongo, Bunkyo-ku, Tokyo 113-0033, Japan · Department of Earth and Planetary Physics, School of Science, The University of Tokyo, 7-3-1 Hongo, Bunkyo-ku, Tokyo 113-0033, Japan · Preferred Networks Inc., 1-6-1, Otemachi, Chiyoda-ku, Tokyo 100-0004, Japan · Universal Biology Institute, Graduate School of Science, The University of Tokyo, 7-3-1 Hongo, Bunkyo-ku, Tokyo 113-0033, Japan
We discuss a connection between a generative model, called the diffusion model, and nonequilibrium thermodynamics for the Fokker-Planck equation, called stochastic thermodynamics. Using techniques from stochastic thermodynamics, we derive the speed-accuracy relations for diffusion models, which are inequalities that relate the accuracy of data generation to the entropy production rate. This relation can be interpreted as the speed of the diffusion dynamics in the absence of the non-conservative force. From a stochastic thermodynamic perspective, our results provide quantitative insight into how best to generate data in diffusion models. The optimal learning protocol is introduced by the geodesic of space of the 2-Wasserstein distance in optimal transport theory. We numerically illustrate the validity of the speed-accuracy relations for diffusion models with different noise schedules and different data. We numerically discuss our results for optimal and suboptimal learning protocols. We also demonstrate the applicability of our results to data generation from the real-world image datasets.
Figures & tables
Figure 1: The organization of this paper. In Sec. II , we explain the basic concepts of diffusion models and discuss their methods, including score-based generative modeling and flow-based generative modeling. In Sec. III , we discuss stochastic thermodynamics and optimal transport theory, and introduce the entropy production rate. We also introduce several thermodynamic trade-off relations, such as thermodynamic uncertainty relations and speed limits, which provide the minimum entropy production in a finite time via optimal transport. In Sec. IV , We derive our main results in analogy to thermodynamic trade-off relations. Our main results reveal the essential role of the entropy production rate in diffusion models. Noise schedules based on optimal transport will lead to the most robust generation in diffusion models, and existing noise schedules such as the conditional optimal transport schedule and the cosine schedule are regarded as suboptimal noise schedules.
Figure 2: Illustration of diffusion models in the case of image data. The forward process, denoted by Pt(x) , is a stochastic process which adds noise to the training data sampled from q(x) . By inverting this forward process in time, we prepare the reverse process Pt~†(x) . The generated data sampled from p(x) is constructed by an estimated reverse process P~t~†(x) that mimics the reverse process. The image in the figure is taken from the dataset [ 91 ] .
Figure 3: Illustration of the speed-accuracy relation for diffusion models [Eq. ( 79 )] in the absence of the non-conservative force
Figure 4: Summary of the main results. We show a hierarchy between the inequalities, and a generalization for incomplete estimation discussed in Sec. V.1 .
Figure 5: Examples of the forward process and the estimated reverse process according to three noise schedules: the cosine schedule [Eq. ( 32 )], the cond-OT schedule [Eq. ( 28 )], and the optimal transport which gives the dynamics along the geodesic in the space of the 2-Wasserstein distance. These figures show the time evolution of samples from a 2 -dimensional Swiss-roll dataset ( nd=2 ) in the flow-based generative modeling. In all noise schedules, we used the same sample for the initial distribution of the estimated reverse process P~τ(x) . Compared with the cosine schedule and the cond-OT schedule, the noise in the generated figure is relatively small when we consider the optimal transport. This result is consistent with the optimality based on the speed-accuracy relations for diffusion models.
Figure 6: Numerical calculations of the speed-accuracy relations for diffusion models [Eqs. ( 79 ), ( 80 ) and ( 83 )]. (a) The time evolution of the probability density functions in the forward process and the estimated reverse process with the cosine schedule, the cond-OT schedule and the optimal transport (OT). Here t denotes time in the forward direction, and we show the probability density functions at t=0 , t=τ/3 , t=2τ/3 and t=τ . The final state of the forward process is Pτ(x)=N(x∣0,1) . We used a Gaussian distribution with a different mean P~0†(x)=N(x∣1,1) as the initial condition of the estimated reverse process. (b-e) We show the situation where the initial condition of the estimated reverse process P~0†(x) is (b) a Gaussian distribution with a different mean, (c) a Gaussian distribution with a different variance, (d) a Gaussian mixture distribution, (e) a uniform distribution. The plot at the top left of each section shows the comparison between P~0†(x) and Pτ(x)=N(x∣0,1) , where N(x∣0,1) is shown in gray. In the top right-hand section, we show the bar graph which displays the quantities η=(ΔW1)2/(τD0) , cl=∫0τdt[vloss(t)]2 , c2=∫0τdt[v2(t)]2 , and this figure shows the validity of the inequality η≤c1≤c2 [Eq. ( 83 )]. In the bottom row we show the time evolution of [v2(t)]2 and [vloss(t)]2 , and this figure shows the validity of the inequality [v2(t)]2≥[vloss(t)]2 [Eq. ( 80 )]. Since [v2(t)]2 and c2 depend only on the forward process, these values are consistent within different initial conditions in (b)-(e). See also Table 3 in Appendix G for more detailed explanations, including the probability density function P~0†(x) used in (b)-(e).
Initial conditions
Noise schedules
Values of η
Gaussian distribution with a different mean
Cosine Cond-OT OT
9.188×10−28.781×10−28.537×10−2
Gaussian distribution with a different variance
Cosine Cond-OT OT
1.078×10−11.047×10−11.022×10−1
Gaussian mixture distribution
Cosine Cond-OT OT
2.764×10−22.674×10−22.611×10−2
Uniform distribution
Cosine Cond-OT OT
2.142×10−22.068×10−22.033×10−2
Table 1: The values of the response function η=(ΔW1)2/(τD0) in the numerical calculations [in Fig. 6 ]. Since we computed 3×108 samples, we display the value with four significant digits.
Figure 7: Numerical experiments of the speed-accuracy relations for diffusion models [Eqs. ( 77 ), ( 80 ) and ( 81 )] in two-dimensional Swiss-roll situation. (a)-(c) The time evolution of the probability density function in the forward process and the estimated reverse process with the noise schedules using g=0,3,5 in Eq. ( 99 ). The value g stands for the degree of non-conservativity and g=0 denotes a situation with conservative external force ft(x) . (d) The time evolution of the probability density function with optimal transport noise schedule. In (a)-(d) t denotes time in the forward direction, and we show the probability density functions at t=0 , t=τ/3 , t=2τ/3 and t=τ . Each axis of the color map denotes the axis of the data space Rnd . Throughout this experiment, we used a Gaussian distribution with different variance P~0†(x)=N(0,(0.9)2I ) as the initial condition of the estimated reverse process. (e) Temporal evolution of TtS˙ttot and [vloss(t)]2 proving the validity of [vloss(t)]2≤TtS˙ttot [Eq. ( 80 )]. (f) A bar graph showing the quantities η=(ΔW1)2/(τD0) , cl=∫0τdt[vloss(t)]2,wD=∫0τdtTtS˙ttot . This figure shows the relationship η≤cl≤wD [Eq. ( 81 )]. For detailed explanations see Appendix G .
Initial conditions
Noise schedules
Values of η
Gaussian distribution with a different variance
g=0g=3g=5 OT
9.301×10−23.430×10−15.615×10−17.221×10−6
Table 2: The values of the response function η=(ΔW1)2/(τD0) in the numerical calculations [Fig. 7 ]. Since we computed 1×108 samples, we display the value with four significant digits.
Figure 8: The application of diffusion models using the real-world image datasets and the validity of Eq. ( 79 ) on the latent space. (a) The forward process and the estimated reverse process for the downscaled version of the CelebA HQ dataset [ 119 ] . (b) The forward process and the estimated reverse process for the bedrooms subset of the LSUN dataset [ 120 ] . (c) The estimated values of η=(ΔW1)2/(τD0),c2=∫0τdt[v2(t)]2 computed on the latent space. The error bars shows the standard deviation among 5 calculations using 3×104 samples for each calculations. nd=4096 denotes the dimension of the latent space. The inequality c2≥η [Eq. ( 79 )] on the latent space is still valid for realistic applications of diffusion models using the real-world image datasets.
Table 3: The initial conditions used in the numerical calculations [Fig. 6 ]. Here N(x∣μ,σ2) denotes the 1 -dimensional Gaussian distribution with mean μ and variance σ2 .
Figure 9: The time evolution of the probability density functions in the forward and estimated reverse processes, for each condition used in the numerical calculations [Fig. 6 ].
We discuss three simple connections between generative diffusion and optimal transport. First, the joint law of the data and of their Ornstein-Uhlenbeck noising is, at every time, the solution of an entropic optimal-transport problem, at a temperature set by the noise level. Backward diffusion can consequently be read as an annealing which gradually decreases the entropic regularization. Second, we study the speciation transition--how the modes of the data distribution emerge during generation--and compare how it happens in diffusion and flow matching with respect to optimal transport. Using Gaussian mixtures as high-dimensional example, we show that the dynamical instability through which diffusion and flow matching separate the modes becomes, for optimal transport, an exponentially thin boundary layer at the beginning of generation, which is rounded by entropic regularization. Third, considering a generic forward dynamics with a general confining potential, we show that the backward process driven by the score also satisfies a drift-based variational principle whose minimal cost is the entropy difference between initial and final distributions.
Giulio Biroli, Marc Mézard
Laboratoire de Physique de l’Ecole Normale Supérieure, ENS, Université PSL, CNRS, Sorbonne Université, Université Paris-Diderot, Sorbonne Paris Cité, Paris, France · Department of Computing Sciences, Bocconi University, Milano, Italy
We develop a principled framework for analyzing and designing noise schedules in diffusion models. We show that one can recast this design problem as an optimal control problem, whose state is the Fisher information of the diffusion process, which evolves according to an ODE, and whose control input is the noise schedule. The objective of the optimal control problem is a functional involving the Fisher information, which is shown to be an upper bound on the Kullback-Leibler sampling error. By solving this optimal control problem, we obtain sufficient conditions on noise schedules under which O~(d/n) sampling error is achievable, where d is the data dimension and n is the number of discretization steps. While existing theoretical works also prove that O~(d/n) sampling error bounds are achievable, these results hold for specific noise schedules that differ from those used in practice. Under additional assumptions on the data distribution, we derive a family of noise schedules that recovers the functional forms of schedules widely used in practice. We show that these schedules achieve the O~(d/n) error bound when their parameters are chosen according to our analysis. Our theory offers a principled approach to noise schedule design, yielding explicit functional forms and guidance for selecting their parameters.
Seo Taek Kong, Weina Wang, R. Srikant
ECE & CSL, University of Illinois Urbana-Champaign · Computer Science Department, Carnegie Mellon University · ECE, CSL & NCSA, University of Illinois Urbana-Champaign
SDE-based generative models, including diffusion models and the Schrödinger bridge, have found broad applications in signal processing tasks such as speech enhancement, image restoration, and time-series generation. This note presents a modeling framework for such models within the context of stochastic thermodynamics. The main results of this note are trajectory-level definitions of work, heat, and entropy production, along with a generalized Jarzynski identity and a second-law-like inequality. The proposed framework extends the original Jarzynski setup to accommodate time-dependent bath temperature and nonconservative driving forces. This thermodynamic perspective may deepen our understanding of diffusion models and the Schrödinger bridge from a nonequilibrium statistical mechanics viewpoint.
Yaowen Zhang
Department of Physics, The Hong Kong University of Science and Technology