Generative Modeling with Phase Stochastic Bridge
Authors: Tianrong Chen
Jiatao Gu
Laurent Dinh
Evangelos Theodorou
Joshua M. Susskind
Shuangfei Zhai
Venue: ICLR 2024 oral (2024)
Keywords: Generative Modeling
Stochastic Optimal Control
Diffusion Model
Primary area: generative models

Abstract: Diffusion models (DMs) represent state-of-the-art generative models for continuous inputs. DMs work by constructing a Stochastic Differential Equation (SDE) in the input space (ie, position space), and using a neural network to reverse it. In this work, we introduce a novel generative modeling framework grounded in \textbf{phase space dynamics}, where a phase space is defined as {an augmented space encompassing both position and velocity.} Leveraging insights from Stochastic Optimal Control, we construct a path measure in the phase space that enables efficient sampling. {In contrast to DMs, our framework demonstrates the capability to generate realistic data points at an early stage of dynamics propagation.} This early prediction sets the stage for efficient data generation by leveraging additional velocity information along the trajectory. On standard image generation benchmarks, our model yields favorable performance over baselines in the regime of small Number of Function Evaluations (NFEs). Furthermore, our approach rivals the performance of diffusion models equipped with efficient sampling techniques, underscoring its potential as a new tool generative modeling.

Published as a conference paper at ICLR 2024
GENERATIVE MODELING WITH PHASE STOCHASTIC BRIDGES
Tianrong Chen', Jiatao Gu², Laurent Dinh Evangelos A. Theodorou'
Josh Susskind Shuangfei Zhai?
Georgia Tech, Apple (tianrong.chen, evangelos, theodoroo gatech.edu, hinb @apple.coom
ABSTRACT
We introduce novel generative modeling framework grounded in phase space dynamics, taking inspiration from the principles underlying Critically damped Langevin Dynamics and Bridge Matching. Leveraging insights from Stochas- tic Optimal Control, we construct a more favorable path measure in the phase space that is highly advantageous for efficient sampling. A distinctive feature our approach the early-stage data prediction capability within the context of propagating generative Ordinary Differential Equations or Stochastic Differ- ential Equations. This early prediction, enabled by the model's unique structural characteristics, sets the stage for more efficient data generation. leveraging addi- tional velocity information along the trajectory. This innovation has spurred the exploration of novel avenue for mitigating sampling complexity by quickly con- verging to realistic data samples. Our model yields comparable results in image generation and notably outperforms baseline methods, particularly when faced with limited Number of Function Evaluations. Furthermore, our approach rivals the performance of diffusion models equipped with efficient sampling techniques, underscoring its potential in the realm of generative modeling.
INTRODUCTION
Diffusion Models (DMs;Song et al. (2020a); Ho et al. (2020)) constitute an instrumental technique in generative modeling, which formulate a particular Stochastic Differential Equation (SDE) linking the data distribution with tractable prior distribution. Initially, a DM diffuses data towards the prior distribution via a predetermined linear SDE. In order reverse the process, neural network is used to approximate the score function which is analytically available. Subsequently, the approx- imated score is utilized to conduct time reversal (Anderson, 1982; Haussmann & Pardoux, 1986) of this diffusion process, ultimately generating data. Recently, the Critical-damped Langevin Dynam- ics (ockooo et al. (2021)) extends the SDE framework of DM into phase space (whereas DMs operate in the position space) by introducing an auxiliary velocity variable, which is defined by tractable Gaussian distributions at the initial and terminal time steps. This augmentation induces trajectory in position space exhibiting enhanced smoothness, as stochasticity is solely introduced into the velocity space. The distinctive structure of CLD is shown to enhance the empirical per- formance and sample efficiency. However, despite the success of CLD, inefficient sampling still persists due to unnecessary curvature of the dynamics (Fig.1) 1) as it has to converge to equilibrium for sampling from the tractable prior.
The remarkable accomplishments of DM have also catalyzed recent advancements in generative modeling, leading to the development of Bridge Matching (BM:(Peluchetti, 2021; Liu et al., 2022; 2023)) and Flow Matching (FM:models(Lipmaan et al., 2022)). These models leverage dynamic transport maps underpinned by the utilization of SDES or ODEs. Unlike DM, Bridge and Flow Matching relaxes the reliance on a forward diffusion process with an asymptotic convergence to prior distribution over an infinite time horizon. Moreover, they exhibit heightened degree of ver- satility, enabling the construction of transport maps between two arbitrary distributions by drawing
work done while Tianrong Chen is an intern Apple

---

Published as a conference paper at ICLR 2024
position
velocity
CLD
58
18
08
16
AM(ours)
03
32
14
6.6
6.0
1.0
60
0. 0.4 0.5 08 18
Figure 1: The pixel-wise trajectories comparison with CLD(Dockhorn et al., 2021). Left figures cor- respond to the trajectories over time w.rt random sampled 16 pixels, for position and velocity. Our model is able to learn straighter trajectories which is beneficial for reducing sampling complexity.
upon insights from domains such as optimal transport (Pooladian et al., 2023), normalizing flow (Tong et al., 2023b), and optimal control (Liu et al., 2023).
In this paper, we focus on enhancing the sample efficiency of velocity based generative modeling (eg, CLD) by utilizing the Stochastic Optimal Control (SOC) theory. Specifically, we leverage the outcomes of stochastic bridge within the context of linear momentum systems (Chen & Georgiou, 2015) to construct a path measure bridging the data and prior distribution. The resulting path ex- hibits more straight position and velocity trajectory compared to CLD (fig. 1), making it more amenable to efficient sampling. Within the broader landscape of dynamic generative modeling (ie, ODE/SDE based generative models), data point can often be represented as linear combinations of scaled intermediate data of dynamics and Gaussian noise. In work, re-establish this property, enabling the estimation of target data points by leveraging both state and velocity information. In the case of DM and FM, the estimation of target data is exclusively reliant on position information, whereas our method incorporates the additional dimension of velocity data, enhancing the preci- sion and comprehensiveness our estimations. It is also worth noting that our model exhibits the capacity to generate high fidelity images at early time steps (fig.2). In addition, we propose sam- pling technique which demonstrates competitive results with small Number of Function Evaluations (NFEs), eg, 5 to 10. 1able.1 demonstrates the design differences among aforementioned models. In summary, our paper presents the following contributions:
1. We propose Acceleration Generative Modeling (AGM) which is built the SOC theory, abling the favorable trajectories for efficient sampling over 2nd-order momentum dynamics gen- erative modeling such CLD.
2. As a result of AGM structural characteristics, it becomes possible to estimate realistic data point at an early time point, concept we refer to as sampling-hop. This approach not only yields significant reduction in sampling complexity but also offers novel perspective accelerating the sampling in generative modeling by leveraging additional information from the dynamics.
3. We achieve competitive results compared DM approaches equipped with specifically designed fast sampling techniques on image datasets, particularly in small NFE settings.
2 PRELIMINARY
Notation: Let X, E Rdd and v, E 'd denote the d-dimensional position and velocity variable of particle m, [x,, vz]T E 2dd at time We denote the discretized time series as 0 S to < ...ty < IN < 1. The Wiener Process is denoted as w,. The identity matrix is denoted as I E Baxdd We define E, as the covariance matrix of x, and v, at time step L.
2.1 DYNAMICAL GENERATIVE MODELING
The generative modeling approaches rooted in dynamical systems, including ODE and SDE, have garnered significant attention. Here, we present three noteworthy dynamical generative models: Diffusion Model (DM), Flow Matching (FM) and Bridge Matching (BM).
2

---

