Non-Denoising Forward-Time Diffusions

1 INTRODUCTION . Denoising diffusion probabilistic modeling ( DDPM ) ( Sohl-Dickstein et al. , 2015 ; Ho et al. , 2020 ; Song et al. , 2021 ) is a recent generative modeling paradigm exhibiting strong empirical performance . Consider a dataset of N samples D = { x ( n ) } Nn=1 with empirical distribution PD . The unifying key steps underlying DDPM approaches are : ( i ) the definition of a stochastic process with initial distribution PD , whose forward-time ( noising ) dynamics progressively transform PD toward a simple data-independent distribution PZ ; ( ii ) the derivation of the backward-time ( denoising / sampling ) dynamics transforming PZ toward PD ; ( iii ) the approximation of the backward-time transitions by means of a neural network . Following the training step ( iii ) , a sample whose distribution approximates PD is drawn by ( iv ) simulating from the approximated backward-time transitions starting with a sample from PZ . Both discrete-time ( Ho et al. , 2020 ) and continuous-time ( Song et al. , 2021 ) formulations of DDPM have been pursued . This work focuses on the latter case , to which we refer as diffusion time-reversal transport ( DTRT ) . As in DTRT , dynamics are specified through diffusion processes , i.e . solutions to stochastic differential equations ( SDE ) with associated drift f ( · ) and diffusion g ( · ) coefficients . A number of approximations are involved in the aforementioned steps . Firstly , as the dynamics are defined on a finite time interval , a dependency from PD is retained through the noising process . Hence , starting with a sample from the data-independent distribution PZ in ( iv ) introduces an approximation . Secondly , while the backward-time dynamics of ( ii ) are directly available for diffusions , they are approximated by means of a neural network in ( iii ) . Thirdly , sampling in ( iv ) is achieved through a discretization on a time-grid , which introduces a discretization error . De Bortoli et al . ( 2021 , Theorem 1 ) links these approximations to the total variation distance between the distribution of the generated samples from ( iv ) and PD . In our first methodological contribution we develop a procedure for constructing diffusion processes targeting PD without relying on time-reversal arguments . The proposed transport ( coupling ) between PZ and PD is achieved by : ( 1 ) specifying a diffusion process X on [ 0 , τ ] starting from a generic x0 ; ( 2 ) conditioning X on hitting a generic xτ at time τ , thus obtaining a diffusion bridge ; ( 3 ) taking a bivariate mixture Π0 , τ of diffusion bridges over ( x0 , xτ ) with marginals Π0 = PZ and Πτ = PD , obtaining a mixture process M ; ( 4 ) matching the marginal distribution of M over [ 0 , τ ] with a diffusion process , resulting in a diffusion with initial distribution PZ and terminal distribution PD . The realized diffusion bridge mixture transport ( DBMT ) between PZ and PD is exact by construction . We thus sidestep the approximation common to all DDPM approaches due to the dependency from PD retained through the noising process . Moreover , the DBMT can be realized for almost arbitrary PZ , f ( · ) and g ( · ) . This increased flexibility is a departure from the DTRT where f ( · ) and g ( · ) need to be chosen to obtain convergence toward a simple distribution PZ . Similarly to the DTRT , achieving the DBMT requires the computation of a drift adjustment term which depends onD . For a SDE class of interest , we develop a unified and interpretable representation of DTRT and DBMT drift adjustments as simple transformations of conditional expectations over D. This novel result provides insights on the target mapping that we aim to approximate and on the quality of approximation achieved by the trained score models of Song et al . ( 2021 ) . Having defined for the DBMT a Fisher divergence objective similarly to Song et al . ( 2021 ) , we leverage on this unified representation to define two additional training objectives featuring appealing properties . In our last methodological contribution we extend the class of SDEs that can be realistically employed in computer vision applications . Specifically , computational considerations have so far restricted the transitions of the stochastic processes employed in DDPM to be fully factorials . We view images at a given resolution as taking values over a 2D lattice which discretizes the continuous coordinate system [ 0 , 1 ] 2 representing heights and widths . Diffusion processes are viewed as spatio-temporal processes with spatial support [ 0 , 1 ] 2 . Doing so , it is possible to leverage on scalable simulation and inference techniques from spatial statistics and consider more realistic diffusion transitions . This paper is structured as follows . In Section 2 we review the DTRT of Song et al . ( 2021 ) and in Section 3 we introduce the DBMT . In order to implement the DTRT and the DBMT it is necessary to specify the underlying SDE , i.e . the coefficients f ( · ) and g ( · ) . We study a class of interest in Section 4 . The unified view of drift adjustments is introduced in Section 5 . Section 6 develops the training objectives and Section 7 reviews the obtained results and finalizes the DBMT construction . In Section 8 we establish the connection with spatio-temporal processes . We conclude in Section 9 . Appendices A to D contain the theoretical framework , assumptions , proofs , and additional material . Notation and conventions : we use uppercase notation for probability distributions ( measures , laws ) and lowercase notation for densities ; each probability distribution , and corresponding density , is uniquely identified by its associated letter not by its arguments ( which are muted ) ; for example P ( dx ) is a distribution , p ( x ) is its corresponding density ; random elements are always uppercase ( an exception is made for times , always lowercase for typographical reasons ) ; if P is the distribution of a stochastic process , we use subscript notation to refer to its finite dimensional distributions ( densities with p ) , conditional or not , for some collection of times ; for example pt′|t denotes a transition density , which is understood to be a function of four arguments pt′|t ( y|x ) = f ( t , t′ , x , y ) ; δx is the delta distribution at x and ⊗ is used for product distributions ; we refer directly to a given SDE instead of referring to the diffusion process satisfying such SDE when no ambiguity arises ; we use [ a ] i and [ A ] i , j for vector and matrix indexing , A > for matrix transposition . 2 DIFFUSION TIME-REVERSAL TRANSPORT . The starting point of Song et al . ( 2021 ) is a diffusion process Y satisfying a generic D-dimensional time-inhomogenous SDE with initial distribution Y0 ∼ PD dYr = f ( Yr , r ) dr + g ( Yr , r ) dWr , ( 1 ) over noising time r ∈ [ 0 , τ ] . Thorough this paper we denote with Q the law of the diffusion solving ( 1 ) and with q the corresponding densities . Thus , let qr′|r ( y|x ) , 0 ≤ r < r′ ≤ τ , be the transition density of ( 1 ) , and let qr ( y ) , 0 < r ≤ τ , be the marginal density of ( 1 ) . As Y0 ∼ PD , we have qr ( y ) = 1 N N∑ n=1 qr|0 ( y|x ( n ) ) . ( 2 ) The dynamics of ( 1 ) over the reversed , i.e . sampling , time t = τ − r , t ∈ [ 0 , τ ] , are given by ( Anderson , 1982 ; Haussmann & Pardoux , 1986 ; Millet et al. , 1989 ) dXt = [ −f ( Xt , r ) +∇ ·G ( Xt , r ) +G ( Xt , r ) ∇Xt ln qr ( Xt ) ] dt+ g ( Xt , r ) dWt , ( 3 ) where r = τ − t is the remaining sampling time , G ( x , r ) = g ( x , r ) g ( x , r ) > and the D-dimensional vector ∇ ·G ( x , r ) is defined by [ ∇ ·G ( x , r ) ] i = ∑D j=1∇xj [ G ( x , r ) ] i , j . That is the processes Xt and Yr = Yτ−t have the same distribution . Approximating the terminal distribution Qτ of ( 1 ) , i.e . the initial distribution of ( 3 ) , with PZ , X0 is sampled from PZ and ( 3 ) is discretized and integrated over t to produce a sample Xτ approximately distributed as PD . The computation of the multiplicative drift adjustment ∇y ln qr ( y ) entering ( 3 ) , i.e . the score of the marginal density ( 2 ) , requires in principle O ( N ) operations . Let sφ ( y , r ) be a neural network for which we would like sφ ( y , r ) ≈ ∇y ln qr ( y ) . It remains to find a suitable training objective for which unbiased gradients with respect to φ can be obtained at O ( 1 ) cost with respect to the dataset size N . As qr ( y ) has a mixture representation , the identity of Vincent ( 2011 ) for Fisher divergences provides us with the desired objective for a fixed r ∈ ( 0 , τ ] LFD , DTRT ( φ , r ) = E Yr∼Qr [ ∥∥∇Yr ln qr ( Yr ) − sφ ( Yr , r ) ∥∥2 ] = E ( Y0 , Yr ) ∼Q0 , r [ ∥∥∇Yr ln qr|0 ( Yr|Y0 ) − sφ ( Yr , r ) ∥∥2 ] . ( 4 ) The key point is that an unbiased , O ( 1 ) with respect to N , mini-batch Monte Carlo ( MC ) estimator for the expectation ( 4 ) can be trivially obtained by sampling a batch Y0 ∼ PD , Yr ∼ Qr|0 ( dyr|Y0 ) , and evaluating the average loss over the batch . In order to achieve a global approximation over the whole time interval ( 0 , τ ] , Song et al . ( 2021 ) proposes uniform sampling of time r LFD , DTRT ( φ ) = E r∼U ( 0 , τ ] , ( Y0 , Yr ) ∼Q0 , r [ Rr ∥∥∇Yr ln qr|0 ( Yr|Y0 ) − sφ ( Yr , r ) ∥∥2 ] , ( 5 ) where Rr = E [ ‖∇Yr ln qr|0 ( Yr|Y0 ) ‖2 ] −1 is a regularization term . A MC estimator for ( 5 ) is constructed by augmenting the MC estimator for ( 4 ) with the additional sampling step r ∼ U ( 0 , τ ] . 3 DIFFUSION BRIDGE MIXTURE TRANSPORT . Our starting point is a generic D-dimensional time-inhomogenous SDE which , in contrast to Song et al . ( 2021 ) , is directly defined on the sampling time t ∈ [ 0 , τ ] dXt = f ( Xt , t ) dt+ g ( Xt , t ) dWt . ( 6 ) We reserve P · |0 ( · |x0 ) to denote the law of the diffusion solving ( 6 ) for a given starting value x0 and p · |0 ( · |x0 ) to denote the corresponding densities .

Summary: In this paper the authors introduce a new model for diffusion-based generative modelling.  Instead of relying on time-reversed processes as in existing works, they propose an approach based on the mixing of diffusion bridges. They show that a mixing of diffusion bridges remains a diffusion with closed-form drift and volatility. Similarly to existing score-matching based generative modelling works the drift in this new formulation can be expressed as a conditional expectation and therefore can be seen as the minimizer of some loss function, allowing the use of neural-network based approximations. In addition, the authors propose some unification of the SDE classes used in [1] and to use non-identity volatility matrices.  [1] Song, Sohl-Dickstein, Kingma, Kumar, Ermon, Poole -- Score-based Generative Modeling through Stochastic Differential Equations