Bespoke Solvers for Generative Flow Models
Authors: Neta Shaul
Juan Perez
Ricky T. Q. Chen
Ali Thabet
Albert Pumarola
Yaron Lipman
Venue: ICLR 2024 spotlight (2024)
Keywords: generative models
flow matching
diffusion models
normalizing flows
ode solver
fast sampling
distillation
Primary area: generative models
TL;DR: A novel method to build costume-made ODE solvers for sampling pre-trained diffusion/flow models. Significantly improves generation quality for low number of function evaluations.

Abstract: Diffusion or flow-based models are powerful generative paradigms that are notoriously hard to sample as samples are defined as solutions to high-dimensional Ordinary or Stochastic Differential Equations (ODEs/SDEs) which require a large Number of Function Evaluations (NFE) to approximate well. Existing methods to alleviate the costly sampling process include model distillation and designing dedicated ODE solvers. However, distillation is costly to train and sometimes can deteriorate quality, while dedicated solvers still require relatively large NFE to produce high quality samples. In this paper we introduce ``Bespoke solvers'', a novel framework for constructing custom ODE solvers tailored to the ODE of a given pre-trained flow model. Our approach optimizes an order consistent and parameter-efficient solver (e.g., with 80 learnable parameters), is trained for roughly 1\% of the GPU time required for training the pre-trained model, and significantly improves approximation and generation quality compared to dedicated solvers. For example, a Bespoke solver for a CIFAR10 model produces samples with Fréchet Inception Distance (FID) of 2.73 with 10 NFE, and gets to 1\% of the Ground Truth (GT) FID (2.59) for this model with only 20 NFE. On the more challenging ImageNet-64$\times$64, Bespoke samples at 2.2 FID with 10 NFE, and gets within 2\% of GT FID (1.71) with 20 NFE.

Published as a conference paper at ICLR 2024
BESPOKE SOLVERS FOR GENERATIVE FLOW MODELS
N. Shaul J. Perez? R. T. Q. Chen A. Thabet A. Pumarola? Y. Kippman Weizmann Institute of Science "GenAI, Meta "FAIR, Meta
ABSTRACT
Diffusion or flow-based models are powerful generative paradigms that are no- toriously hard to sample as samples are defined as solutions to high-dimensional Ordinary or Stochastic Differential Equations (ODEs/SDES) which require a large Number of Function Evaluations (NFE) approximate well. Existing methods to alleviate the costly sampling process include model distillation and designing dedicated ODE solvers. However, distillation is costly train and sometimes can deteriorate quality, while dedicated solvers still require relatively large NFE to produce high quality samples. In this paper we introduce "Bespoke solvers". a novel framework for constructing custom ODE solvers tailored to the ODE of given pre-trained flow model. Our approach optimizes an order consistent and parameter-efficient solver (e.g., with 80 learnable parameters), is trained for roughly 1% of the GPU time required for training the pre-trained model, and sig- nificantly improves approximation and generation quality compared to dedicated solvers. For example, a Bespoke solver for a CIFAR1O model produces samples with Fréchet Inception Distance (FID) of 2.73 with 10 NFE, and gets to 1% of the Ground Truth (GT) FID (2.59) for this model with only 20 NFE. On the more challenging ImageNet-64x666 Bespoke samples at 2.2 FID with 10 NFE, and gets within 2% of GT FID (1.71) with 20 NFE.
INTRODUCTION
Diffusson models (Sohl-Dickstein et al., 2015; Ho et al., 2020), and more generally flow-based models (Song et al., 2020b; Lipman et al., 2022: Albergo & Vanden-Eijnden, 2022), have become prominent in generation of images (Dhariwal & Nichol, 2021; Rombach et al., 2021), audio (Kong et al., 2020; Le al., 2023), and molecules (Kong et al., 2020). While training flow models is relatively scalable and efficient, sampling from a flow-based model entails solving a Stochastic or Ordinary Differential Equation (SDE/ODE) in high dimensions, tracing velocity field defined with the trained neural network. Using off-the-shelf solvers to approximate the solution of this ODE to a high precision requires a large Number (i.e., 100s) of Function Evaluations (NFE), making sampling one of the main standing challenges in flow models. Improving the sampling complexity of flow models, without degrading sample quality, will open up new applications that require fast sampling, and will help reducing the carbon footprint and deployment cost of these models.
Current approaches for efficient sampling of flow models divide into two main groups: (i) Distil- lation: where the pre-trained model is fine-tuned to predict either the final sampling (Luhman & Luhman, 2021) or some intermediate solution steps (Salimans & Ho, 2022) of the ODE. Distillation does not guarantec sampling from the pre- pre-trained model's distribution, but, when given access the training data during distillation training, it is shown to empirically generate samples of comparable quality to the original model (Salimans & Ho, 2022; Meng et al., 2023). Unfortunately, the GPU time required to distill a model is comparable to the training time of the original model Salimans & Ho (2022), which is often considerable. (ii) Dedicated solvers: where the specific structure of the ODE is used to design a more efficient solver (Song et al., 2020a; Lu et al., 2022a;b) and/or employ suitable solver family from the literature of numerical analysis (Zhang & Chen, 2022; Zhang al., 2023). The main benefit of this approach is two-fold: First, it is consistent, i.e., as the number of steps (NFE) increases, the samples converge to those of the pre-trained model. Second, it does not require further wwapiinn//ie-tuinnn. of the pre-trained model, consequently avoiding long additional training times and access to training data. Related to our approach, some works have tried to learn an ODE solver within a certain class (Watson et al., 2021; Duan et al., 2023); however, they do not guarantee consistency and usually introduce moderate improvements over generic dedicated solvers.

---

Published as a conference paper at ICLR 2024
In this paper, we introduce Bespoke solvers, framework for learn- ing consistent ODE solvers custom custom-tailored to pre-trained flow models. The main motivation for Bespoke solvers is that different models exhibit sampling paths with different characteristics, lead- ing to local truncation errors that specific to each instance of trained model. A key observation of this paper is that optimizing solver for a particular model can significantly improve quality of samples for low NFE compared to existing dedicated solvers. Fur- thermore, Bespoke solvers use a very small number of learnable parameters and consequently are efficient to train. For example, we have trained n (5, 8, 10] steps Bespoke solvers for pre- trained ImageNet-64 x 64 flow model with (40, 64, 80} learnable parameters (resp.) producing images with Fréchet Inception Dis- tances (FID) of 2.2, 1.79. 1.71 (resp.), where the latter is within 2% from the Ground Truth (GT) FID (1.68) computed with ~ 180 NFE. The Bespoke solvers were trained (using a rather naive im- plementation) for roughly 1 1% of the GPU time required for training the original model. Figure compares sampling at 10 NFE from a pre-trained AFHQ-256 256 flow model with order 2 Runge-Kutte (RK2) and its Bespoke version RKK2-Bes), along with the GT sam- ple that requires ~ 180 NFE. Our work brings the following contributions:
noise
Basekne
Bespok
samples
Figure 1: Using 10 NFE to sam- ple using our Bespoke solver im- proves fidelity w.r.t. the baseline (RK2) solver. Visualization of paths was done with the 2D PCA plane approximating the noise and end sample points.
1. A differentiable parametric family of consistent ODE solvers.
2. A tractable loss that bounds the global truncation error while allowing parallel computation.
3. An algorithm for training Bespoke n-step solver for specific pre-trained model.
4. Significant improvement over dedicated solvers in generation quality for low NFE,
BESPOKE SOLVERS
We consider a pre-trained flow model taking some prior distribution (noise) p to a target (data) dis- tribution g in data space Rdd The flow model (Chen al., 2018) represented by a time-dependent Vector Field (VF) u [0, 1] x Rd + Id that transforms a noise sample ~ p(2o) to data sample T1 ~ (x11 by solving the ODE
x(t) = un (x(t)),
(1)
with the initial condition z00) = zu p (z))), from time t until time t 1, and ((t) x: st))
The solution at time t = 1, i.e., xx1) ~ (=(()). is the generated target sample.
Numerical ODE solvers. Solving equation is done in practice with numerical ODE solvers. A numerical solver is defined by an update rule:
Algorithm Numerical ODE solver.
Require: fa for 0, 0, 22
(0,11,2,11) step 2; u.) end for return 2.,
(Inext; Tnea)) = step(t, u.).
(2)
The update rule takes as input current time and approximate so- lution æ, and outputs the next time step fnxx and the corresponding approximation Tnext to the true solution 2(trext) at time Enex- To approximate the solution at some desired end time, i.e., = 1, one first initializes the solution at = O and repeatedly applies the update rule in equation 2 n times, as presented in Algorithm 1. The step is designed so that fn = 1.
An ODE solver (step) is said to be of order i if its local truncation error is
Tnexill O ((tnext +)*+1),
(3)
asymptotically
fnexa where [0. 1) is arbitrary but fixed and fnexts Tnext are defined by the solver, equation 2. A popular family of solvers that offers a wide range of orders is the Runge-Kutta (RK) family (Iserles, 2009). Two of the most popular members of the RK family are (set h n ));
RKI (Euler order 1):
step(l, 1; u.) (1 h, hvve
(4)
RK2 (Midpoint- order 2): step(1,x;uun = + h, x + hu4 (
h (x)(x))
))
2
(5)
2

---