Published as a conference paper at ICLR 2024
Table 1: Comparison between models in terms of boundary distributions Po and Pt- Our AGM generalizes DM beyond Gaussian priors to phase space, similar to CLD. However, unlike CLD, AGM does not need to converge the Gaussian at equilibrium which causes curved trajectory(see 1), instead, velocity distribution will be the convolution of data distribution with Gaussian.
Models
DM/FM
CLD
AGM()
pol-) (z)
Pdttta (x) x N N Ia)
N(, Do x 12d)
Pi()) I(O,
N Ia) N La) Pdata(z) Pdata (z) * (O Σ1 @ I2a)
Diffusion Model: In the framework of DM, given Xo drawn from data distribution Pdata the model proceeds to construct a SDE,
dx;- = f((x; )dt + g(t)dw, Xo Pasta (x)
whose terminal distributions at approach an approximate Gaussian, i.e. x1 (O, Ia). This accomplishment is realized through the careful selection of the diffusion coefficient Ir and the base drift fx)). It is noteworthy that the time-reversal (Anderson, 1982) of (1) results in another SDE:
dx,= = [f((x.) - o. p(x,, t)] di + g())dw, X1 ~ I,)
(2)
where p(:,1) t) is the marginal density of (1) at time / and Vx log is known as the score function. SDE (2) can be regarded as the time-reversal of (1) in such a manner that the path-wise measure is almost surely equivalent to the one induced by (I). As a consequence, these two SDEs share identical marginal over time. In practice, it is feasible to analytically sample X, given & and Xo- Additionally, we can leverage a neural network to learn the score function by regressing scaled Stein Score Ex,, |l/ [xx, t; -) Vx & log p(x, t t/xo) 12 for the purpose of propagating (2). This learned score can then be integrated into the solution of the aforementioned SDE(2) to simulate the generation of data that adheres to the target data distribution from the prior distribution. Meanwhile, (2) also corresponds to ODE which shares the same path-wise measure:
1 9 Vx logp(x,, t) dt, X1 ~ NoO
dx,=
((()))
(3)
which motivates the popular sampler introduced in (Zhang & Chen, 2022; Zhang et al., 2022; Bao et al., 2022) to solve the ODE (2) efficiently.
Bridge Matching and Flow Matching: An alternative approach to exploring the time-reversal of forward noising process involves the concept of building bridges' between two distinct distributions Po(.) and Pı (-). This method entails the learning of a mimicking diffusion process, commonly
referred to as bridge matching, as elucidated in previous works (Peluchetti, 2021; Shi et al., 2022).
Here we consider the SDE in the form of:
(4) which is pinned down at an initial and terminal point To, Iy which are independently samples from predefined Po and p1. This is commonly known as the reciprocal projection of zu and æ1 in the literature (Shi et al., 2023; Peluchetti, 2023; Liu al., 2022; Léonard et al., 2014). The construction of such SDE is accomplished by meticulous design of v,. A widely adopted choice for v, is V, := (x1 x1 - 1), which induces the well-known Brownian Bridge (Liu et al., 2023; Somnath et al., 2023). Similar to the approach in DM and owing to the linear structure of the dynamics, one can efficiently estimate this drift by employing a neural network parameterized by weights for regression on: [Ex,,llvv e (x,, t; 0) v, (xe, t)ll2 given X and t. As extensively discussed in previous studies (Liu et al., 2023; Shi et al., 2022), this bridge matching framework takes on the characteristics of FM (Lipman et al., 2022) when the diffusion coefficient 9: tends to zero.
dx, (x, t)dt g,dw, s.t. (10, ~ Io, (x0. x1) == Po X P1
Remark 1. The practice of constraining stochastic process to specific initial and terminal condi- tions is a esallestabl setup in SOC. For a gentle introduction of it's connection with Brownian Bridge, Schrödinger Bridge please see Appendix. C. From this perspective, Bridge, as elaborated in Appendix. D. 1 for comprehensive elucidation. It is imperative to note that the SOC framework will serve as the fundamental basis upon which we will develop our algorithm.
can derive Brownian
3

---

Published as a conference paper at ICLR 2024
3 ACCELERATION GENERATIVE MODEL
We apply SOC to characterize the twisted trajectory of momentum dynamics induced by CLD(Dockhom et al., 2021). becomes evident that the mechanisms encompassing flow matching, diffusion modeling, and Bridge matching collectively facilitate the construction of an estimated tar- get data point, denoted as x1, by utilizing the intermediate state of the dynamics, Xx. Our additional objective is to expedite the estimation of plausible x by incorporating additional dynamics-related information, such as velocity, thereby curtailing the requisite time integration.
In this section, we introduce the proposed method, termed as the Acceleration Generative Model (AGM), rooted SOC theory. Building upon (Chen & Georgiou, 2015), we extend the frame- work by incorporating time-varying diffusion coefficient and accommodating arbitrary boundary conditions, ultimately arriving at an analytical solution suited for the generative modeling. We demonstrate its efficacy in rectifying the trajectory of CLD, concurrently showcasing its aptitude for accurately estimating the target data at an early timestep 4, thereby enabling expeditious sampling.
As suggested by BM approach, there is necessity to formulate trajectory that bridges the two data points sampled from Po and #1 respectively. Desirably, the intermediate trajectory should exhibit optimal characteristics that facilitate smoothness and linearity. This is essential for the ease of simulating the dynamics system to obtain the solution. In our endeavor to tackle this challenge and enhance the estimation of the data point X1 by incorporating velocity components, we encapsulate the problem within SOC framework, specifically formulated in the phase space which reads:
Definition 2 (Stochastic Bridge problem of linear momentum system (Chen & Georgiou, 2015)).
min Ae
[
|| 1[att + (m1 m2) m1)_Rmmm m2)
dxe dv.
Vt ar(Xe, v,
dt
0
dw
0
gt
dm,
f(m.t)
g+
m,
X,
Ir
R
r 0
Las
31 ~ Pdata.
(5)
In this context, the matrix R is recognized as the terminal cost matrix, serving to assess the proximity between the propagated m, and the ground truth mi at the terminal time = 1. As the parameter approaches positive infinity, the trajectory converges toward the state xi, prompting transition to constrained dynamics wherein the system becomes constrained by two predetermined boundaries, namely my and mi. This configuration aligns scamlessly with the principles of constructing feasible bridge, as advocated by the tenets of BM. is worth noting that this interpolation approach essentially represents natural extension (Chen & Georgiou, 2015) of the well-established concept of the Brownian Bridge (Revuz & Yor, 2013), which has been employed in trajectory inference (Somnath al., 2023; Tong al., 2023a) and image inpainting tasks (Liu et al., 2023) and its connection with Diffusion has been discussed in Liu et al. (2023). Indeed, it is evident that the target velocity lacks a precise definition within this problem, allowing for flexibility in the design space for our approach. To address this, we opt for the linear interpolation of the intermediate point and the target point, represented as V1 (x1 - x1) - t), as the chosen terminal velocity, which also is the optimal control in the original space (see Appendix..D.), This choice is made due to its ability to construct a trajectory characterized by straightness. Conceptually, the acceleration a, continually guides the dynamics towards the linear interpolation of the two data points, serving to mitigate the impact of introduced stochasticity. In contrast to previous bridge matching frameworks, the velocity's 's boundary condition in our approach varies over time since it depends the state x, and t. The velocity variable serves solely as an auxiliary component aimed at straightening the trajectories. Regarding this SOC problem formulation, the solution is,
Proposition 3 (Phase Space Brownian Bridge). When +oc, The solution w.t optimization problem 5 is,
a" amm 1) = g
X1 - X2
v.) 1-1
where: : Pn =
-4
(6)
02 -
Proof. Please see Appendix.D.2.
4

---

Published as a conference paper at ICLR 2024
EDM
x
S
X
X V
t
AGM(ours)
Sq
X1
0
1 0.3
0.7
Figure 2: Data estimation comparison with EDM (Karras al., 2022). When the network is en- dowed with supplementary velocity, AGM gains the capacity to estimate the target data point during the early stages of the trajectory. One can use estimated image x₁ at t, < fty as generated results and allocated more NFE between time [0, f.] which results to smaller discretization error.
Remark 4. Pn1 denotes the second diagonal component the matrix Pt, a solution derived from the Lyapunov equation (see Lemma.9) serving as an implicit representation of the optimality of the control. This value is dependent upon the uncontrolled dynamics, where a, set to the zero vector in (5), and will vary accordingly when uncontrolled dynamics change.
3.1 TRAINING
By plugging the optimal control (6) back to the dynamics (5), can obtain the desired SDE. As been suggested by (Song et al., 2020b; Dockhorn et al., 2021), such SDE has a corresponding probablistic ODE which shares the same marginal over time in which the drift term will have additional score term V. log p(m, t). Here we summarize the force term for SDE and ODE as:
dxe dv.
=
Vz F,
dt
0
0 h4
dw:
s.t.
mg =
Xo
~ NNuo, En),
Vo
Bridge Matching SDE F; == F (m, t) = at (m,
h(t) = g(t),
Probablistic ODE F. := (m, () = a, (m, 5 V. log (og,(m,
I((1) : 0.
Henceforth, we refer to the dynamics associated with the Bridge Matching SDE as AGM-SDE, and its corresponding ODE counterpart as AGM-ODE. Meanwhile, the linearity of the system implies the intermediate state m, and the close form solution of score term are analytically available. In par- ticular, the mean μ,, and covariance matrix E, of the intermediate marginal [m, x1) N E,)
Σx Σ
Dav
of such system can be analytically computed with Σ,
Σ
2 Ia. and
pro-
vided we have the boundary conditions p4 and Eg in place, as outlined in Särkkä & Solin (2019).
Please Appendix.D.3 for detail. In order to sample from such multi-variant Gaussian, one need to decompose the covariance matrix by Cholesky decomposition, and m, reparamertized as:
LE L En
m, = pr + Lye 14
V. log = 1,61
L E1
cg] 61
where L. f
~ -N I2a) and ly v
5

---

Published as a conference paper at ICLR 2024
Parameterizatton: The Force term can be represented as composite of the data point and Gaussian noise. Specifically,
a" am 1) = 4x1 (1 1)² - P11 [1--
L
1
60
14 61
We express the force term as F S z₁. Here, 2 assumes the role of regulating the output of the network si, ensuring that the variance of the network output is normalized to unity. For the detailed formulation of the normalizer 3r, please refer Appendix.D.8 In a manner similar the BM approach, one can formulate the objective function for regressing the force term as follows:
min Ex 1xx: Em (m. (xx) A(t) [ t; 0) FF 1)
(10)
0
Where A(t) is known as the reweight of the objective function across the time horizon. We defer the derivation of & and the presentation of Lt, X(t) and a in Appendix.D.
3.2 SAMPLING FROM AGM
Once the paramterized force term F/ is trained, we are ready to simulate the dynamics to generate the samples by plugging back to the dynamics (7). One can use any type of SDE or ODE sampler to propagate the learnt system. Here we list our choice of sampler for AGM-SDE and AGM-ODE.
Stochastic Sampler: To simulate the SDE, prior works are majorly relying on Euler- Maruyama(EMM) (Kloeden et al., 1992) and related methods. We adopt the Symmetric Splitting Sampler(SSS) from Dockhorn et al. (2021) in our AGM-SDE. This selection based on the com- pelling performance it offers when dealing with momentum systems.
Deterministic Sampler: is imperative to acknowledge that this system is inherently underactuated because the force term is exclusively injected into the velocity component, while velocity serves as the driving factor for the position-a variable of primary interest in generative modeling context. More specifically, time step t,, the impact of force does not immediately manifest in the position but rather takes effect at a subsequent time step, denoted as f,41 after discretizing time horizon. At time (o. it becomes undesirable to propagate the state Xn using an initially uncontrolled velocity over an extended time interval on- The presence of this delay phenomenon can also exert an influence when the time interval 6, is large, thereby impeding our ability reduce the NFE during sampling- We propose the adoption of an Exponential Integrator (EI) approach, as claborated in Zhang & Chen (2022). Empirical evidence suggests that this method aligns well with our model. We provide an illustrative example of how the AGM-ODE, in conjunction with the EI technique, can be employed to inject the learnt network into both velocity and position channels simultaneously:
X1,+1 = (11++1,4)
X
(l++1 - -T) 2r Ma5(TTdr si (m,, i s:41-)))
Vi++1
Vt
Shittz, z, (., s (m...ssti-))
==0
(II)
Where Mis(T) (T)
II kØj
T T- ti
and ((t,s)
1 t t-8
زيتا - -
0
1
In 11, ((s, !) denotes the transition matrix for our system, while Mr, represents the w-order multistep coefficient (Hochbruck & Ostermann, 2010). For comprehensive derivation of these terms, please refer to Appendix.D.9. It is worth noting that the mapping of Sy into both the po- sition and velocity channels significantly emulates the errors introduced by discretization delays. Sampling-hop: In the context of CLD (Dockhorn et al., 2021), their focus is on estimating the score function w.r.t. velocity, which essentially corresponds to estimating scaled E, in our notation. However, relying solely on the aforementioned information is not sufficient for estimating the data point x₁. Additional knowledge regarding eg is also required in order to perform such estimation. In our case, the training objective implicitly includes both 6n and é (see eq.9), hence one can manage to recover X1 by Proposition.5. Remarkably, our observations have unveiled that when the network is equipped with additional velocity information, it acquires the capability to estimate the target data point during the early stages of the trajectory, as illustrated f2g.2. This estimation can be seamlessly integrated into AGM-SDE and AGM-ODE and we name it sampling-hop. Specifically,

---

Published as a conference paper at ICLR 2024
Proposition 5 (Sampling-Hop) Given the state, velocity and trained force term Fú at time step in sampling phase, The estimated data point X2 can be represented as
(1 FF + vv)) + Xt.
F. Pu + B.v.)
SDE X1
x
(12)
P1
11t Pu (ast + (B(f))
for AGM-SDE and AGM-ODE dynamics respectively, and Be = - 2P11 ot
L
Proof. See Appendix. D.10
This property empowers us to allocate the NFE budget selectively within the time interval / E [0, 4,]. where & tv. effectively reducing the discretization error while maintaining the sampling quality. This insight paves the way for efficient low NFE sampling strategies later. Here we summarized the training and sampling procedure of our method in Algorithm.1 and Algorithm.2 respectively.
Algorithm Training
Algorithm Sampling
1: Input: data distribution Pdata()
1: Input: trained F(.,;; 0)), discretized time step [to.... :f,]. Choose the sampler from [SSS(SDE), EI(ODE)). Choose prior mean and covariance μ4o Eg
2: while not converge do
3: t~ 4(O0, 1)), X1 ~ (X1)
4: Compute mean and covariance lr and E,
(Appendix.De D.3)
2: Sample mg ~ (m(: Mo; En).
5:
Sample m, = |4; (.E.(eq.8)
for n = 0 to do 4: estimate F (m. (n)
6: Compute target F; (cq.7) using optimal acceleration (eq.9)
m -Sampler ( F (m)
7:
Compute loss IE [A|LE? - F, 10).
6: reconstruct x using Proposition.5.
8: Take gradient descent with respect to
end for 8: Return 81
F t,
9: end while
4 EXPERIMENTAL RESULTS
k=-0,9
k=-0.2
k=-0.5
data scale
(1::0)0) O) using modified NCSN+ model as provided in Karras al. (2022).
Architectures and Hyperparameters: We parameterize (..
2.0
15-
We employ six input channels, accounting for both position and velocity variables, as opposed to the standard three channels used in the CIFAR-10 (Krizhevsky al., 2009), AFHQv2 (Choi et al., 2020) and ImageNet (Deng et al., 2009) which leads to a negligi- ble increase of network parameters. For the purpose of comparison with CLD in the toy dataset, we adopt the same R ar- chitecture utilized in CLD. Throughout all of experiments, we maintain a monotonically decreasing diffusion coefficient, given by g(t) 3(1 t). For the detailed experimental setup, please refer further to Appendix.E.
1.0
05
0.0
-3
Figure 3: The standard devia- ton of the terminal marginal for uncontrolled dynamics. We empirically selected the hyperparameter A = -0.2. This choice induces termi- nal marginal distribution with that covers the data range with uncontrolled dynamics.
Evaluation: To assess the performance and the sampling speed of various algorithms, we employ the Fréchet Inception Distance score (FID:Heusel al. (2017)) and the Number Function Evaluations (NFE) as our metrics. For FID evaluation, we utilize reference statistics of all datasets obtained from EDM (Karras et al., 2022) and use 50k generated samples to evaluate. Additionally, evaluate the FID of CLD and EDM using the same reference statis- ties to ensure consistency in our comparisons. For all other reported values, we directly source them from respective referenced papers.
re-
Selection of Eg: The choice of initial covariance Eg directly influ- ences the path measure of the trajectory. In our case, we set Eg := [ [] with hyperparameter k. We observe that trajectories tend exhibit pronounced curvature under specific conditions: when

---

Published as a conference paper at ICLR 2024
NFE=S
NFE-i
NFE=15
NFE-2i
NPE-5
NFE
NPE-15
NPE-21
EDM (Karras al.)
AGM(ours)
FID=16.g8
FID=4.71
FID=3-91
FID=3.68
FID-100
FID-13.6
FID FID-4-73
FID FID-2.99
Figure 4: Comparison with EDM (Karras al., 2022) on AFHQv2 dataset. AGM-ODE exhibits superior generative performance when NFE is exceedingly low, owing to its unique dynamics archi- tecture that incorporates velocity when predicting the estimated data point.
Lncondiiiionll Dynumiks
Atroke YO
Stroke Conditional Dynamics
Inpainting Conditional Dynamics
We showcase that AGM generate conditional results from unconditional model by injecting the conditional information into the velocity Vo. thus leading to new initial velocity vgond
Figure 5:
the is positive, the absolute value of the position is large. This behavior is particularly notice- able when dealing with images, where the data scale ranges from 1. We aim for favorable uncontrolled dynamics, as this can potentially lead to better-controlle dynamics. Our strategy is to design in such way that the marginal distribution uncontrolled dynamics f.y 1 effectively covers the range of image data values meanwhile & keeps negative. We can express the marginal of uncontrolled dynamies by leveraging the transition matrix (11, 0), which gives us x1 : Xg + Vo- Figure 3 illustrates the standard deviation of x, for various values of k. Based on our empirical observations, we choose & = -0.2 for all experiments, as effectively covers the data range. The subsequent controlled dynamics (eq.7) will be constructed based on such desired uncontrolled dy- namics as established.
Table 2: FIDJ, Comparison with CLD(Dockhorn et al., 2021) using same SSS Sam- pler on CIFAR-10.
Stochastic Sampling: In experiments, emphasize the advan- tages of using the AGM-SDE compared with CLD. Firstly, we show that model exhibits superior performance when NFE is signifi- cantly lower than that of CLD, particularly in toy dataset scenarios. For evaluation, utilized the multi-modal Mixture of Gaussian and Multi-Swiss-Roll datasets. The results obtained from the toy dataset, shown in 8ig.8, demonstrate that AGM-SDE is capable of generating data that closely aligns with the ground truth, while requiring NFE that is around one order magnitude lower than CLD. Furthermore, our findings reveal that AGM-SDE outperforms CLD in the context of CIFAR- 10 image generation tasks, especially when faced with limited NFE, illustrated in Table 2.
NFE, CLD-SDE AGM-SDE
20
> 100
7.9
50
19.93
3.21
150
2.99
2.68
1000
2.44
2.46
Deterministic Sampling: We validate our algorithm on high-dimensional image generation with
deterministic sampler. We provide uncurated samples from CIFAR-I0, AFHQv2 and ImageNet- 64 with varying NFE in Appendix.H. Regarding the quantitative evaluation, Table.3 and Table.4 4 summarize the FID together with NEE used for sampling CIFAR-10 and ImageNet-64, Notably, AGM-ODE achieves 2.46 FID score with 50 NFE on CIFAR-10, and 10.55 FID score with 20 NFE in unconditional ImageNet-64 which is comparable the existing dynamical generative modeling.
8

---

Published as a conference paper at ICLR 2024
Table 3: Unconditional CIFAR-10 generative per- Table Unconditional ImageNet-64
formance
generative performance
Model Name
NFE,
FID
Model
NFE,
FID
EDM (Karras et al., 2022)
35
1.84
FM-OT(Lipman et al., 2022)
138
14.45
CLD+EI (Zhang et al., 2022)
50
2.26
ODE
FM-OT (Lipman et al., 2022)
142
6.35
MFM(Pooladian et al., 2023)
132
11.82
AGM-ODE(ours)
50
2.46
MFM(Pooladian et al., 2023)
40
12.97
SDE
VP (Song et al., 2020b)
1000
2.66
AGM-ODE(ours)
40
10.10
VE (Song et al., 2020b)
1000
2.43
CLD (Dockhorn et al., 2021)
1000
2.44
AGM-ODE(ours)
30
10.07
AGM-SDE(ours)
1000
2.46
AGM-ODE(ours)
20
10.55
Table 5: Performance comparing with fast sampling algorithm using FID,, metric on CIFAR-10
NFE 5
10
20
Dynamics Order
Model Name
EDM (Karras et al., 2022)
> 100
15.78
2.23
1st order
VP+EI (Zhang & Chen, 2022)
15.37
4.17
3.03
dynamics
DDIM (Song et al., 2020a)
26.91
11.14
3.50
Analytic-DPM(Bao al., 2022)
51.47
14.06
6.74
2nd order
CLD+EI (Zhang et al., 2022)
N/A
13.41
3.39
dynamics
AGM-ODE(ours)
11.93
4.60
2.60
We underscore the effectiveness of sampling-hop, especially when faced with constrained NFE budget, in comparison baselines. We validate it on the CIFAR-10 and AFHQv2 dataset respec- tively. Fig.4 illustrates that AGM-ODE is able to generate plausible images even when NFE= 5 and outperforms EDM(Karras et al., 2022) when NFE is extremely small (NFE<15) visually and nu- merically on AFHQv2 dataset. We also compare with other fast sampling algorithms built upon DM in table.5 on CIFAR- 10 dataset where AGM-ODE demonstrates competitive performance. Notably, AGM-ODE outperforms the baseline CLD with the same El sampler by large margin. We suspect that the improvement is based on the rectified trajectory which is more friendly for the ODE solver.
Conditional Generation We showcase the capability of AGM to generate conditional samples ing an unconditional model (fig.5) by incorporating conditional information into the prior velocity variable Vn- Instead of employing randomly sampled Vo.
linear combination of Vo and the desired velocity Vi = (x1 Xt, )/(1111 (n). where Xı is conditioned data. Thus, 4o, the initial velocity defined as vgonn = v + (v1, with serving mixing coefficient. Fig.5 shows that AGM can generate conditional data without augmentation and additional fine-tuning Such property can be extended to the inpainting task as well and the detail can be found in appendix.F.
CONCLUSION AND LIMITATION
In this paper, we introduce a novel Acceleration Generative Modeling (AGM) framework rooted in SOC theory. Within this framework, we devise more favorable, straight trajectories for the mo- mentum system. Leveraging the intrinsic characteristies of the momentum system, we capitalize on additional velocity to expedite the sampling process by using the sampling-hop technique, signifi- cantly reducing the time required to converge accurate predictions of realistic data points. Our experimental results, conducted on both toy and image datasets in unconditional generative tasks, demonstrate promising outcomes for fast sampling.
However, it is essential to acknowledge that approach's 's performance lags behind state-of-the-art methods scenarios with sufficient NFE, This observation suggests avenues for enhancing AGM performance. Such improvements could be achieved by enhancing the training quality through the adoption of techniques proposed in Karras et al. (2022) including data augmentation, fine-tuned noise scheduling, and network preconditioning. among others.

---

Published as a conference paper at ICLR 2024
REFERENCES
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Ap- plications, 12(3):313-326, 1982.
Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytie-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. arXiv preprint arXiv:2201.0666 .6600,, 2022.
Arthur Earl Bryson. Applied optimal control: optimization, estimation and control. CRC Press, 1975.
Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos A Theodorou. Deep momentum multi- marginal schr)." odinger bridge. arXiv preprinl arXiv:2303.0155/, 2023.
Yongxin Chen and Tryphon Georgiou. Stochastic bridges of linear systems. IEEE Transactions
Automatic Control, 61(2):266-355 531, 2015.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings the IEEE/CVF conference on computer vision and pattern recognition, pp. 8188-8199 2020.
Valentin De Bortoli, Guan-Horn Liu, Tianrong Chen, Evangelos A Theodorou, and Weilie Nie. Augmented bridge matching. arXiv preprint arXiv:2311.06977 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255, Icee, 2009.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. arXiv preprint arXiv:2/05.05233, 2021.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically- damped langevin diffusion. arXiv preprint arXiv:2/ 2021.
Ulrich G Haussmann and Etienne Pardoux. Time reversal of diffusions. The Annals of Probability, pp. 1188-1205, 1986.
Jeremy Heng, Valentin De Bortoli, Amaud Doucet, and James Thornton. Simulating diffusion bridges with score matching. arXiv preprint 1xXx:2211107222 2021.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.1125: 2020.
Marlis Hochbruck and Alexander Ostermann. Exponential integrators. Acta Numerica, 19:209-286, 2010.
The Math Works Inc. Matlab version: 9.13.0 (r2022b), 2022. URL URLhttpss//ww.mtttweee com.
HJ Kappen. Stochastic optimal control theory. ICML, Helsinki, Radbound University, Nijmegen, Netherlands, 2008
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 35:26565-26577 2022.
Donald Kirk. Optimal control theory: an introduction. Courier Corporation, 2004.
Peter E Kloeden, Eckhard Platen, Peter Kloeden, and Eckhard Platen. Stochastic differential equations. Springer, 1992.
10

---

Published as a conference paper at ICLR 2024
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Christian Léonard, Sylvie Rælly, and Jean-Claude Zambrini. theoretical point of view. 2014.
Reciprocal processes. a measure-
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima
Anandkumar. 12sb: Image-to-image schr)," odinger bridge. arXiv preprint arX772230..0887, 2023.
Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208.1400 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
1rX1vv1711.05101, 2017.
Neil O'Connell. Conditioned random walks and the rsk correspondence. Journal of Physics A: Mathematical and General, 36(12):3049, 2003.
Bernt Øksendal. Stochastic differential equations. In Stochastic differential equations, pp. 65-84. Springer, 2003.
Kushagra Pandey, Maja Rudolph, and Stephan Mandt. Efficient integrators for diffusion generative models. arXiv preprint arXiv:2310.07884 2023.
Stefano Peluchetti. Non-denoising forward-time diffusions. 2021.
Stefano Peluchetti. Diffusion bridge mixture transports, schr\," odinger bridge problems and gener- ative modeling. arXiv preprint arXiv:2304.00917, 2023.
Aram-Alexandre Pooladian, Heli Ben-Hamen Carles Domingo-Enrics Brandon Amos, Yaron Lip- man, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304. /4772, 2023.
Daniel Revuz and Marc Yor. Continuous martingales and Brownian motion, volume 293. Springer Science & Business Media, 2013.
Simo Särkkä and Amo Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Amaud Doucet. Conditional simulation using diffusion schrödinger bridges. In Uncertainty Artificial Intelligence, pp. 1792-1802. PMLR, 2022.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schr\" odinger bridge matching. arXiv preprint arXiv:2303.16852 2023.
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion schr\" "odinger bridges. arXiv preprint arXiv:2302. 1419, 2023.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020a.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint WwXiv:2011.13455 2020b.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training score- based diffusion models. arXiv e-prints, pp. arXiv-2101, 2021.
Robert F Stengel. Optimal control and estimation. Courier Corporation, 1994.
11

---

Published as a conference paper at ICLR 2024
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation- free schr\" odinger bridges via score and flow matching. arXiv preprint arXiv:237.03372, 2023a.
Alexander Tong. Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023b.
Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations, volume 43. Springer Science & Business Media, 1999.
Qinsheng Zhang and Yongxin Chen, Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204. /39022 2022.
Qinsheng Zhang, Molei Tao, and Yongxin Chen. gddim: Generalized denoising diffusion implicit models. arXiv preprint arXiv:2206.05558 2022.
Qinsheng Zhang, Jiaming Song, and Yongxin Chen. Improved order analysis and design of expo- nential integrator for diffusion models sampling. arXiv preprint arXiv:2308.022 (2200.0011: 157, 2023.
12

---

Published as a conference paper at ICLR 2024
A SUPPLEMENTARY SUMMARY
We state the assumptions in Appendix.B. We provide the technique details appearing in Section.3 Appendix.D) D. The details of the experiments can be found in Appendix.E. The visualization of generated figures can be found in Appendix.H.
B ASSUMPTIONS
We will use the following assumptions construct the proposed method. These assumptions are adopted from stochastic analysis for SGM (Song al., 2021; Yong & Zhou, 1999; Anderson, 1982),
(i) po and pi with finite second-order moment.
(ii) Dr. is continuous functions, and [a()) > 0 is uniformly lower-bounded w.r.t. t.
(iii) 14 E [0, 1]. we have V. (e () Lipschitz and at most linear growth w.r.t. and V.
Assumptions (i) (ii) are standard conditions in stochastic analysis ensure the existence-uniqueness of the SDEs; hence also appear in SGM analysis (Song et al., 2021).
C STOCHASTIC OPTIMAL CONTROL (SOC) IN THE WILD
In this section, we are going to provide gentle introduction of the Stochastic Optimal Control (SOC). Our work is majorly relying on the prior work Chen Georgiou (2015) in which some technical details are missing. Here we first clarify some core derivations that may help the broader audience to understand Chen & Georgiou (2015) and our work.
C.I LINEAR QUADRATIC STOCHASTIC OPTIMAL CONTROL
SOC has wide applications in financial, robotics, and manufacturing. Here will focus Linear Quadratic SOC which usually refers to Linear Quadratic Regulator because the dynamic is linear and the objective function is quadratic (Bryson, 1975; Stengel, 1994). The problem states as:
min f 1u + +xxxxxxxx Rxi
2
(13)
s.1. dx, [A(t)x, + t. u + gi X = To-
In this formulation, X, means the state, and u, is the control variable. Conceptually, the SOC problem is aiming to design the controller u, to drive the system from point To to Iı 0 with minimum effort. In the case of first order system, the control will be the optimal vector field v and for the second order system, the control is denoted as the optimal acceleration at The presence of stochasticity, introduced by the Wiener Process denoted as dw,, prevents the system from precisely converging to the Dirac mass æ1 In order to strike balance between the objective of converging to an and minimizing overall control effort f'|u, dld the terminal cost x] Rxx has been imposed.
One special case R -> oc. Intuitively, means the controlled dynamics should precisely converge to Fi. However, one can notice that the stochastic trajectory which connects 2o and æ is not unique in this case. Based on this constraint (pinned down at #1 and To at two boundaries), the optimization problem of SOC finds the optimal solution with minimum effort u, which can be understood the regularization the trajectories, hence, such stochastic trajectory is unique while the regularization of controller is still applied. One can also draw connection with such pinned-down SDE with well- known Doob-h transform. For the people who are not familiar with these, here are some interesting paper (Heng al., 2021; O' Connell, 2003).
The classical procedure to solve the SOC problem includes:
1. write down the Hamilton-Jacobi- Bellman equation (HJB PDE) which explicitly represents the propagation of value function over time.
2. Construct the Ricatti/I yyapuoo Equation.
3. Solve Ricatti/
Ricatti//_yapunov yapunov Equation and obtain the optimal control.
13

---

Published as a conference paper at ICLR 2024
C.2 VALUE FUNCTION, HAMILTOONJACOBIAN (HAMILTON -AAOOIJ--OOLLLLLLA EQUATION) AND RICATTI EQUATION
We adopt the classical notation in the SOC for the value function. Specifically, the underscript of the value function V represents for the partial derivative of it. For example, V,, Va and Vze represent for the first order derivative of V w.r.t time t, state x and second order derivate of V x. We first define the value function as:
V(G,,) t) inf E S. I 1127 + x Rxi
and the dynamics is,
dx, (Ax, + go dt + gidw;
From Bellman's principle the value function, one can get:
V(L dt, Xx+de) + [. r dt
V(t,x,) inf IE
2 u dr
inf IE
u, [2d + x() + Vi(1, xz)d V2 (t, x)xx 7/ [V2=99]) dt
Plug in the dynamics dx,
inf E
u Edt + (tt + xttt + V2tt, ((Ax, + gou)ddt + gdw))
+ tr [V.n.g "
dt
inf [
|uuel [dtt v(t, + x,)t, + v.(ttx))(AAs (Ax: + gru))
2
+ t tr [V...99] dt
One obtain:
1
V2 inf
lu 12 + V.T(Ax, + geus)
+ tr [Va=9g 0
The optimal control can be obtained by
u1 = -
Plugging it back, obtain the HJB PDE:
Vi 2 Vagg"ve V. + V Ax; + 1 tr [Vangg]] =0
We assume that there exist certain matrix Q. s.t. V(x,t) = 1 different power term of HJB, one can write:
By matching the
x x -xx
1 Togg Tx +x A +
1 1r (999
(14)
2
with boundary condition:
E(1) = 0, Q(1) =
(15)
Due to the fact that x
TTQQ== = TQAxx one arrives Riccati Equation:
- A- Q99 TQ
(16)
Recall that the optimal solution is u: V and V = + E(t,, the optimal control can be expressed in the way of the solution of Ricatti equation: un = -Q)((
14

---

Published as a conference paper at ICLR 2024
C.3 RICATTI EQUATION AND LYAPUNOV EQUATION
Here we provide the connection between Ricatti Equation and Lyapunov Equation in the current setup.
Lemma 6. Define P(t) Q(t) in which Q(t) is the solution of Ricatti equation (eq.16) Then P(I) solve the Lyapunov equation:
AP PA 9g
(17)
For notation consistency, name the elements in P matrix as,
Poo
Poi
P =
P1o
P11
Proof. By plugging in the Lyapunov equation P(t) == Q(t) -1, one can get:
Q AQ + Q- 99
99 Q QQ = - -99*
# Q = QA + A'Q Q9g Q
By Lemma.6, 6, the optimal control can also be represented as the solution of the Lyapunov equation: -. = -g" P(t)-1x, x, which is indeed the optimal control term used in Chen & Georgiou (2015) after adopting their notation, and is same as the optimal control term we used in the Lemma. 12 without base dynamics compensation.
C.4 SOC CONNECTION SCHRÖDINGER BRIDGE
The optimal control solution is also the solution of Schrödinger Bridge when the terminal condition degenerate to the point mass (see example of Brownian Bridge in Appendix.D.1), is also the solution of the Schrödinger Bridge when the optimal pairing is available see proposition.2 De Bortoli et al. (2023).
So in our case, we are not solving the momentum Schrödinger Bridge as shown in Chen et al. (2023) (also see. 6ig.6), even tough the problem formulation is similar. Specifically, AGM a special case of momentum Schrödinger Bridge when the boundary conditions are degenerated to Dirac Distribtuions.
Po
Remark: Optimal Transport needs correct pairing
( Po 2 (Fo, ~ I1 P,) + SOC
solution
Phase Schrödinger Bridge
P1
SOC
Figure 6: momentum Schrodinger Bridge versus AGM.
15

---

Published as a conference paper at ICLR 2024
D TECHNIQUE DETAILS IN SECTION.3
D.1 BROWNIAN BRIDGE AS THE SOLUTION OF STOCHASTIC OPTIMAL CONTROL
We adopt the presentation form Kappen (2008). We consider the control problem:
min J. [
ur (l2
2 X
x1112
u.
s.t. dx; u,dd,
Where is the terminal cost coefficient. According to Pontryagin Maximum Principle (PMP; (2004)) recipe, one can construct the Hamiltonian:
H(t, x, u, 7) = - 1 yu
By setting:
0H
0,
Su,
the optimized Hamiltonian is:
X(,,,, = 2? where = y
Then we solve the Hamiltonian equation of motion:
dx;
OH
di
Oy
dy
0H
di
8x
where
Xo zo
and
?1 r (x1 -x1)
One can notice that the solution for 7t is the constant for x, is X2 X1 76.
= 7 - -xx 21), hence the solution
-r(x- - x1) -r(x)
+ (1 - t)7 x1)
u = r(z1 Xo)
r(1 t)
When r > +00, we arrive the optimal control as U. also the optimal control law for
21 XD I-t
Due to certainty equivalence, this is
dx, u + dw,
By plugging it back into the dynamics, we obtain the well-known Brownian Bridge:
dx; = .1 X; dt + dw;
Remark 7. there is not stochasticity dw,, one get 21-x4 I1 Xo which is the vector field constructed by Lipman et al. (2022) during traning.
D.2 PROOF OF PROPOSITION.3
Proposition 8. The solution of the stochastic bridge problem of linear momentum system (Chen & Georgiou, 2015) is
ammm,, 1) = = P1
1x- -1--x
vo v.)
where : Pu1
(18)
1--
g 1)
16

---

Published as a conference paper at ICLR 2024
Proof. From Lemma. 12, can get the optimal control for this problem is
u -g P.Pl (m, (t, 1)m1)
where state transition function 4 can be obtained from Lemma. and P; is the solution of Lyapunov equation and P," can be found in Lemma.9.
Then we have:
".
8g
P
m (11, 1)m)
gg
P
m, gg PP-1t(ll
(1, 1)m;
0
0
g²
P
m, gg PP:
0
m1
0
0
Puo
Poi
1
t-
Pio
Pn
m.
0
Piu
Puj
1
mı
[
t-1 -1
gi
Pio
P11
m.
Pio
Pn1
m1
0
0
-
Pie
P11
m. + gi
Pio
Pio
1)
P11
m1
1 vz)
g P1o(xI x1) + oe Pio 1) V1
Pni
X1 Xt
Plug in
--t
0
PP1
x1=Xi
v.)
1-t
Lemma 9. The Lyapunov equation corresponding the optimization problem showed in Lemma. 12:
1 |lull dt + R
u, arg min E
2
well
0
8.
dm,
midi + u,di + gdw.
0
0
mg
ma.
m
m
depited as
P AP+PPP 9g"
(19)
00]
When
the solution for Lyapunov equation above, with terminal condition
g
0
0
0 0
P
lim
(20)
However, one does not need the force to converge exactly Vı because only care about the generate quality of X1. Here we give a general case which the r keeps a small value w for the velocity channel:
0
P
R
(21)
0
W
17

---

Published as a conference paper at ICLR 2024
Then the solution is given by
w(t - 1)² -1 2t 1)3 w(1 - 1) 2 - 1)2
wit - 1) g - 1)²
P.
g² - 1) + +
and the inverse of Pe is,
12 (t--))((-11)) (t 112
6 20+g2((11)
P.L g + g²² - 1))(t -1
1-1
6 (- 1-
12w - 4g (t - 1)
Thus,
-12w + 6t 1)
-12
Pio g21-4w 4w + g2 1)](t - 1)² g21-4w + g² -1)2
|- + g² - 1)](t 1)
12w 4g² 1)
12w
-4
P1 g Aw + (t 1)] 1)
g²1-4uu 4w + ge - [)] - +
[-4w + g2 - 1)]
Proof. One can plug in the solution of P, into the Lyapunov equation P. and it validates P. is indeed the solution.
Remark 10. Here we provide general form when the terminal condition of the Lyapunov function not a zero matrix. explicitly means that it allows that the velocity does not necessarily need to converge to the exact predefined V1. It will have the same results shown in the paper by setting O.
Lemma 11. The state transition function Φ(t, s) of following dynamics.
[
dm,
m,df
0
0
is,
1 t t- s)
(
0
0
Proof. One can easily verify that such & satisfies 04/01- =
Φ.
0
Lemma 12 (Chen & Georgiou (2015)). When R + co, The optimal control uI of following prob- lem,
0
[" - [u dt + x3 Rxi
u
E arg min m EM
Jo 2
a,
[
s.t
dm,
mtdf + u,dt + g,dw;
0
mg
mo
given by
=. = -- Where P follows Lyapunov equation (eq. 19) with boundary condition P function ( s) is the transition matrix from time-step s to time-step t given uncontrolled dynamics.
T PP-''mm
- (tt, 1)m;)
And is indeed the stochastic bridge of following system:
0
dm, =
m.df + u,dt + gdw:
(22)
mg mo:
m, = m1
(23)
Proof. See page 8 in Chen & Georgiou (2015).
18

---

Published as a conference paper at ICLR 2024
D.3 MEAN AND COVARIANCE SDE
By plugging the optimal control into the system, one obtain the system
dm:
dt + gidw;
Fr
dt + ggdw,
9
P11
xI-X1
vz)
1-t
1
Xt
P11
Vi
dt +
1-t x
Pu1
dt g.dw;
F.
D.
We follow the recipe of Särkkä Solin (2019). The mean μ4 and variance Σ, of the matrix of random variable m, obey the following respective ordinary differential equations (ODEs):
dup: F, di + D.dt
dE, = F, + | dt + gg dt
One can solve by numerically simulating two ODEs whose dimension is just two. Or can software such as Inc. (2022) to get analytic solutions. If you opt to the later approach, you can get:
Pi xt 4t 6)
D = 4/x1
- +
3
1
MEE = 9
EN 1
8
Nev =
(33
Remark 13. The expressions above are too complicated. Hence, we provide the python functional bracket in Appendix. with general initial covariance and diffusion coefficient for easy copy-paste. Equations above are ones we used through this paper and feel free to play around with other hyper- parameters.
D.4 DERIVATION FROM SDE TO ODE FOR PHASE DYNAMICS
One can represent the dynamics the form of,
dx;
dv;
F
V.
dí
0
9t
dw, s.t
mo
Vo
X ~N(1po: Eo)
(24)
dm, fm:) + g,dw,
And its corresponding Fokker-Planck Partial Differential Equation Øksendal (2003) reads,
Op.
[f, (m.)) + 222 mm,m
02
Σ
g.g pe pe(m)
(25)
06
0m:
19

