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: The paper introduces a methodological framework for generative modeling through diffusion processes without time-reversal arguments. By utilizing diffusion bridges the authors consider processes that start and end in pre-determined points $x_0$ and $a_{\tau}$. Then by extending this to mixtures of diffusion bridges they show how to transport a distribution $\Pi_0 = P_D$ to $\Pi_{\tau} = P_Z$ without time-reversal of the diffusion process. For specific classes of SDEs by drawing connections to the paradigm introduced by Song et al 2021 and provide a unified view on the drift adjustment in forward and backward SDEs. This allows for defining a time and state-dependent class probability function (given by conditional expectation) providing further insight into the inner-workings of DTRT and DBMT. The authors then introduce two other training objectives $L_{\text{FD}},L_{\text{CE}}$ and discuss their favorable properties. Finally, methods for moving beyond fully factorial noise models are presented by extending the process and the noise term from functions of time to be functions of both time and space. This allows for incorporating "priors" with better spatial dependency models more suitable for specific datasets.