Published as a conference paper at ICLR 2024
Approach outline. Given pre-trained u, and target number of time steps n our goal is to find a custom (Bespoke) solver that is optimal for approximating the samples z defined via equation 1 from initial conditions sampled according to z=(0) = xo ~ p(ro)- To that end we develop two components: (i) differentiable para- metric family of update rules step". with parameters
Algorithm 2 Bespoke solver.
Require: fu. pre pre-trained me. for i 0,
do end (4e41: (441)))))))))) + step" ((.
return 2i
E 12
(where p very small), where sampling done by replacing step with step" in Algorithm 1, see Al- gorithm 2; and (ii) tractable loss bounding the global truncation error, i.e., the Root Mean Square Error (RMSE) between the approximate sample zf, and the GT sample x(1),
Global truncation error: CRMSE(() Bropp(ro) [[x(1) =
(6)
where z is the output of Algorithm 2, and ||ا|||| = ( E 1001)1111
2.1 PARAMETRIC FAMILYY OF ODE SOLVERS THROUGH TRANSFORMED SAMPLING PATHS
Our strategy for defining the parametric family of solvers step" is using a generic base ODE solver, such as RKI or RK2, applied to a parametric family of transformed paths.
Transformed sampling paths. We transform the sample trajectories a(t) by applying two components: a time reparametrization and an arbitrary invertible transformation. That is,
s(1)
P1
((r) (p)(t),)),, re r 1].
(7)
Br (2(rz)
where ty. (xr) are arbitrary functions in a family F defined by the [0. 1] diffeo- morphism! and yp [0, 1 x Rd - Rd is C and a diffeomorphism in x. We also assume ra and of are Lipschitz continuous with a con- stant L > O. (ii) Boundary conditions: i, satisfies lo = 0 and 4 1,
x(r:)
following conditions: (i) Smoothness; t, : =00,1]
(x(x()))
Pr.
(f)
x(0)
Figure 2: Transformed paths.
and (-)(-) is the identity function, i.e., po((x) = x for all æ E R& Fig-
ure 2 depicts a transformation of path, z((). Note that z(0) xx0). however the end point a(1) does not have to coincide with x (1). Furthermore, as fr [0, 1] - [o. 1] is a diffcomorphism, t, is strictly monotonically increasing.
The motivation behind the definition of the transformed trajectories is that it allows reconstructing ar(t) from Ø(r). Indeed, denoting r, the inverse function of t t, we have
(() प., (T(r)))
(8)
Our hope is find transformation that simplifies sampling paths and allows the base solver to provide better approximations of the GT samples. The transformed trajectory F, is defined by VF (x) that can be given an explicit form as follows (proof in Appendix A):
Proposition 2.1. Let x(1) be a solution to equation 1. a(r) defined in equation 7 is solution to the ODE (equation 1) with the VF
Denote 9p. == dr pr and it = h Then
d
(æ) = (,(( (x)) + (xp
(z))ur, (1p (1p7
(9)
We are now ready to define our parametric family of solvers step (4, x; ux): First we transform the input sample x) according to equation 7 to
Solvers via transformed paths.
(r, (re, (r,
(10)
Next, we perform a step with the base solver of choice,
(Frest: Tnext) = step(r, T;; ū.),
(11)
and lastly, transform back using equation 8 to define the parametric solver step" via
(Inext: Frex) step" (x. 1; (e)= (Ir PT (Fnext))
(12)
The parameters 0 denote the parameterized transformations f, and yr satisfying the properties of F and the choice of a base solver step. In Section 2.2 derive the explicit rules we use in this paper.
A diffcommrphism is a C¹ continuously differentiable function with C* continuous differentiable inverse.
3

---

Published as a conference paper at ICLR 2024
Consistency of solvers. An important property of the parametric solver step" is consistency. Namely, due to the properties of F, regardless the particular choice of tr, Pr E F, the solver step" has the same local truncation enror the base solver.
Theorem 2.2. (Consistency of parametric solvers) Given arbitrary t, 4y in the famity of functions F and a base ODE solver of order ki the corresponding ODE solver step" is also of order k, i.e.,
||x((nex))
Imeal || = (t(ene - t) =)
(13)
The proof is provided in Appendix B. Therefore, as long as 4r 9Pr (x(()) are in F, decreasing the base solver's step size h + 0 will result in our approximated sample an, converging to the exact sample x(1) of the trained model in the limit, i.e., x² + z x=1) as n oc.
2.2 TWO USE CASES
We instantiate the Bespoke solver framework for two cases of interest (a full derivation is in Ap- pendix E), and later prove that our choice of transformations in fact covers all "noise scheduler" configurations used in the standard diffusion model literature. In our use cases, we consider a time- dependent scaling as our invertible transformation PPr.
(x) S, and its inverse Φp (x) = 2//5r;
(14)
where [0, 1| + Ry0 is strictly positive C' scaling function such that 1 (i.e., satisfying the boundary condition of (p). The transformation of trajectories, i.e., equations 7 and 8, take the form
(r) = s,,,),,, and sr()) #(rr)/st,..
(15)
and we name this transformation: time The transformed VF ūr (equation 9) is thus
Sx + 8, ū (x)
(16)
8r
Use case I: RKI-Bespoke. We consider RKI (Euler) method (equation 4) the base solver step and denote r3 = ih, E [n], where [n] = {0, 1.....,] and h = n Substituting equation 4 in equation 11, we get from equation 12 that
step (tc, 2,;4;) u.) ti+1,
9i h5; ns + hi, St 81+1
un (z:)
(17)
8i41
where denote 1r.- 4, brr 8; Sr, 8 r-r, 8r, and E [n 1]. The learnable parameters 0 e IR" and their constraints are derived from the fact that the functions tr, are members of F. There are p 4n parameters in total: 0 (0*, 0*), where
g¹: 0 to t1 < tn < to =1
31
0
80
0²
(18)
in-1 Lio
80
Note that we ignore the Lipschitz constant constraints in F when deriving the constraints for A. Use case I: RK2-Bespoke. Here choose the RK2 (Midpoint) method (equation 5) the base solver step. Similarly to the above, substituting equation in equation 11, we get
step 0 (4,,2;; uz) = (4+1,
9i
h
6 1 + 4 W
(19)
8141
Si+1
8i+
8i+
where we set Ti+3 = Ti +
2, and accordingly 444g. din f. 8.43. and 8 are defined, and
+ h 2 sitiu (2,).
=
8 +
(20)
2
In this case there are = 8n learnable parameters, (0', 0*) TRP,, where
0¹: 0 = to t < ln
=
8 $1. 81,..., 8n
80 =1
(21)
ln. / , I 1 >
8n, ,
4

---

Published as a conference paper at ICLR 2024
Equivalence of scale-time transformations and Gaussian Paths. We note that our scale-time transformation covers all possible trajectories used by diffusion and flow models trained with Gaus- sian distributions. Denote by (z()) the probability density function of the random variable z(t), where stt) is defined by a random initial sampling x²0) = so ~ p( and solving the ODE in equation 1.
When training Diffusion or Flow Matching models, Pt has the form (r( = - where pr = N (xlauT1, of f). A pair of functions a, a: [0, 1] + [0, 1] satisfying
ag 0 01 a1 1 = do. and strictly monotonic sur(t) az/o,
(22)
is called scheduler". We use the term Gaussian Paths for the collection of probability paths (x( achieved by different schedulers. The velocity vector field that generates pr (2) and results from zero Diffusion/Flow Matching training loss is
Pr pe(alz1) dr
Ue (x( = ue (z(()))
(23)
pc(xr)
21, as derived in Lipman et al. (2022). Next, we generalize a result by Kingma et al. (2021) and Karras et al. (2022) to consider marginal sampling paths xx() defined by ux(z), and show that any two such paths are related by scale-time transformation:
where (a(())))) = de + (rr oz
at
oe
Theorem 2.3. (Equivalence of Gaussian Paths and sicale-time transformation) Consider Gaussian
Path defined by a scheduler (as, at). and let xxt) denote the solution of equation 1 with u, defined in equation 23 and initial condition xx()) Izg- Then,
(i) For every other Gaussian Path defined by scheduler (ār, σr) with trajectories Ø(r) there exists a scale-time transformation with $ = such that æ(r) = s,,((tr).
(ii) For every scale-time transformation with 81 = there exists a Gaussian Path defined by a scheduler Ox) with trajectories æ(r) such that s,.((t,) = a(r).
(Proof in Appendix C.) Assuming an ideal velocity field (equation 23), i.e., the pre-trained model optimal, this theorem implies that searching over the scale-time transformations equivalent to searching over all possible Gaussian Paths. Note, that in practice we allow 81 A 1, expanding beyond the standard space of Gaussian Paths. Another interesting consequence of Theorem 2.3 (simply plug in t 1) is that all ideal velocity fields in equation 23 define the same coupling, i.e., joint distribution, of noise zo and data ZT.
2.3 RMSE UPPER BOUND LOSS
Optimizing directly the RMSE loss (equation 6) is theoretically possible but would require keeping a full computational graph of Algorithm 2, i.e., nxorder compositions of u, leading to a large memory footprint. Therefore, we instead derive an upper-bound to the RMSE loss that enables parallel computation over the steps of the solver, considerably reducing memory consumption. To construct the bound, let us fix an initial condition Fy ~ p(ze) and denote before æ(1) to be the exact solution of the sample path (equation 1). Furthermore, consider candidate solver step", and denote its and x coordinate updates by step" = (step{)stepp)) Applying Algorithm 2 with to = 0, Fi produces series of approximations at each corresponds to time step f,, i E [n]. Lastly, we denote by
=1(()(()) - lI d (1(()) - step (4- (44-))))))))):):))) u.)l
(24)
the global and local truncation errors at time t, respectively. Our goal is to bound the global error at the final time step f. i.e., eg. Using the update rule definition (equation 2) and triangle inequality we can bound
e S | sta x (tt); un ++)++ |||ep= x((;;; un) - step aiuu)) < & + Lie..
where L is defined to be the Lipschitz constant of the function step, (t., ; u.). To simplify notation set by definition L 1 (this is possible since L does not actually participate in the bound).
We use the convention of noise time 0 and data at time
5

---