---

Published as a conference paper at ICLR 2024
According to c3.( in Song et al. (2020b), One can rewrite such PDE,
Öpe Σ
si (pm
[p(m.)Vm (gig)) + p(m)) /V Vn logp(m))
at
dm:
(26)
due to the fact g:
(27)
0 gt
Σ 0m:
J(mee (m;))))) jr(m:) 192 V. log
(28)
Then one get the equivalent ODE:
dm, [(mm.)) - 5 V. logp(m,
(29)
D.5 DECOMPOSITION OF COVARIANCE MATRIX AND REPRESENTATION OF SCORE
Here follow the procedure in Dockhorn et al. (2021). Given the covariance matrix E,, the decomposition the positive definite symmetric matrix is,
Σ, L.
(30)
Where,
V Dg 0
L L
L
L
(31)
L
VE
We borrow results from Dockhorn et al. (2021), the score function reads,
Vm log p( mn mm)) = - V m:
(m. mx) E. 2m, p4)
2
= "(m, pe)
Cholesky decomposition of E,
= LL- (mm me)
-T
The form of L reads,
V VI V
L
F2s VE
and the transpose inverse of L reads,
L-"
(DEE)
(EX") -
Hence, the score function reads,
Vy log pm ( mi) =- ((E + Coz (Σ + (vo) (I
20

---

Published as a conference paper at ICLR 2024
D.6 REPRESENTATION OF ACCELERATION ay
As been shown in Proposition.3, the optimal control can be represented as,
aj = 96 P1
Xr
V/
1
9. P11
P1
1-t +V1
X1
1-t
o P11
P1
Mi
+ L ED
+ (pg Li ED + Li E1)
1-1
1
[xx1 1-t
96 Pı1
pi)
Eo
+ Li Eu + Li
61
1
4/x1
solving D.3 we can get u
xx
t2
4t + 66
(t2 3t + 3)
3
3
Plug inx;, V
X1 3 x112 12 (6
41+11 +12
4|x1
1
= 9 Pn
(t2 - 3t +3) -
1- E Lt + 11d
1-t
3
ta + 4tt 6t + 3)
4t
2 Pn
(t2 31 + 3) x1
Co + I 6D + L, 61
3(1 - t)
1
(t - 1)( 3t + 38 3)
4t
L
of Pu
(t2 3t + 3)
X1
l-t Co
+ L Eu L E1
3(1
3
(t 32 + 3t + 3)
1
L
of Pn
(4t
-12t 126)
x1
E0 + 2 ED + 4 E1
LEz
96 Pn
(1- (1-1)xx
Eg + 11 60 + L
1-1
Law
4(1 x) + P1
Eo + _ ED + L E1
D.7 LOSS REWEIGHT
In practice, we use the following loss function
[ min Exejo. Ex, palata E (m)
[ (m,, t: 0) F, F,,,, t)
(32)
X min EE 10. Exx Me (mm (x) T-7
F (m,, 0) F. 12
(33)
We admit that this might not be optimal selection. The motivation behind this is simply increasing the weight of training when + 1 and normalize the label with normalizer 2t-
D.8 NORMALIZER OF AGM-SDE AND AGM-ODE
Since the optimal control term can be represented as,
L
Fx
a" (m,, t) = 4x (1 - 1)2 - 92 P11
1-1
L
RF
6n
+
L PM
E1
Then we introduce the normalizer as
Li
2
ZSDE
(4(1 - 1112 D2 + 9 +
1-1 +
L Lix
+ (Livgz
2
1 27
ZODE
(4(1 - t)2 -2doda 22 + P1
+ 20
+
gf
PuLL - -
1-t
21

---

Published as a conference paper at ICLR 2024
Where 8== Spp (2
D.9 EXPONENTIAL INTEGRATOR DERIVATION
As suggested by Zhang Chen (2022), one can write the discretized dynamies as,
Xt: V4:11
= ((t++1,te)
X Vi
+
Σ
Cis
so (m,, 5: [i-
0
==
(34)
0
link Li-j li-k
1
t - s
Where Oi. =
+(t+
dr, st,,
Zr
0
After plugging the transition kernel ((t, s), one can easily obtain the results shown in (11).
Remark 14. In light of the momentum system, there are mumerous methods for achieving high accuracy its resolution. However, the practical performance in generative modeling remains untested. We recommend that readers consult the classical numerical physics text book or recent momentum dynamics solver (Pandey al., 2023; Dockhorn al., 2021).
D.10 PROOF PROPOSITION.5
The estimated data point can be represented as
SDE (1 (F + vi). X1, xDDE=
F.P + P1 + +Bie))
(35)
P1
4(t 1)2 + g. P11 + 111p))
14
--LL
for SDE and probablistic ODE dynamics respectively, and 3, I 2P11 at
1,a
Proof. It is easy to derive the representation of x of the SDE due to the fact that the network is essentially estimating:
X1 -
Fl o P11
1 1-1
(1
FF? + vv)
+ x1
+ Xt
P11
It will become slightly complicated for probabilistic ODE cases. We notice that
Mr Le
X; = 10z + L E1, V, = 1he + L Eo + Li El
In probabilistic ODE case, the force can be represented as,
L 1- +
F F(m,, 4x, (1 t)² P1
+ L c1
,giles
In order to use linear combination of x, and represent F one needs to match the stochastic term F, by using
L 3
BLLE
L + 2P1
The solution can be obtained by:
Bt = G
L
& B, L
0t
L
22

---

Published as a conference paper at ICLR 2024
By subsitute it back to F2, can get:
F(m,, t) = 4x (1 - t)² - Pi (x -2)) + Be (.viv-. -1
[[4(1 - + +) 9 + ((pp)) PI [axx, + 3.ve]
F. + Pn + (eve)
X1
4(t 1)2 + P P11 (a + ((pt)
E EXPERIMENTAL DETAILS
Training: We stick with hyperparameters introduced in the section.4. We use & Hutter, 2017) as our optimizer and Exponential Moving Averaging with the exponential decay rate of 0.9999. We use 8 Nvidia A100 GPU for all experiments. For further, training setup, please refer to Table.6.
Table Additional experimental details
dataset toy
Training Iter
Learning rate
Batch Size
network architecture
0.05M
le-3
1024
ResNet (oockhrn al., 2021)
CIFAR-10
0.5M
le-3
512
NCSNa+ et al., 2022)
AFHQv2
0.5M
1e-3
512
(Karaa et al., 2022)
ImageNet-64
1.6M
2e-4
512
ADM(Dhariwal Nichol, 2021)
Sampling: For Exponential Integrator, we choose the multistep order w 2 consistently for all experiments. Different from previous work (Dockhorn al., 2021; Karras al., 2022; Zhang al., 2023), we use quadratic timesteps scheme with A = 2:
N N
t,
N
Which is opposite to the classical DM. Namely, the time discretization will get larger when the dynamics propagated close data. For numerical stability, use to 1E-5 for all experiments. For NFE = 5, we use tN = 0.5 and NFE = 10, Tn = 0.7. For the rest of the sampling, we use ty 0.999.
Due to the fact that EDM(Karras et al., 2022) is using second-order ODE solver, in practice, we allow it have an extra one NFE as reported for all the tables.
E. CODE EXAMPLE FOR COVARIANCE
We will abuse the notation in this coding section. Here provide the example code for compute -kymm
the covariance matrix. Here we consider the general case where Eg = -kymn n
and
the diffusion coefficient is g(t) == p(tt - t) where p is the scaling coefficient and the damping coefficient.
def Sigmaxx (t, P, tt, n) return
(t **2* (30 +m* (t**3 3+ + 3xt + 3) +*21
60*p**2* (t 1) 3* (1 t)
t* sqo* (m*n) x (t**5 - 6*tx*4 + 15xt**3 15xt*42 + 9)
30*22 - 3xt 3) **2 + p**2* (t**5* (t*A5x (5+tE**2 + 3xtt + 1)
6<tt+44 3*tt + 1)
+ 15+t++3+ (6+tt++22 3+tt + 1) 100t*+2> (9*tt**2
11) + 150xt 60))
23

---

Published as a conference paper at ICLR 2024
def Sigmaxv tt, return
(1/270 t/270) (30=k=np. (m+n) *8*t**6 *8*
+ 124 -135xt**3 45+t**2 27+t 9) +1
150*p**2* =1)
- t)
t*(1120xmm (t**5 6*t**4 15*t**3 15xt**2 9)
(00***(xxx*< - 24*t**4 + 60*t**3 75*t**2 + 45xt 9)
+ (4***<(xt**** 3+tt + 1) 5x 2 18+tt + 55)
3xtt 1) 24+t** (6+II=*2 + 3+tt 1
+ 60+t**3<<6*tt***
+ 15*t* + 25) 150))
def Sigmavv (t, tt, n) return
n* (-4+t**3 + 12*t**2 - 12xt + *
**p**2** ***tottchh 1og (1 t)/91
+ t*(-120=**np.sqrt (m*n) (4*t*+5 24+t**4 + =0*t**31
15*t**2 45xt 9) 240xm*t: (t**2 -3xt 3)
+ (6 3xtt + 1) 5* (180 72*tt 53)
+ 3xtt + 1) + 48*I**4* (Exttx*0 + 3xtt + 1)
+ 9*tt + 20) + 135+tt**2 + 120)
F CONDITIONAL GENERATION DETAILS
Here we provide the detail of conditional generation details.
F.1 STORKE BASED GENERATION
For stroke based generation, we provide two types of conditional generation.
initial Velocity (IV): refer section.4.
Dynamics Velocity (dyn- Sinice the mean and variance of velocity and position are available, one can specify the velocity which valid. In this case, we can set the velocity as
1t 1h
velm la
D lote
(36)
In which,
14
Yee
(x:
pi
(37)
(38)
when & S c. The e the guidance length. We typically set it to be c 0.25.
F.2 INPAINTING
In the inpainting case, apply the similar strategy dyn-V. Specifically, in this case, the X; will be represented as:
x == MASK C n + (1 MASK) O X
(39)
where MASK represents for the mask matrix which zero-out the pixel of the original image. Such &1 will serve the source estimate ue in eq.37.
F.3 INPAINTING BASED GENERATION
For stroke based generation, provide two types of conditional generation.
24

---

Published as a conference paper at ICLR 2024
G ABLATION STUDY OF STOKE-BASER CONDITIONAL GENERATION
In order to investigate the diversity and faithfulness of stoke-based conditional generation, we con- duct the ablation study with respect the hyperparameter &.
$ 0
= 0.05
ğ = 0.10
ğ 0.12
Reference stoke
c 0.15
5 0.20
Figure 7: Ablation study the stoke-based conditional generation. When = is unconditional Inmeration.Notable the diversity of the generation will decay when we increase &. In order to achieve a balance between faithfulness and diversity, one needs to tune the hyperparameter, ç.
H ADDITIONAL FIGURES
We demonstrate the samples for different datasets with varying NFE,
H.1 TOY DATASET COMPARED WITH CLD
500, ME=10. SDE
Ground Trubh
MEE AGMiours) NFE-10, SDE
Ground Truth
Figure 8: The comparison with CLD((ockhorn et al., 2021) using same network and stochastic sampler SSS, for Multi-Swiss results with order less NFEs.
Sirii-Swis--ooo Roll and Mixture of Gaussian datasets. We achieve visually better
25

---

Published as a conference paper at ICLR 2024
H.2 AFHQV2 INPAINTING GENERATION
Conditional Data
Generated Data
Figure 9: AGM-ODE Uncured inpainting generation
H.3 AFHQV2 STROKE BASED GENERATION
Conditional Data
Generated Data
Figure 10: AGM-ODE Uncured stroke based generation
26

---

Published as a conference paper at ICLR 2024
H.4 CIFAR-10
Figure 11: AGM-ODE Uncurated CIFAR-10 samples with NFE-5
27

---

Published as a conference paper at ICLR 2024
Figure 12: AGM-ODE Uncurated CIFAR-I0 samples with NFE=10
28

---

Published as a conference paper at ICLR 2024
Figure 13: AGM-ODE Uncurated CIFAR-I0 samples with NFE=20
29

---

Published as a conference paper at ICLR 2024
Figure 14: AGM-ODE Uncurated CIFAR-I0 samples with NFE=50
30

---

Published as a conference paper at ICLR 2024
H.S AFHQV2
Figure 15: AGM-ODE Uncurated AFHQv2 samples with NFE-5
31

---

Published as a conference paper at ICLR 2024
Figure 16: AGM-ODE Uncurated AFHQv2 samples with NFE=10
32

---

Published as a conference paper at ICLR 2024
Figure 17: AGM-ODE Uncurated AFHQv2 samples with NFE=20
33

---

Published as a conference paper at ICLR 2024
Figure 18: AGM-ODE Uncurated AFHQv2 samples with NFE=50
34

---

Published as a conference paper at ICLR 2024
H.6 IMAGENET-64
Figure 19; AGM-ODE Uncurated Imagenet-64 samples with NFE=10
35

---

Published as a conference paper at ICLR 2024
Figure 20: AGM-ODE Uncurated Imagenet-64 samples with NFE=20
36

---

Published as a conference paper at ICLR 2024
Figure 21: AGM-ODE Uncurated Imagenet-64 samples with NFE=50
37


Reviews:
[{"review_id": "pFGsU7EJ91", "reviewer": "Reviewer_gcJ1", "summary": "This paper proposed the acceleration generative model (AGM), which is Bridge Matching method with a dynamics-based diffusion model with stochastic optimal control (SOC) theory that rectifies the trajectory of the second-order momentum dynamics which is first introduced in CLD. First, the optimal acceleration function is derived to the solution of the stochastic bridge problem, which is given by minimizing the SOC objective function. Different from CLD whose velocity field is defined by the score function of the pre-defined critically-damped Langevin diffusion process, the velocity field of AGM is learned to rectify the particle trajectory. This SOC problem is designed to minimize the distance between the ground-truth (GT) destination and the trajectory destination. Then like in CLD, this paper took advantage of the momentum-based approaches and proposed that the sampling-hop, the estimated data point $x_1$ given the early sampling stage outputs, is predicted more accurately compared to existing methods. In the empirical experiments, the sampling quality is improved especially in the low-NFE regime.", "questions": "* It will be helpful if the acceleration coefficient $a_t$ for image datasets is depicted, as the momentum of how the image data is being generated in Figure 1 or Figure 2. It is expected that the acceleration coefficients show similar semantic features like $x_t$ and $v_t$, but have varying scales.\n* Can you provide elementary introduction of the stochastic optimal control? While this paper works only the simple case of the SOC (no regularization case), introducing some details, or at least some introductory materials will help the readers to follow up the backgrounds.\n* I guess that the ImageNet64 performance is not yet optimized: the generative performance of the SoTA models are expected to be much better than the paper have proposed. I think at least the performance should be compared with CLD-SGM from the same architecture.\n\n================================\n\n* It will be helpful for the readers' understanding if you use the colored hyperlinks by reference (\\ref) or citation (\\cite, \\citep, \\citet) commands.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The idea of rectifying the particle trajectory with the velocity field is an intuitive approach, which is widely used in the literature. Existing works used handcrafted way to design the velocity field, but this paper aimed to both optimize this part of the stochastic process by using the SOC theory.\n* The objective is well-defined: when the control regularization term approaches to zero, then the objective directly turns into the square mean\n* Even though the velocity should be trained, the whole training process is simulation-free: we do not need any further simulation process like in current SoTA models that require further self-distillation for high-quality image generation in low-NFE regime. Furthermore, this method can be pipelined with the distillation techniques like other diffusion model methods.", "weaknesses": "* The clarity of the paper will be better if the conditions of the Lemmas and Propositions written in this paper is stated more concretely and with full notations, especially in the appendix.\n  - In the sampling-hop part, the writing does not fully cover how the sampling-hop is more accurately evaluated compared to the CLD case. Both this method and CLD make predictions of the data from both the current state and the velocity, while the compared EDM (Figure 2) does it from state alone.\n  - In the Probabilistic ODE part of (7), an additional notation rather than $g(t)$ is recommended to be used, like $g_t \\to h_t$ in the matrix notation and $g(t)=h_t$ for BM-SDE part. Because the notation $g(t)$ or $g_t$ is abused, it can be misleading that the score term of the probabilistic ODE is neglected.\n* The SOC theorem is only used limitedly; the regularization in terms of $\\int ||a_t||^2$ is ignored and this can threaten the stability of the acceleration space, even though this is not directly revealed in the paper.\n* Whereas the theoretical background is sound and the improved performance is guaranteed, the hyperparameters such as the diffusion coefficients and the SDEs are not optimized, which causes its lacking performance compared to EDM (look at Figure 5). However, this is expected to be enlightened with further works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9XwMoPKMEu", "reviewer": "Reviewer_rgEV", "summary": "The work proposes Acceleration Generation Modeling (AGM) as an extension to Critical-damped Langevin Dynamics (CLD) based on the theoretical results of stochastic optimal control. The proposed acceleration term has the effect of straightening the sample trajectories in the sampling process and reducing sampling complexity. The linearity of sampling trajectory enables the AGM generation process to take less number of evaluations and make sampling hops. AGM is compatible with both deterministic (ODE) and stochastic (SDE) samplers. Experiment results on CIFAR-10, AFHQ, and ImageNet show that AGM demonstrate competitive results with less number of evaluations with smaller number of evaluations.", "questions": "In Table 4 which shows experiment results on ImageNet 64, DDPM uses a stochastic sample while the other approaches including FM-OT, MFM, and AGM-ODE all use deterministic samplers. This may not be a fair comparison because even for the same type of diffusion model, the sample quality and sampling efficiency could be very different with different types of sampler and deterministic samplers based on ODE numerical methods generally take less number of steps than stochastic sampler. I would suggest the author include both AGM-SDE and AGM-ODE results under different number of evaluations.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. Overall speaking, the proposed idea is simple yet effective and the motivation is backed by solid theoretical results in the domain of stochastic optimal control. \n2. Both quantitative and qualitative results support the motivation of AGM. Quantitative results under different settings show AGM is able to achieve competitive or better results with similar to less number of evaluations. Qualitative results also show the better ability of AGM to make sampling hops and recover the denoised images at an early stage compared to CLD. \n3. The presentation of the work is also of high quality. The introduction of the theoretical results is concise but also critical to motivate the proposal of AGM. The rest of Section 3 presenting AGM in technical details is also well-structured and easy to follow.", "weaknesses": "The work does not have significant weakness. Minor weakness points include\n1. The work only shows experiment results on CIFAR-10, ImageNet 64, and AFHQv2 without scaling to higher resolution images.\n2. As the author points out in limitations, AGM is not performing as good as some existing methods especially when the number of evaluations is large. I don't think this is a major weakness as the major benefit of AGM and straight sampling trajectories is the reduced number of evaluations during sampling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "saz6dq8ncL", "reviewer": "Reviewer_76FN", "summary": "The papers uses the tools of stochastic optimal control theory to define the forward pass for a kind of generative diffusion model strictly related to Diffusion Schrödinger Bridge Matching. The approach combines the velocity augmentation used in Critical-damped Langevin Dynamics with the bridge approach by solving a linear Gaussian control problem in closed-form. This solution leads to relatively straight paths that are suitable for fast-sampler acceleration both in the stochastic and in the deterministic case. The method has competitive performance for small numbers of functional evaluations, but it lags behind other methods when more evaluations are used (>100).", "questions": "I find the pseudocode in Algorithm 1 and 2 to be rather uninformative. A good pseudo-code should allow the reader to implement the algorithm almost without referring to the rest of the paper. In this case, the most important parts of the code (e.g. the form of the loss) are omitted. Could you update it to make it more self-contained? \n\n- The idea of the initial velocity conditioning  is interesting, but it is difficult to evaluate its potential without quantitative results and comparisons. Intuitively, it seems to me that it will likely lead to a substantial drop in diversity. Can you report the FID for the conditional sampler?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper uses stochastic control theory effectively in order to construct a forward process with the desired properties. I believe that this is highly promising as optimal control and diffusion modeling are deeply related and many advanced control techniques can be imported in the diffusion literature using a similar approach. \n\n- The paper provides a rigorous description of the algorithm and the math behind it without requiring an excessive level of mathematical sophistication in the reader. \n\n- The experiments are rigorous and comprehensive and properly show the performance profile of the method and most relevant baselines under different conditions.", "weaknesses": "- The exposition is rather dense and, as a consequence, the paper is somewhat difficult to read. This is a pity since the underlying concept are rather intuitive and can be understood by a wide audience. \n\n- As also stated by the authors, the performance of the method is inferior to several baselines for a large number of functional evaluations. However, I do not think that this is a major issue since this class of models are generally designed to work well in the low NFE range, and the results are good in this relevant range. It is quite intuitive to me that there should be a trade off between straight paths and high NFE performance, since the smoothness constraints can limit the expressivity and probabilistic coverage of the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SSFef9xelE", "reviewer": "Reviewer_5NhC", "summary": "The paper presents a novel generative modeling framework called Acceleration Generative Modeling (AGM), which is grounded in phase space dynamics. The authors leverage insights from Stochastic Optimal Control to construct a path measure in the phase space that enables efficient sampling. The framework demonstrates the capability to generate realistic data points at an early stage of dynamics propagation, which sets the stage for efficient data generation by leveraging additional velocity information along the trajectory. The model yields favorable performance over baselines in the regime of small Number of Function Evaluations (NFEs) and rivals the performance of diffusion models equipped with efficient sampling techniques.", "questions": "Can the AGM framework be applied to other domains beyond image generation, such as natural language processing or time series data?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The proposed AGM framework offers a new perspective on accelerating sampling in generative modeling by leveraging additional velocity information.\n\n2. The model demonstrates competitive results compared to diffusion models equipped with efficient sampling techniques, particularly in small NFE settings.\n\n3. The paper provides a clear and detailed explanation of the AGM framework, its training, and sampling procedures.", "weaknesses": "1. The paper could provide more insights into the potential applications of the AGM framework beyond image generation.\n\n2. The paper could discuss potential improvements to the AGM framework, such as enhancing the training quality through data augmentation, fine-tuned noise scheduling, and network preconditioning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)