Published as a conference paper at ICLR 2024
Using the above bound times and noting that ej = 0 we get
Algorithm Bespoke training.
Require: pre-trained number of steps initialize / R'
en S Mfde, where M 3
(25)
while converged
p(uo)
sample noise GT path
xx1) <- solve ODE
Motivated by this bound we define our R.SE- -Bond loss:
init loss
parallel for 0.
do 5 step, (a (0.)
-+= 1 (f:41)
CRMSE-B ((0) pp(oo)
Σ Mºa.. 1=1
(26)
end for
VVL
optiiizzation step
end while return
where de is defined in equation 24 and Mf defined in equation 25. The constants L, depend both on the parameters 0 and the Lipschitz constant Ly of the network u2 As L,, is difficult to estimate, we treat Lu as a Hyyer-parameter denoted L, (in all experiments we use L, 1), and compute Lę in terms of 0 and L, for our two Bespoke solvers, RK1 and RK2, in Appendix D. Assuming that L, 2 Lu, an immediate consequence of the bound in equation 25 is that the RMSE-Bound loss bounds the RMSE loss, i... the global truncation error defined in equation 6, CRMSE (0) I CRMSE (0).
Implementation of the AMSE-Bound loss. We provide pseudocode for Bespoke training in Algo- rithm 3. During training, we need to have access to the GT path x=(t) at times f;, i E [n], which we compute with generic solver. The Bespoke loss is constructed by plugging step" (equations 17 or 19) into di (equation 24). The gradient Vo CRLSEE. (0) requires the derivatives (zx(1))0 Comput- ing the derivatives of sr14;) can be done using the ODE obeys, i.e., ((44) = dt, (2:4). Therefore, a simple way to write the loss ensuring correct gradients w.r.t. is replace zt,, with 2." (t.) where
2 2 (t) z ([t.]) "[+:] (=((t.])) (t [t,]]).
(27)
where [-] denotes the stop gradient operator; i.e., x; (t) is linear in and its value and deriva- tive w.r.t. t coincide with that of (t,)) time t = t,. Full details are provided in Appendix F. In Appendix K.1 we provide an ablation experiment comparing different Bespoke losses and corre- sponding algorithms including the direct RMSE loss (eq. 6) and our
loss (eq. 26).
3 PREVIOUS WORK
Diffusion models (Sohl- Dickksien et al., 2015; Ho al.. 2020) are a powerful paradigm for gen- erative models that for sampling require solving a Stochastic Differential Equation (SDE), or its associated ODE, describing a (deterministic) flow process (Song et al., 2020a). Diffusion models have been generalized to paradigms directly aiming to learn a deterministic flow (Lipman et al., 2022; Albergo & Eajden.,ijnden, 2022; Liu et al., 2022). Flow-based models are efficient to train but costly to sample. Previous works had tackled the sample complexity of flow models by building dedicated solver schemes and distillation.
Dedicated Solvers. This line of works introduced specialized ODE solvers exploiting the structure of the sampling ODE. Lu et al. (2022a); Zhang & Chen (2022) utilize the semi-linear structure of the score/e-based sampling ODE to adopt method of exponential integrators. (Zhang et al., 2023) further introduced refined error conditions to fulfill desired order conditions and achieve better sam- pling, while Lu et al. (2022b) adapted the method guided sampling. Karras et al. (2022) suggested transforming the ODE to sample a different Gaussian Path for more efficient sampling, while also suggesting non-uniform time steps. In principle, all of these methods effectively proposed-based on intuition and heuristics to apply a particular scale-time transformation to the sampling trajec- tories of the pre-trained model for more efficient sampling, while Bespoke solvers search for an optimal transformation within the entire space of scale-time transformations.
Other works also aimed at learning the solver: Dockhorn et al. (2022) (GENIE) introduced a higher- order solver, and distilled the necessary JVP for their method; Watson et al. (2021) (DDSS) opti- mized a perceptual loss considering a family of generalized Gaussian diffusion models; Lam et al. (2021) improved the denoising process using bilateral filters, thereby indirectly affecting the effi- ciency of the ODE solver; Duan et al. (2023) suggested to learn a solver for diffusion models by replacing every other function evaluation by a linear subspace projection. Our Bespoke Solvers be- long to this family of learnt solvers, however, they are consistent by construction (Theorem 2.2) and minimize bound on the solution emror (for the appropriate Lipschitz constant parameter).

---

Published as a conference paper at ICLR 2024
Distillation. Distillation techniques aim to simplify sampling from trained model by fine-tuning or training a new model to produce samples with fewer function evaluations. Luhman & Luhman (2021) directly regressed the trained model's samples, while Salimans & Ho (2022); Meng et al. (2023) built a sequence of models each reducing the sampling complexity by a factor of 2. Song al. (2023) distilled consistency map that enables large time steps in the probability flow; Liu al. (2022) retrained a flow-based method based on samples from a previously trained flow. Yang et al. (2023) used distillation reduce model size while maintaining the quality of the generated images. The main draw banks of distillation methods is their long training time (Salimans & Ho, 2022), and lack of consistency, i.e., they do not sample from the distribution of the pre-trained model.
RK]
4 EXPERIMENTS
AK2
Models and datasets. Our method works with pre-trained models: we use the pre-trained CIFARIO (Krizhevsky & Hin- ton, 2009) model of (Song et al., 2020b) with published weights from EDM (Kar- ras et al., 2022). Additionally, we trained diffusion/flow models on the datasets: CI- FAR10, AFHQ-256 (Choi et al., 2020a) and ImageNet-664/128 (Deng et al., 2009). Specifically, for ImageNet, as recommended by the authors (ima) we used the official face-blurred data x64x64 downsampled using the open source preprocessing scripts from Chrabaszez et al. (2017)). For dif- fusion models, we used an e- Variance Preserving (c-VP) parameterization and schedule (Ho et al., 2020: Song et al., 2020b). For flow models, used Flow Matching (Lipman et al., 2022) with Conditional Optimal Transport (FM-OT), and Flow Matching/e-predietiin with Cosine Scheduling (FM/v-CS) (Salimans & Ho, 2022; Albergo & Vanden-Fijjneen 2022). Note that Flow Matching methods directly provide the velocity vector field (x),, and we converted c-VP to a velocity field using the identity in Song et al. (2020b). For conditional sampling we apply classifier free guid- ance (Ho & Salimans, 2022), so each evaluation uses two forward passes.
NFE
NFE
Figure 3: Bespoke RK1/2, Figure 4: Bespoke solver
ImageNet-64 FM-OT.
applied to EDM's (Karras al., 2022) CIFARIO pub- lished model.
Bespoke
and optimization. As our base ODE solvers, we tested RKI (Euler) and RK2 (Midpoint). Furthermore, we have two hyper-parameters n the number
Method
NPE FID
Divillation
Zheng el al. (2023)
3.78
Lahman & Luhman (2021)
9.36
of steps, and L,, the Lipschitz constant from lemmas D.2,
Salimans (2022)
D.3. We train our models with n [4, 5, 8, 10, 12) steps and fix L,, = 1. Ground Truth (GT) sample trajectories, ar(t4),
2
4.51
DDIM et al., 2020a)
10
13.36
are computed with an adaptive RK45 solver (Shampine, 1986). We compute FID (Heusel et al., 2017) and valida- tion RMSE (equation 6) is computed on a set of 10K fresh example
6.84
DPM al., 2022a)
20
3.99
noise samples To ~ p(zu); Figure 12 depicts of RMSE vs. training iterations for different n values. Un- less otherwise stated, below iteration and show samples on best RMSE validation itera- tion. Figures 21, 22, 23 depict the learned Bespoke solvers' parameters 0 for the experiments presented below; note the differences across the learned schemes for different models and datasets.
Dedicated solvers
DEIS (Zhang & Chen, 2022)
20
10
2.86
4.17
GENIE (Doekhorn al., 2022)
10
5.28
report results best FID
20
3.94
DDSS (Watson al., 2021)
10
7.86
20
4.72
RK2-BES
c-Vp
10
3.31
2.75
RK2-RES
FM/+CCS
10
2.89
FM/e-CS
20
2.64
RK2-BES
FM-OT
10
2.73
FM-OT
20
2.59
Bespoke RK1 vs. RK2. We compared RKI and RK2 and their Bespoke versions on CIFAR10 and ImageNet-6i mod- els (FM-OT and FM/v-CS). Figure 3 and Figures 10, show best validation RMSE (and corre- sponding PSNR). Using the same budget of function evaluations RK2/RK2-Bespok. produce con- siderably lower RMSE validation compared to RK1/RK/ I- -Bespokee respectively. We therefore opted for RK2/RK2-Bespook for the rest of the experiments below.
Table 1: CIFARIO sampling.
CIFAR10. We tested our method on the pre-trained CIFAR1O e- c-VP model (Song et al., 2020b) released by EDM (Karras et al., 2022). Figure 4 compares our RK2-Bespoke solver to the EDM method, which corresponds to a particular choice of scaling, s;, and time step discretization, f, Fuler and EDM curves computed as originally implemented in EDM, where the latter achieves
7

---

Published as a conference paper at ICLR 2024
ImageNet-64: e-pred
ImageNee-64: FMre-CS
ImageNet-64: FM-OT
ImageNet-125 FM-OT
14:00
10:00
3.0
NFE
NFE
NFE
NFE
700
MFE
NFE
NFE
NFE
Figure 5: Bespoke RK2 solvers vs. RK1/2/4 solvers on CIFAR-10 ImageNet-64, and Image-Net128. FID
vs. NEE (top row), and RMSE vs. NEE (bottom row). PSNR vs. NFE is shown in Figure 13.
FID=3.05 at 35 NFE, comparable to the result reported by EDM. Using our RK2- Solver, achieved an FID of 2.99 with 20 NFE, providing 42% reduction in NFE. Additionally, we tested our method on three models we trained ourselves on CIFAR1O, namely c- c-VP, FM/r-CS, and FM-OT. Table compares our best FID for each model with different baselines demonstrating superior generation quality for low NFE among all dedicated solvers; e,g. for NFE=10 we improve the FID of the runner-up by over 34% (from 4. to 2.73) using RK2-Bespoke FM-OT model. Table 4 lists best FID values for different NFE, along with the GT FID for the model and the fraction of time Bespoke training took compared to the original model's training time; with 20 NFE, our RK2-Bespoke solvers achieved FID within 8%, 1%, 1% (resp.) of the GT solvers' FID. Although close, our Bespoke solver does not match distillation's 's performance, however our approach is much faster train, requiring ~ % of the original GPU training time with our naive implementation that re-samples the model at each iteration. Figure 11 shows FID/MMSE/PSNR vs. NFE, where PSNR computed w.r.t. the GT solver's samples.
ImageNet 64/128. We further experimented with the more challenging ImageNet-64 / 128 x 128 datasets. For ImageNet-64 we also trained 3 models as described above. For ImageNet- 128, due to com- putational budget constraints, we only trained FM-OT (training requires nearly 2000 GPU days). Figure 5 compares RK2-Bespoke to various baselines includ- ing DPM 2nd order (Lu et al., 2022a). As can be seen in the graphs, the Bespoke solvers improve both FID and RMSE. Interestingly, the Bespoke sampling takes all methods to similar RMSE levels, fact that can be partially explained by Theorem 2.3. In Ta- ble 2, similar Table we report best FID per NFE for the Bespoke solvers we trained, the GT FID of the model, the % from GT achieved by the Bespoke solver, and the fraction of GPU time (in %) it took train this Bespoke solver compared to training the original pre-trained model. Lastly, Figures 6, 7, 27, 28, 29, 25, 26 depict qualitative sampling examples for RK2-Bespoke and RK2 solvers. Note the signifi- cant improvement of fidelity in the Bespoke samples to the ground truth.
ImageNet-64
NFE
FID
GT-FID/% TTime
RK2-HES
c-VP
8
3.63
1.83
229
3.5
c-VP
10
2.96
163
3.6
<-VP
16
2.14
120
3.6
c-VP
20
1.93
109
3.5
c-PP
24
1.84
101
3.6
RK2-BES
FM/e-CS
1.68
FM/o-CS
10
2.20
131
1.6
FM/e-CS
16
1.79
107
1.8
FMRe-SS
1.71
102
FM/o-CS
24
1.69
101
2.0
RK2-HES
FM-OT
3.10
1.68
185
1.6
FM-OT
10
2.26
135
FM-OT
16
1.84
110
1.7
FM-OT
20
1.77
105
1.7
FM
ImageNNt-122 128
NFE
FID
GT-FID/%
% Time
RK2-BES FM-OT
8
5.28
2.30
230
1.1
FM- OT
FM-OT
16
2.64
115
1.2
FM-OT
20
2.45
107
1.2
FM.OT
Table 2: ImageNet Bespoke solvers.
AFHQ-256. We tested our method on the AFHQ dataset (Choi et al., 2020b) resized to 256x256 where as pre-trained model we used a FM-OT model we trained as described above. Figure 14 depicts PSNR/RMSE curves for the RK2-Bespoke solvers and baselines, and Figures 7 and 24 show qualitative sampling examples for RK2-Bespoke and RK2 solvers. Notice the high fidelity of the Bespoke generation samples.
8

---

Published as a conference paper at ICLR 2024
GT
NFE=10 NFE=20
NFE=8
GT
NFE=20
NFE=10
NFE-
NFE-20
NPE=10
NFE=8
RK2
FM-0TT
RK2-BES
RK2
FMMe-CS
RK2-BER BESS
Figure 6: Comparison of FM-OT and FM/u-CS ImageNet-64 samples with RK2 and bespoke-RK2 solvers. Comparison DPM-2 samples are in Figure 30. More examples are in Figures 27, 28, and 29, The similarity of generated images across models be explained by their identical noise data coupling (Theorem 2.3).
GT NFE=20
NFE=10
NEE=
NFE-20
NFE=8
RK2
RK2-BES
ImogeNer- 128
RK2
RK2
AFIQ 256
Figure FM-OT ImageNct- 128 (top) and AFHQ-256 (bottom) samples with RK2 and bespoke- solvers.
More examples are Figures 25, and 24.
Ablations. We conducted two ablation experiments. First, Figure 16 shows the effect of training only time transform (keeping 8, = 1) and scale transformation (keeping f, r). While the time transform more significant than scale transform, incorporating scale improves RMSE for low NFE (which aligns with Theorem 2.2), and improve FID. Second, Figure 18 shows application of RK2-Bespoke solver trained on ImageNet-64 applied to ImageNet-128. The transferred solver, while sub-optimal compared to the Bespoke solver, still considerably improves the RK2 baseline in RMSE, and improves FID for higher NFE (16,20). Reusing Bespoke solvers can potentially be a cheap option to improve solvers.
5 CONCLUSIONS, LIMITATIONS AND FUTURE WORK
This paper develops an algorithm for finding low-NFE ODE solvers custom-tailored to general pre- trained flow models. Through extensive experiments we found that different models can benefit greatly from their own optimized solvers in terms of global truncation error (RMSE) and generation quality (FID). Currently, training Bespoke solver requires roughly 1% of the original model's training time, which can probably be still be made more efficient, e.g., by using training data or pre-pro processing sampling paths. A limitation of our framework is that requires separate training for each target NFE and/or choice of guidance weight. For general NFE solvers one may consider more
combined loss and/or continuous representation of fr,tr, while guidance weight general conditions can be used to directly condition fr. More general/expressive models for Pr, have the potential further improve fast sampling of pre-trained models.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
NS is supported by grant from Israel CHE Program for Data Science Research Centers.
REFERENCES
Imagenet website. https: /twp:///wwwwwww http::/://ww..nggg-net.orrr////
Michael S. Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic inter- polants, 2022.
Ricky Q. Chen. torchdi ffeq.
torchdiffeq, 2018. URL
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. Stargan v2: Diverse image synthesis for multiple domains, 2020a.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung- Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2020b.
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. downsampled variant of imagenet as an alternative to the cifar datasets. arXiv preprint arXiv:1707.08819 2017.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei Li. Imagenet: A large-scale hi- erarchical image database. In The IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 11:8780-8794 2021.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers. Advances in Neural Information Processing Systems, 35:30150-30166, 2022.
Zhongjie Duan, Chengyu Wang, Cen Chen, Jun Huang, and Weining Qian. Optimal linear subspace search: Learning to construct fast and high-quality schedulers for diffusion models. arXiv preprint 1wX7 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Jonathan Ho and Tim Salimans. 2rXiv:2207.12590 2022.
Classifier-free
diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851 2020.
Arich Iserles. A first course in the numerical analysis of differential equations. Number 44, Cam- bridge university press, 2009.
Tero Karras, Miik 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.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Ad- vances in neural information processing systems, 34:21696-21707, 2021.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
10

---

Published as a conference paper at ICLR 2024
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint (0Xi::2009.0966,, 2020.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. In University of Toronto, Canada, 2009.
Max WY Lam, Jun Wang, Rongjie Huang, Dan Su, and Dong Yu. Bilateral denoising diffusion models. XrXiv preprint arXiv:2108. 11514, 2021.
Matthew Le, Apoorv Vyas, Bowen Shi, Brian Karrer, Leda Sari, Rashel Moritz, Mary Williamson, Vimal Manohar, Yossi Adi, Jay Mahadeokar, al. Voicebox: Text-guided multilingual universal speech generation at scale. arXiv preprini arXiv:2306. 15602 2023.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:22/0.02777, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 00000,,,,0 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 1155775-5787 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095, 2022b.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02338 02388, 2021.
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models, 2021.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.0053 00 2022.
Lawrence F Shampine. Some practical runge-kutta formulas. Mathematics of computation, 46(173); 135-150, 1986.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2265. PMLR, 2015.
Jiaming Song. Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. preprint wwXiv:2010002502. 2020a.
arXiv
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020b.
Yang Song. Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023.
Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning fast samplers for diffusion models by differentiating through sample quality. In International Conference on Learning Representations, 2021.
Xingyi Yang, Daquan Zhou, Jiashi Feng, and Xinchao Wang. Diffusion probabilistic model made slim. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22562, 2023.
11

---

Published as a conference paper at ICLR 2024
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:220..13900 2022.
Qinsheng Zhang, Jiaming Song, and Yongxin Chen. Improved order analysis and design of expo- nential integrator for diffusion models sampling. arXiv preprint vrXiv:2308.021555 2023.
Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference on Machine Learning, pp. 42390 42402. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2024
A TRANSFORMED PATHS
(Appendix to Section 2.1.)
Proposition A.1. Let xxt) be solution equation Denote &. d
dr Sr and f, tr. Then
ā(r) defined equation is a solution to the ODE (equation 1) with the VF
iy (x) (, i,0 (1 (9, (x)).
Proof. Differentiating æ(r) in equation 7, i.e. Z(r) (x((tt)) w.r.t r and using the chain rule gives
d
a(r) dr (pr(a(t.)))
(,(((t,) + 0 -)(x((()) + (m, = ( + 0, (4tp:///hh))000p/p/////////////)
where the third equality we used the fact that s²() solves the ODE in equation and therefore a(t) (x()))));; and in the last equality we applied op to both sides of equation 7, i.e. (t,,) = fr (i(r).. The above equation shows that
k() up (F(r))
(28)
where (2:) is defined in equation 9, required.
B CONSISTENCY OF SOLVERS
(Appendix Section 2.1.)
(Inext; Freext
(ttnextt)
(Tnext Fnext
4 Preext
Te (Tuext))
Lr
(1,
(r,
Figure 8: Proof notations and setup.
Theorem 2.2. (Consistency of parametric solvers) Given arbitrary by Pr the family of functions F and a base ODE solver of order k. the corresponding ODE solver step" is also of order k, i.e.,
||x((nex))
Eneul || = O((tneu t) k+11)
(13)
Tnext + h and fnext fr Furthermore, by definition ã (xr) is sample at time rı #F(nest) the solution to the ODE defined by ü, starting from (r; #next is approximation to @(Tnext) as generated from the base ODE solver step. Lastly, Inext = = (Xraa) and (Innex)) = 97. (#(Tnext)) See Figure for illustration visualizing this setup.
Proof. Here (f, is our input sample E d at time E [0, 1]. By definition
Now, since step of order we have that
3 (neeex) Fnext #(Foext) ssteptt, r: ūr))
O((reet r) h+1).
(29)
13

---

Published as a conference paper at ICLR 2024
Now,
x(Inex)
Tnext
x((nex))
Prane (Fnext)
x((nex))
rrmm (#(reet) + ((TTexx - -))
nttnee) Prme (5 (F(reaa)) O 17) 1)
((
r)ht1)
O((nee t)(+t),
where in the first equality we used the definition of Tnext: in the second equality we used equation 29; in the third equality we used the fact that p is Lipschitz with constant L (for all r); the fourth equality we used the definition of the path transform, ar(eee)) & (#((nesx) mentioned above; and in the last equality we used the fact that rg is also Lipschitz with a constant T. and therefore
Tnext
Tt Otosst t).
C EQUIVALENCE OF GAUSSIAN PATHS AND SCALE-TIME TRANSFORMATIONS
(Appendix Section 2.2.)
Theorem 2.3. (Equivalence of Gaussian Paths and scale-time transformation) Consider a Gaussian Path defined by a scheduler (as, ot), and let æ(t) denote the solution of equation I with us defined in equation 23 and initial condition z²0)= = To- Then,
(i) For every other Gaussian Path defined by a scheduler (ār, (ā),ū,) with trajectories ø(r) there exists a acale-time transformation with = 1 such that Ø(r) s,x(tr).
(ii) For every scale-time transformation with = 1 there exists a Gaussian Path defined by scheduler (ār, σ.) with trajectories æ(r) such that s,,((t,) ((r).
Proof of theorem 2.3. Consider arbitrary schedulers (z) and (ā,, We can find sr,t such that
o S,00,
or St
(30)
Indeed, one can check the following are such 8 ty:
or 8y
tp = smr
(31)
where we remember snr is strictly monotonic defined in equation 22, hence invertible. On the other hand, given an arbitrary scheduler (ar, and an arbitrary scale-time transformation (4,, 8r) with $1 1, we can define scheduler G,) via equation 30.
For case (i), we are given another scheduler ā,, 3. and define a scale-ting time transformation Sy, fy with
equation 31. For case (ii), we are given scale-time transformation s,, and define scheduler Gr, G. by equation 30.
Now, the scheduler ā,, o, defines sampling paths a(r) given by the solution of the ODE in equa- tion with the marginal VF (1)) (z) defined in equation 23, i.e.,
(f)))=) ūr (x/(21) (((x)))((())) dx
(32)
(x)
where (x(Ø1) å=x j0 ã, F1. 0,
The scale-time transformation Sr, t, gives rise to second VF ū (I) as in equation 16,
š, x + S,
(33)
Sr
where is the VF defined by the scheduler (or, σI) and equation 23.
14

---

Published as a conference paper at ICLR 2024
By uniqueness of ODE solutions, the theorem will be proved if we show that
x u, (z),
Var R E [0, 1].
(34)
For that end, use the notation of determinants express
1
x
3
ú, (x|(21) = 0r Or 5, Or
(35)
Or
where x; x1 Rd and (r, or, âr, o, E R as in vector cross product. Differentiating Or, Or w.r.t. gives
à = at + 8, lrs - b,, + 8,
(36)
Using the bi-linearyty of determinants shows that:
0
1
u, (2)21)
1
Or
1
0
I
Sr
St
1
8,0t,
01 + S,
s, + spart tr
0
D1
0
21
8 S,
1
+
Sp Spre,tr
8,0t,
1
t,0tr
Bt 3,
0
S,
Sp
är x +
Sptr
3r
Otr
Ctr
1
Ötr
Ott
0
är + syt, Ute
x
2
Sr
where in the second equality we substitute ö,, o, as in equation 36, in the third and fourth equality we used the bi- bi-linearity of determinants, and in the last equality we used the definition of un (rr(ar))
der ae
as oe
xı expressed determinants notation. Furthermore, since
(x(x)))
S,00 x
2 01, = Pt. 8,
I1
(37)
have that
(x(2)2)) Pt.
I Sy
(38)
Therefore,
ʃ ür (x)221) Pr ( (o((()))((m)) (x)
dri = (x ur
Sr
E p.. (#i 2))
x + SptpUle
8r
8r
$ r + s,t pe (2r + +)
x
uer
8p
8, + s,t, us
where in the first equality we used Bayes rule, in the second equality we substitute ū, (x)2))) and (=(l) as above, and the last equality we used the definition of d as in equation 23. We have proved equation 34 and that concludes the proof.
15

---

Published as a conference paper at ICLR 2024
D LIPSCHITZ CONSTANTS OF STEP".
(Appendix to Section 2.3.)
We are interested in computing L. a Lipschitz constant of the bespoke solver step function (t., Namely, LQ should satisfy
||step (4;,;;:;)) - step y: u.) L || wll Var, E R
(39)
We remember that step, by computing Lipschitz constant for Uy, denoted (r;) in auxiliary lemma:
( (44, ;)) is defined using base solver and the VF ū,;()): hence, we begin
Lemma D.1. Assume that the original velocity field u, has a Lipschitz constant L,, 0. Then for every ry E [0, 11, Lr 2 Lu, and x, y E R=
[r (x) Tr (1)l I La(ri) - -lll
(40)
where
La(r;) 1341 igLr
(41)
$i
Proof of lemma D.I. Since the original velocity field w has a Lipshitz constant Lu > 0, for every t [0, and x, E Rd
[1)x(x) ue(2))l - Lu 1 yll
(42)
Hence
81
8;
y/
[rr, (1) the (v)ll
x. + 4uu0
y
(43)
8i
8;
8;
81
$4
(x:
+
uta
y
(44)
81
8;
81
x
||3
yl
Ut
Uta
(45)
8:
8i
8:
i
||:0
yl
(46)
8;
i
|
(47)
8i
We first apply the auxiliary lemma D.1 compute Lipschitz constant of step, (t,, u.) with RK1 (Euler method) as the base solver in lemma D.2 and for RK2 (Midpoint method) as the base solver in lemma D.3.
Lemma D.2. (RKI Lipschitz constant) Assume that the original velocity field un has Lipschitz constant T., > 0. Then, for every Top 2 T.w.
L
34 (1+ hLah)arr))),
(48)
8i+1
is a Lipschitz constant of RKI 1Kl-Bespoke update rule, where
La(r:) = |8:1 + L
(49)
Proof of lemma D.2. We begin with writing explicit expression of step. (ts, 2;;44 for Euler solver in terms of the transformed velocity field ū,. That is.
1
step (14, F2;
[spo
(8.2))
(50)
8:41
16

---

Published as a conference paper at ICLR 2024
So that applying the triangle inequality and lemma D.1 gives
1
|step_ 24 x; u4) step'. (1i, M; un)l
||s,ar + butr; (1127)
18sy + hittr (8:y)]
8;41
81
||
yl +
h
| (11ar) they (8:1))l
8;+1
8i+1
8i
h
[88l + tilerr ||stet [|stet styll
||
yl
+
8;+1
8
8i
+ t(iler) ) 1| 2ll.
1 + h
8i+1
94
Lemma D.3. (RK2 Lipschitz constant) Assume that the original velocity field we has Lipschitz constant L, > Then for every L, 2 Lu
h 8i 1 (+a(((ret)) (1 La(re)) )
1
(51)
81+1
a Lipschitz constant of RK2-RK2BRRe -e2Besk update rule, where
[i
La(r:)
+ feL..
(52)
8;
Proof of lemma D.3. We begin by writing explicit expression of step (t;, x;; uz) for RK2 (Midpoint)
method in terms the transformed velocity field ü.. We set
= 5;2 + h Ur (s(z), w s.y + 2
ür (s:t).
(53)
2
then
step, . (1,,,
1
3t + hūū (z)
(54)
8,+1
and
1
step (14, y5 up)
s.y + hur (10)
(55)
8;-1
So that applying the triangle inequality and lemma D.1 gives [ttep 2; u.) step", (t;, y: un]l
|| yll +
9 | (2) 3r (w)l
91+1
8i
|| - yll +
La (a((++))) }) || w|l
(56)
8i+1
8
We apply the triangle inequality and the IemmaD. again to || - w|l. That is.
h 8,2 +
h ür; (sw)
2| w|
uy (s:x)) sty +
2
< 8t ||
yll
+
||ū, (8(x) ūr (s.y))
2
S 8: ||
wll +
La
Ln(r)));
||:0 wll
2
1 h
La(ri)
|| vll
2
Substitute back in equation 56 gives
h
(1+ (r:))) ) || vll
|
yll
+
L, (Tee
||2 w|l <
1 + hL (F+4g)
8,41
Si+1
8:41
17

---

Published as a conference paper at ICLR 2024
E DERIVATION OF PARAMETRIC SOLVER STEP*
(Appendix Section 2.2.)
This section presents a derivation of n-step parametric solver
step" (t,r;u)) = (step? (t, ;;; st). (t,
(57)
for scale-time transformation (equation 15) with two options for a base solver: (i) RK1 method (Fuler) as the base solver; and (ii) RK2 method (Midpoint). We do so by following equation 10- We begin with RK1 and derive equation 17. Given (t), equation 10 for the scale time transformation is,
ã = 9:
(58)
Then according to equation 11,
Iz+1 step,, (ri, Fis Ur:))
(59)
Ii hūr, (T:)
(60)
$1 - List
Z
I; + h
(61)
8:
s,3; + h (á,2, + iissue (x+)).
(62)
where in the second equality we apply an RK1 step (equation 4), in the third equality we substitute Ur, using equation 16, and in the fourth equality substitute T; as in equation 58. According to RKI step (equation 4) we also have r++1 = r, + h. Finally, equation 12 gives,
=4+1
(63)
8: + h8s 23 + i,su (zi),
step (ti, ;1))
(64)
81t1
8:41
in equation 17.
Regarding the second case, equation 11 for the RK2 method (equation 5) is,
H141 step, (T's Tyg Upx)) I. + hūrr | (1++g)
(65)
(66)
where
E Hz + ür. (5.))
(67)
the RKI step from (r,, R.) with step size h/2. Now substituting ã; as defined equation 58 and i, defined in equation 16 get
$ B 1 R,+ + E 15 we
T++1 = 3,2, +
(68)
8143
8
where
8 =
8i + 3. xs + (sdiuur
(69)
Lastly, according to equation 12 we have
step 2 (t., =
8;
xx; +
h
Sith Zith
+ E
814 184 uett
(70)
si+1
8i41
84+
8
in equation 19 where 2, F
18

---

Published as a conference paper at ICLR 2024
F IMPLEMENTATION DETAILS
(Appendix to Section 2.3.)
This section presents further implementation details, complementing the main text. Our parametric family of solvers step" is defined via base solver step and transformation (l,, (r) as defined in equation 12. We consider the RK2 (Midpoint, equation 5) method as the base solver with n steps and (tr, (r)) the scale-time time transformation (equation 15). That is, (x) =, where [0, 1] R>o, as in equation 14, which is our primary use case.
Parameterization of t.. Remember that t, is strictly monotonic, differentiable, increasing func- tion :: [0, 1] + [0, 1]. Hence, & must satisfy the constraints as in equation 21, i.e.,
= to t
< t
(71)
to: $
nn >0. 0.
(72)
To satisfy these constrains, we model t, and i via
= L
i, = =
(73)
where 0f and 0f, i = 0, 2....,, are free learnable parameters.
Parameterization of sy. Since s, is strictly positive, differentiable function satisfying a boundary condition at = 0, the sequence s; should satisfy the constraints as in equation 21, i.e.
$1.
Sp
1,
(74)
and à are unconstrained. Similar to the above, we model 8; and & by
0
i=
8: exp Of otherwise
8; 0g
(75)
where 0 and 07, 0, 2: ..... are free learnable parameters.
Bespoke training. The pseudo-code for training a Bespoke solver is provided in Algorithm 3. Here we add some more details on different steps of the training algorithm. We initialize the parameters 9 such that the scale-transfommen transformationn is the Identity transformation. That is, for every i = 0, 1
2 ... n,
t i
i ==
(76)
8 = 1, 85 = O.
(77)
Explicitly, in terms of the learnable parameters, for every = 0, 2, 2, n,
0 = 1, 0
(78)
0 = 0, 0 0.
(79)
To compute the GT path ztt), we solve the ODE in equation 1 with the pre-trained model uy and DOPRI5 method, then use linear interpolation to extract att+),
1,.,, n (Chen, 2018). Then, apply (t) (equation 27) to correctly handle the gradients w.r.t. to 0f. See Table 3 for number of trajectories used during training. To compute the loss LRMSEB (equation 26) we compute 2;+1 = step? (4:), 4; ue) with equations 19,20, and compute M, via lemma D.3 with L, = 1. Finally, we use Adam optimizer Kingma Ba (2017) with learning rate of -3
Efficient sampling. When sampling using a Bespoke solver (Algorithm 2) each step involves applying क. and Pr; consecutively. In case we use scale transformation, equation 14 (as is done in all examples in this paper), this does not introduce any difficulty, however if a more compute intensive is used the following sampling pseudo-code (Algorithm 4) provides an equivalent sampling while avoiding this unnecessary step.
19

---

Published as a conference paper at ICLR 2024
CIFARIO
ImageNet-64
ImageNet- 128
AFHQ 256
Total number of trajectories
72k
48k
48k
4k
Batch size Number of iterations
12
8
6k
6k
6k
4k
Table Hyper-parameters of Bespoke solvers training on CIAA4AU0mmuggeeee////gg///////40 256.
Algorithm 4 Bespoke sampling (efficient).
Require: pre-trained Uye trained
20 p(ro)
sample noise
To - 0, To f To >> initial conditions
for 0, l.....n do
(ri+1. T++1) t step(r;, T.; u
end for
return 101 (in)
G BESPOKE RK1 VERSUS RK2
ImageNet 64: FM/v-CS
ImageNet 64: FM-OT
RK1
BES
RKI
BES
RK2
Base
RX2
Base
RMSE
RMSE
60
20
30
10
12
1.4
16
16
10
1,2
14
NFE
NFE
25
PSNR
PSNR
20
15
10
12
1.4
36
16
1.0
12
34
16
NFE
NFE
Figure Bespoke RK1, Bespoke RK2, RK1, and RK2 solvers on ImageNet-64 models: RMSE vs. NFE (top row), and PSNR vs. NFE (bottom row).
20

---

Published as a conference paper at ICLR 2024
CI e-VP
CIFARIO: FM/U-CS
CIFARIO: FM-OT
15
RK1
BES
RK1
BES
RKI
BES
RK2
Base
RK2
Basc
RK2
Basc
RMSE
RMSE
RMSE
10
25
35
PSNR
PSNR
PSNR
NFE
Figure 10: Bespoke RK1, Bespoke RK2, RKI, and RK2 solvers CIFARIO: RMSE vs. NFE (top row), and PSNR vs. NEE (bottom row).
H CIFAR10
CIFARIO
NFE
GT-FID/%
%Timm
RK2-BES
c-PP
8
4,26
2.54
168
1.4
c-VP
10
3.31
130
1.5
c-VP
16
2.84
112
1.5
e-VP
20
2.75
108
1.4
RK2-BES
FM/e-CS
3,50
2.61
0.5
FM/e-CS
FM/e-CS
103
FM/e-CS
20
2.64
101
0.6
RK2-BES
FM-OT
8
3.13
2.57
122
0.5
FM-OT
10
2.73
106
0.6
FM-OT
16
2.60
101
0..
FM-OT
20
2.59
101
0.6
Table 4: CIFARIO Bespoke solvers. We report best FID NFE, the ground truth FID (GT-FID) for the model FID/GT-FID in and fraction of GPU time %) required train bespoke solver w.r.t training the original model.
21

---

Published as a conference paper at ICLR 2024
CIFARIO: e-VP
CIFARIO: FM/e-CS
CIFARIO: FM-OT
RK4
RKAI
RXA
RK1
50.6
RK1
RK1
RK2-BES
RK2-BES
RK2-8ES
50.3
10.0
20.3
5.0
15
12x
NFE
RMSE
RMSE
RMSE
25
20
NFE
NFE
15
PSNR
PSNR
PSNR
10
12
34
26
10
12
14 NFE
36
20
18
20
NFE
NFE
Figure 11: CIFARIO sampling with Bespoke RK2 solvers vs. RK1,RK2,RK4; FID vs. NEE (top row), RMSE vs. NFE (middle row), and PSNR vs. NEE (bottom row).
T IMAGENET-64/122
ImageNet-64: e-pred
ImageNet-64: FM/e-CS
ImageNet-64: FM-OT
ImageNet-1288 FM-OT
MrE-B
WrE-B
WFE-B
NrE-0
MIE-10
MIE-15
WIE-12
ME--10
NFE-16
MEE116
MEE-16
NFE-16
NDE-20
MF-
MFE
NFE-20
NFE-24
MFE-24
MFE-24
1300
Iberation
600
100
Iteratioo
600
100
Iteratioo
2000 1000
Figure Validation RMSE vs. training iterations Bespoke RK2 solvers on ImageNet and ImageNet- 128.
ImageNet-64: e-pred
ImagoNet-64:: FM/a-CS
ImageNet-64: FM-OT
ImageNet-1228 FM-OT
RC4
RC eEs
Figure 13: Bespoke RK2, RK1, RK2, and RK4 solvers ImageNet-64, ImageNet- PSNR vs.
22

---

Published as a conference paper at ICLR 2024
J
AFHQ-256
RK0
RX2
RK2 B15
Figure 14: Bespoke RK2, RK1, RK2, and RK4 solvers on AFHQ-256; PSNR vs. NFE (left), and RMSE vs. NFE (right).
23

---

Published as a conference paper at ICLR 2024
K ABLATIONS
K.1 LOSS ABLATION
We consider here three losses for optimizing the Bespoke solvers: RMSE-Bound (the parallel loss we advocate in the paper), RMSE (optimizing directly equation 6), and a simplified version of the RMSE-Bound: sum of Local Truncation Errors (LTE-parallel). That is, LTE is defined in equation 26 but taking M = 1 for all i. Algorithm 5 provides the pseudo-codes for all three losses. We have run all three algorithms on the ImageNet 64 dataset and compared their FID, and RMSE, where RMSE is computed w.r.t. GT samples (see Section 4 for details). For the non-parallel RMSE loss. we needed to use Activate Checkpointing to run this loss. Figure 15 shows the results. As can be seen in the graphs, RMSE loss, as expected, reaches lowest RMSE values per NFE, the second best is the RMSE. loss and worst in terms of RMSE is the LTE. As for FID, RMSE performs worst, while for FM-OT model RMSE-Boun and LTE perform equivalently for NFE> 10, and LTE has advantage as far as FID goes otherwise. Since our goal is to reduce RMSE and provide memory training algorithm we opted to use the memory efficient RMSE-Bound the paper.
reduce memory consumption in order to be able
Algorithm 5 Bespoke training (parallel).
Require: pre-trained uy, number of steps nt
initialize E TR?
while not converged do
To ~ p(zo)
b sample noise D GT path
a=t) < solve ODE
RRSE loss then
x Bespoke sampling
Alg.
L <-
xx1) - 2
else if RMSE-paralle
loss then
L +-00 0
>> init loss
parallel for i 0, 0..., do
2 <- step (2rtes (t;), ux)
C -- M 7+1 (1++1) x
end for else if LTE ATE-parallel loss then 6-0
> init loss
parallel for i = 0, 0,..., n do (t,), f4; uz) L -- || m (1++1) - =2-1ll
2 + step (t,)
end for end if + 0 VeL end while return A
optimization step
K.2 SCALE-TIME ABLATIONS
This section presents an ablation experiment on the effect of each component in the scale-time transformation. We train Bespoke-RK2 solvers with three choices of transformation: (i) time-only: train fr., fr and freeze s, = 1, = 0, (ii) scalc-only: freeze f, = =,tt = 1 and train sy, $,, and (iii) scale-time: train both I, and 8y, Sr- All experiments are performed on ImageNet 64 FM- wawe-ooool/scloloysssoooeee
OT model. Figure 16 shows the FID, and RMSE of Bespoke-RK2 solvers and the base RK2 solver. First, all three Bespoke- RKK solvers improve upon the base RK2 solver. Second, the time component seems more significant, but the scale component improves FID for all NFEs and RMSE for NFE < 20. Third, in RMSE, the significance of the time component increases NFE increases. In addition, Figure 17 shows the trained time-only (top) and scale-only (bottom) tramsformations. Interestingly, we see that even scemingly small changes (c.g., scale-only with NFE E (16, 20}) can affect dramatically the FID.
24

---

Published as a conference paper at ICLR 2024
ImageNet 64: e-VP
ImageNet 64: FM/U-CS
ImageNet 64: FM-OT
LTE
LTT
LTE
RMSE
RMSE
RMSE
4.5
RMSE-8ound
4.3
RMSE-Bound
4.5
4.0
4.3
25
2.3
25
2.0
1.
10
RMSE
RMSE
RMSE
NFE
NFE
NFE
Figure 15: Different Bespoke losses for ImageNet 64. We compare the RMSE-Bound (the paraliel loss advo- cated the paper), direct RMSE loss, and Local Truncation Error (LTE) which a slightly simplified version of RMSE-Boum loss. All three variations are provided in Algorithm 5. In terms of RMSE the direct RMSE loss is as expected best however has large memory footprint, while RMSE-Bound is the runner-up and paral- lelizable. FID is perfectly correlated with RMSE and shows somewhat opposite trend (partially excluding the FM-OT model where RMSE- Bound LTE are almost FID equivalent). We opted for RMSE-Bound loss the paper since is memory-scalabee and provides best RMSE among the parallel loss options considered.
30.00
RI2-BES; MM2-BE::Scale-Time Scale-Time
RK2-BES: Time Only
RK2-BE5: Scale Only
RMSE
2.50
10
2.50
1.55
NFE
Figure Bespoke ablation I: RK2, Bespoke RK2 full scale seale-time optimization. time-only optimization (keeping 8r = fixed), and scale-only optimization (keeping f, = r fixed) on FM-OT ImageNet-64: FID vs. NFE (left), and RMSE vs. NFE (right). Note that most improvement provided by time optimization where scale improves FID for all NFFs, and RMSE for 20 NFEs.
K.3 TRANSFERRING BESPOKE SOLVERS
This section presents an ablation experiment demonstrating trained Bespoke solvers' generalization different models. We train a Bespoke-RKe -Kok solver on an ImageNet-64 FM-OT model and evaluate it on an ImageNet- 128 FM-OT model. We compare its FID, and RMSE vs. NFE against the base RK2 solver evaluated on ImageNet- 128 FM-OT and a Bespoke-RK2 solver trained and evaluated on ImageNet-128 FM-OT. The results is shown in Figure 18.
K.4 DISTILLATIIO-TYPEE PARAMETRIZATION
This section presents a distillation-type experiment in our framework. Our parametric family of solvers is defined by composition of scale-time transformation and the RK2 solver, resulting in step" as in equations 19 and 20, where the weights of the pre-trained model u, are frozen. A natural comparison to distillation-like approach is to let the weights of ut uf change during
25

---

Published as a conference paper at ICLR 2024
NFE=8
NFE=10
NFE=16
NFE=20
Naushon Deviiten
Soreleee Rowaln
Time Only
Scale- Only
Figure 17: Trained of scale-time ablation: Bespoke-RK' time-only optimization (top), Bespoke-RK3 scale- only optimization (bottom), ImageNet-64 FM-OT for NEE 8/10/16/20.
Transterred RK2-BE5 RK2--ES
RMSE
5.00
4.00
3.00
535
10
NFE
NFE
Figure 18: Bespoke ablation II: RK2 evaluated on FM-OT ImageNet-128 model, Bespoke RK2 trained and evaluated on FM-OT ImageNet-128 model, and Bespoke RK2 trained on FM-OT ImageNet-64 and evaluated FM-OT ImageNet-123 model (transferred): FID NFE (left), and RMSE vs. NFE (right). Note that the transferred Bespoke solver is still inferior to Bespoke solver but improves considerably RMSE compared the RK2 baseline. In FID the transferred solver improves over the baseline only for NFE=16.20.
optimization instead of using the scale-time transformation, that is
Lais (0) (xxx))
Σ [04+1(()))) step, (xi, 4i u |I
(80)
where I, ih is fixed. We perform this experiment ImageNet 64 FMr-CS. For fair comparison, we use the same compute budget as used to train our Bespoke-RKe KK2 solvers on these models: a total of 48K generated trajectories. 6k iterations, and report best FID. Figure 19 shows the RK2 base solver, Bespoke-RK. solver, and Distillation-Rke on ImageNet 64 FMu-CS: FID vs. NFE (left), and RMSE NFE (right). Note that while distillation is able improve from the baseline solver (RK2) it does not match the performance of the Bespoke solver. Two potential explanations why distillation not as performant as Bespoke in this experiment are: First, the amount of trajectories/compute budget we use to optimize Bespoke is not sufficient for effective distillation; and second, successful distillation methods require access to training data. Training only on generated data distillation less effective.
K.5 TRAINING STOPPING CRITERIA
In this ablation experiment, we qualitatively compare samples when changing the stopping criteria of the bespoke solver training. In Figure 20, we compare samples with bespoke-RK2 solvers where the training stopping criteria is best FID versus best RMSE on four models, ImageNet 64 e- VP/P/- CS/FM-OT and ImageNet 128 FM-OT, and NFE {8, 10, 20). In all of the cases the differences are practically indistinguishabl..
26

---

Published as a conference paper at ICLR 2024
RX2 RX2 Ditillation
16-
RK2-8ES
RMSE
10
20
22
34
16
25
NFE
NFE
Figure Distillation-type experiment: RK2 solver, Bespoke-Re solver, and Distillation using same number of trajectorics and compute as Bespoke, on ImageNet 64 FM/v-CS: FID vs. NEE (left), and RMSE vs. NFE (right).
GT
2F
NFE-
N
GT
NFT=20
NFE=10
NFT-
GT
NFF-20
NFE=10
NFF=8
RMSE RTD
FM-OT
FMPe-CS
RASE
FID
CYP
EMSE
FID
FM-OT
RMSE
FTD
FM-OT
RASS
Figure 20: Different stopping criteria experiment. We compare samples with bespoke-RKe MKakee-KK2oo solvers with train-
ing stopping criteria at best FID and best RMSE; ImageNet64 top rows) and ImageNet- 128 (2 bottom rows).
L TRAINED BESPOKE SOLVERS
In this section, we present the trained Bespoke solvers by visualizing their respective parameters 0. Figures 21, 22, and 23 show the learned scale-time transformation of Bespoke-RK2 solvers trained on ImageNet- 128, ImageNet-64, and CIFARIO (resp.) for NEE {8, 10, 16, 20}. First, we note the significant differences between the learned scale-time transformation of VP versus FM- on ImageNet 64 in Figure 22 top and bottom rows (resp.). Second, we note that the scale-time transformations trained on the same model type but different datasets to have similarities
27

---

Published as a conference paper at ICLR 2024
to some degree but still different from another, Figure 21 and Figures 22, 23 bottom rows, for FM-OT trained on ImageNet- 128, ImageNet-64, and CIFARIO (resp.). These two obser- vations showcase the advantage of custom-man -adi solver for each model. We also tested the latter point empirically in the ablation experiment in Section K.3, where we tested a Bespoke-RR K2 solver trained on ImageNet 64 FM-OT model to ImageNet 128 FM-OT model and noticed drop in performance compared to Bespoke solver trained directly on the ImageNet 128 FM-OT model. In addition, we note the resemblance in the form of scale-time transformations trained on the same model type and same dataset across different NFE (i.e., row in Figures 21, 22, and 23). This phe- nomenon suggests there may be a well-defined scale-timg transformation in the limit of NFE + oo. Furthermore, for l; and 8;, even and odd parity points the grid seem to converge to different curves, possibly due to their different role in the RK2 solvers.
NFE
NFF=10
1F
NF
Trasken Deskrter
Pranenn Devirrte
Paaponn Deriertee
-
PM-OT
Figure 21: Trained 0 of Bespoke- RK2 solvers on ImageNet- 128 FM-OT for NFE 8/10/16/20.
NFE=8
NFE=10
NFE=16
NFE=20
EMI6-CS
Taybbm Derietve
Taybm Sereetve
Mambre Servetive
Trautm Detivtbe
Figure 22: Trained 0 of Bespoke-RKi RK2 solvers on ImageNet-64 for NFE 8/10/16/20; VP (top), FM/v-CS (middle), and FM-OT (bottom).
M PRE-TRAINED MODELS
All our FM-OT and FM/e-CS models were trained with Conditional Flow Matching (CFM) loss derived in Lipman et al. (2022),
CCFM (0) (x. [,,; 0) (ō,30 O, 111
(81)
where t ~ U(O 1)). (2n) = / (zo]0, ((x)) is the data distribution, vy (z,;; 8) is the network, (or, 0x) is the scheduler defined in equation 22, and 2s d 0- For FM-OT the scheduler
28

---

Published as a conference paper at ICLR 2024
NFE=8
NFE=10
NFE=16
NFE=20
FMI
Tesbbm
Truurm Dutivtiv
Sutiutive
Trastmn Sutiutive
PM--T
Figure 23: Trained of Bespoke-RK2 solvers CIFARIO for NFE 8/10/16/20; e-pred (top), FM/e-CS (mid- dlc), and FM-OT (bottom).
is
04 t,
di o= t,
(82)
and for FM/e-CS the scheduler is
0t sin
at cos 2
(83)
2
All our c- VP models were trained with noise prediction loss derived in Ho et (2020) and Song et al. (2020b),
Luoise ( = (mo)) (|<((1;))) - zoll?
(84)
where the VP scheduler is
O: / 11-1-- Gi 8 (-
04 11-
(85)
and B = 20, b = 0.1. All models use U-Net architecture as in Dhariwal & Nichol (2021), and the hyper-parameters are listed Table
29

---

Published as a conference paper at ICLR 2024
CIFARIO eVP
CIEARIO FM-OT:FM/--CS
ImageNet-64 -NP.PM--O..MII-CS
ImageNet-126 FM-OT
AFIQ 256
FM-OT
Channels
128
128
196
256
256
Depth
4
4
3
2
2
Channels multiple
2,2,2
2,2,2
1,2,3,4
1,1,2,3,4
1,1,2,2,4,4
Heads
Heads Channels
64
64
Attention resolution
32,16,8
32,16,8
64,32,16
Dropout
0.3
1.0
0.0
Effective Batch size
512
512
2048
2048
256
GPUs
8
8
64
64
64
Epochs
2000
3000
1600
1437
862
Iterations
200k
300k
IM
900k
50k
Learning Rate
5e-4
le-4
le-4 constant
le-4
le.4
Learning Rate Scheduler
constant
constant
Poly Decay 5k
Polyn Decay
Warmup Steps
P-Unconditional
0.2 0.20 (vp.es), 0.15 (ot)
0.2
0.2
Guidance weight
0.5
0.1
Total parameters count
55M
55M
296M
421M
537M
Table 5: Pre-trained models' hyper-parameters.
N
MORE RESULTS
In this section present sampling results using RK2-Bespoke solvers, the RK2 baseline and Ground Truth samples (with DOPRI5).
30

---

Published as a conference paper at ICLR 2024
GT
NFE=20
NFE=10
NFE=8
GT
NFE=20
NFE=10
NFE-8
RK2
RK2-BES
RK2
RK2- BBS
RK2
RK2-BES
RK2
RK2- BES
RK2
RK2-BES
Figure 24: Comparison of FM-OT AFHQ-256 GT samples with RK2 and Bespoke-RKe RK2 solvers.
31

---

Published as a conference paper at ICLR 2024
GT
NFE=20
NFE=10
NFE=8
GT
NFE=20
NFE=10
NFE-8
RK2
RK2-BES
RK2
RK2- BES
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
Figure 25: Comparison of FM-OT ImageNet- 128 GT samples with RK2 and Bespoke-RKe RK2 solvers.
32

---

Published as a conference paper at ICLR 2024
GT
NFE=20
NFE=10
NFE=8
GT
NFE=20
NFE=10
NFE=8
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BESS
RK2
RK2-BES
Figure 26: Comparison of FM-OT ImageNet- 128 GT samples with RK2 and Bespoke-RKe solvers.
33

---

Published as a conference paper at ICLR 2024
GT
NFE-20
NFE-10
NFE-8
GT
NFE-20
NFE-10
NFE-8
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
RK2
BES
RK2- RK2-BB3
Figure 27: Comparison FM-01 OT ImageNet-64 GT samples RK2 and Bespoke-RKe solvers.
34

---

Published as a conference paper at ICLR 2024
GT
NFE-20
NFE-10
NFE-8
GT
NFE=20
NFE=10
NFE-8
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2- BES
Figure 28: Comparison FM/v-CS ImageNet-6o 64 GT samples with RK2 and Bespoke-RK' solvers.
35

---

Published as a conference paper at ICLR 2024
GT
NFE-20
NFE=10
NFE-8
GT
NFE-20
NFE-10
NFE-8
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BES
RK2
RK2-BESS
RK2
RK2-BES
RK2
RK2- BES
Figure Comparison e-pred ImageNet-64 GT samples with RK2 and Bespoke- RK2 solvers.
36

---

Published as a conference paper at ICLR 2024
NFE=20
NFE=10
NFE=8
GT
NFE=20
NFE-10
NFE-8
GT
NFE-20
NFE=10
NFE=8
CVP
RKQ BES
2
FMW--S
RK2-BIS BES
Figure 30: Comparison of VPV and FM/e-CS ImageNet 64 samples with DPM-2 and bespoke- RK2 solvers.
37


Reviews:
[{"review_id": "fSCP4KSr7V", "reviewer": "Reviewer_vykk", "summary": "This work addresses the issue of slow sampling in Diffusion models. The authors consider deterministic sampling of Diffusion models via the Probability Flow ODE. They address the issue by introducing a learnable (continuous) reparameterization of the ODE. Rather than learning the continuous parameterization directly, they discretize the reparameterized ODE using Runge--Kutta methods, which allows for learning only a discrete set of parameters (e.g. $4n-1$ parameters for Euler's method, where $n$ is the number of steps in the discretization). The parameters can then be learned by minimizing the discrepancy of the learned Runge--Kutta discretization to ground truth trajectories of the Diffusion model (which are in practice simulated using a higher-order adaptive ODE solver). The authors claim that the training time of the reparameterized ODE is roughly 1% of the training time of the original Diffusion model. The authors compare their method to the literature on Cifar-10 and ImageNet (64x64).", "questions": "**Depth of Experiments**: Since running the method is quite cheap, it would have been nice to include larger values of $n$ (currently only up to $n=10$) and higher order methods (e.g. RK3). Is there any particular reason why this has not been done?\n\n**Limited Outlook**: I think the idea of repameterizing the ODE with learnable parameters is neat. How could this method potentially be scaled/improved? Are more elaborate reparameterizations possible? \n\n**Guidance scale(s)**: What guidance scale(s) is/are used in the experiments?\n\nAlso see some questions entangled with the Weaknesses above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem of accelerating inference in Diffusion models is important\n- The structure of the paper is good, and the paper is generally well written (I am very happy that the authors opted to include Figure 2, and Algorithm boxes 2&3 in the main paper which help to understand the method)\n- The idea of learning parameters for a reparameterization of a Neural ODE is original (to the best of my knowledge)\n- The authors cover a lot of previous work (however I think some works are misrepresented, see also below)", "weaknesses": "**Missing details**: I think some details of the training are missing. How many GT trajectories are used and is the time for computing GT trajectories accounted for when claiming that the method only needs \"roughly 1% of the GPU time\" compared to training of the Diffusion model? This seems like a very important detail.\n\n**Single guidance value training**: As far as I understand, the authors need to retrain fresh parameters for each guidance value (and compute GT trajectories for each guidance value). I think the authors should have addressed this as drawback of the method (compared to other methods). \n\n**Misrepresentation of Distillation approaches**: I think the authors are misrepresenting the work on Diffusion distillation. For example, they say \"Distillation does not guarantee sampling from the pre-trained model's distribution\" - while this is correct the same applies to the proposed method. In fact, the inference distribution of a Diffusion model is always coupled to a numerical discretization; any solver will result in a different distribution of the generated samples.\n\nIn general, I also think the authors should have tried to compare the performance of their approach to distillation. For example, they could have done Progressive / Consistency distillation using a fixed compute budget. Unfortunately, there are no comparisons at all.\n\n**Interpretation of Figures 17/18/19**: It is actually quite nice that the authors can visualize the learned parameters (since there are so few) but unfortunately the work is missing a discussion on the figures. I think it's quite interesting that $s_r \\approx 1$ for most experiments and that $t_r$ seems to increase linearly with $r$. Could the method for example be used to explain the success of the DDIM ODE reparameterization?", "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": "eyPyy6pdGE", "reviewer": "Reviewer_7aEk", "summary": "This work introduces Bespoke Solvers, a method for efficiently sampling flow models by optimizing an instance-specific solver for the flow's particle ODE. Given a fixed budget of $n$ function evaluations, fitting a bespoke solver involves optimizing over a parametric family of $n$-point discretization schemes to minimize the average discretization error over new sample generations. To identify the relevant parametric family, the authors derive how a generic ODE solver changes under space- and time-reparametrization of the underlying particle trajectories. While these reparametrizations may be complicated functions of the underlying data, it is sufficient to know their values and derivatives _only_ at the $n$ discretization points, so that optimizing the bespoke solver only requires fitting $O(n)$ parameters. \n\nMinimizing the global discretization error directly may be computationally challenging due to the recursive dependence of the errors on previous time steps. Instead, the authors propose to minimize a weighted sum of one-step discretization errors, which can be computed at different times independently (and hence in parallel) and which is also an upper bound on the discretization error when the weights are chosen appropriately. The proposed method is shown to outperform existing dedicated solvers and it is shown to be competitive with distillation-based methods (ex. Ho and Salimans 2022) while requiring substantially less training time and fewer parameters.", "questions": "- Do you have any understanding of the limiting behavior of the bespoke solver parameters as NFE is increasing? In Figure 18, the parameters seem to be converging to a nontrivial limit as NFE increases. This is related to my comment in 'Weaknesses,' since one way to tune NFE could be to identify the limit of the solver parameters and to discretize it 'on demand' for different choices of NFE. \n- The scale-time parametrization seems like an arbitrary choice. Are there any benefits to using more complicated parametrizations? For example, did you run any experiments with time-dependent affine transformations (eg. adding an additive bias parameter to equation 14) or with higher order polynomials?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Clarity: the proposed method is elegant and it is well-explained in the paper. The derivations are correct to the best of my knowledge.\n- Highly practical and effective: the method is also cheap to train, adaptable to many existing architectures, and it leads to significant benefits for simulating flow-based image samplers. The experiments in Tables 1 and 2 are especially convincing", "weaknesses": "- No tunability for NFE: one minor weakness of this approach is that training the bespoke solver requires choosing up-front the number of function evaluations to be used in sampling. In contrast, non-instance dependent schemes can adjust to different NFE budgets at sample time, or they can be run 'until convergence' (choosing NFE adaptively for each particle trajectory).", "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": "uiGb6Hux26", "reviewer": "Reviewer_Zfm7", "summary": "This paper introduces the Bespoke solvers for generative flow models, which is customized to the pretrained flow models. To apply this, a model that returns a small number of parameters (<80) is introduced to learn the hyperparameters for the consistent ODE solvers. To define the parametric family of solvers, the transformed sampling path from an invertible transformation $\\varphi_r$ is first applied. \n\nFor transforming the sample paths, two components are applied: the time reparameterization $r\\to t_r$ and the invertible transformation $\\varphi_r$. This framework finally turns into the generalized noise scheduler for all possible diffusion models and flow models, by normalizing the inputs with the invertible transformation and denoting the timestep with the time parameterization function. This is realized by learning the hyperparameters that consists the scaling function $s(t)$ and time parameterization function $t_r(t)$.\n\nNext, the tractable RMSE loss that computes the global truncation error between the approximate sample and the ground truth (GT) sample is enabled: the upper bound of the RMSE loss is formulated by the weighted sum of the local truncation errors, weighted with the Lipschitz constants. By minimizing this RMSE loss, one can reduce the gap between the approximate sample (given by the solver) and the ground truth (given by the oracle ODE solver).", "questions": "(1) If this bespoke solver is trained well with the higher-order solver, it is curious about the optimal required order of the solver with the learning-based coefficients: this means, if the higher-order solver is not required, the higher-order coefficient $t_r ''$ (if exists) or $s_r ''$ is expected to have low value, which is nearer to zero.\n\n(2) Is there any ablation with the $s_i$-only or $t_i$-only cases? (In this cases, $t_i$ is uniform (or EDM-like) or $s_i$ follows the VP, VE, or EDM preconditioning...) This additional will help understanding the advantages of training these hyperparameters.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "(1) This method solves the important problem in the diffusion/flow models, the time-step and input scaling problem by learning hyperpamareters of the pretrained model ubiquitously, by optimizing the time steps and the input scaling with just some data-driven optimization with the pretrained model. If trained properly, and the bound between RMSE loss and the global truncation loss is rightly narrowed, then sampling from this learning-base parameterization derives good sampling results, as the paper proposes.\n\n(2) The writing is compact and sound; one can easily understand why the learning-based parameterization is required and how this benefits the sampling process of the diffusion/flow models, with abundant experimental results and theoretical supporting materials.\n\n(3) According to Table 3, this method requires much less training time of the hyperparameters (about 0.5~2% of the original training time) compared to the existing distillation-based method for fast sampling.", "weaknesses": "(1) Even though the RMSE objective is upper bounded by the Lipschitz loss, there can be some gap between these two loss; unless the loss converges to zero.\n\n(2) There results are not yet validated with larger-scale (than size $64\\times 64$) datasets, like FFHQ or LSUN.", "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": "4TIz0yE5Rg", "reviewer": "Reviewer_KhFN", "summary": "This paper presents a distillation approach diffusion models that is based on ODE solvers. The family of Runge-Kutta ODE integrators is used as a basis for performing a \"time integration\" for the diffusion time with learned coefficients for the different time steps. The number of time steps is chosen quite low such that correspondingly a low number of function evaluations (NFEs) of the pre-trained diffusion model is needed. A small network is trained that parametrizes the time integrator, with the benefit of a small number of weights and fater training.\n\nThe NFE count is used as a central \"dimension\" along which the RMSE and FIG performance of different versions is evaluated given a certain number of NFEs. The approach seems to achieve a very good performance with relatively low NFE and fast training. \n\nThe shown samples are close to the \"GT\" versions (with larger NFEs from the original model), and seem to perform better than running an RK-integrator directly on the pre-trained model. The images are typically shown for 20,10 and 8 evaluations, and the versions of the proposed method usually show few variations, i.e. stay close to the GT version. This is of course raises the question how even fewer steps would perform, and when the visual difference would start to grow. This is unfortunately not (yet) demonstrated, but should be easy to add.", "questions": "1) I was surprised about the statement that best FID iterations are \"reported\", but best RMSE iterations is shown. So the results shown in figures 6,7,... are from a different model than the graphs and tables shown? Can the authors can explain why?\n\n2) The introduction claims a \"very small number of learnable\" parameters, but I didn't find a table listing the actual parameter counts of the original models and the additional parameters for the \"bespoke\" versions. Can the authors provide these?\n\n3) Can the authors re-run some cases of \"bespoke\" training with a regular RMSE loss, e.g., using GT values from a more finely sampled reference trajectory? This would highlight the influence of the loss.\n\n4) Can the authors show some results from models with fewer steps than 8? E.g., 6, 4 or 2 steps?\n\n----\n\nI think the authors have done a nice job addressing the remaining open questions in their rebuttal, and included a collection of additional background information and illustrative results. As such, in line with my previous score, I would recommend an accept for this paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The paper targets an important area, the runtime of generative models from diffusion approaches. The distillation via a parametrized and learned ODE is a new idea as far as I can tell, and the results are convincing. The quantified performance of the models as well as the qualitative examples look very good to me.\n\nOverall, I find the results convincing, and would suggest an \"accept\" given my current understanding of the work. There are a few smaller open points below which I hope the authors can clarify in the rebuttal.", "weaknesses": "While the main approach with custom RK integration makes a sound impression, the loss from section 2.3 seems a bit weaker in comparison. It essentially seems to yield a weighting the the MSE terms with the M_i coefficients, and replacing the acutal Lipschitz constants with 1 everywhere seems ad-hoc. I was also missing an ablation showing how much the training benefits from this loss over a \"regular\" RMSE loss.", "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": "mlCyrsSItR", "reviewer": "Reviewer_ovsR", "summary": "The paper proposes a framework for creating custom ODE solvers tailored to the ODE of a given pre-trained flow model. The framework is efficient in terms of additional parameters and training times and can produce high-quality images with a low number of function evaluations.", "questions": "- Is proposing a learnable transformed path instead of fixed and cherry-picking ones, as in previous works, the main contribution of the paper?\n- Does increasing NFE to 35 boost the performance of your method to be comparable to EDM in [1]? More generally, what happens to your method's performance when NFE is larger than 20? The same questions apply in the case of FID < 10.\n- Can your method possibly be applied to solve Stochastic Differential Equations (SDE) directly?\n\n**References**\n\n[1]  Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. \"Elucidating the design space of diffusion-based generative models.\" Advances in Neural Information Processing Systems, 2022.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The method is sound and demonstrates good results in fast sampling. \n- The training time is short, and the number of parameters is low.", "weaknesses": "**Major concerns**\n- Both the quality and the speed of sampling (the number of function evaluations) of the proposal are not as competitive as distillation methods.\n- Karras et al., 2022 [1], propose a time discretization method to determine $t_i$, achieving an FID of 1.97 with unconditional image generation on CIFAR10 when NFE = 35. Your proposed method also seeks the sequence of $t_i$, so I strongly compare it with Karras's method.\n- Minor Point: There is still room for improvement in the presentation of the paper, especially in notations and function definitions in Section 2.1. It takes me a considerable amount of time to understand.", "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 (spotlight)