Bellman Optimal Stepsize Straightening of Flow-Matching Models
Authors: Bao Nguyen
Binh Nguyen
Viet Anh Nguyen
Venue: ICLR 2024 poster (2024)
Keywords: flow matching
generative model
efficient sampling
distillation
responsible ML
Primary area: generative models
TL;DR: This paper introduces Bellman Optimal Step-size Straightening (BOSS) technique for distilling flow-matching generative models while adhering to a computational budget constraint.

Abstract: Flow matching is a powerful framework for generating high-quality samples in various applications, especially image synthesis. However, the intensive computational demands of these models, especially during the finetuning process and sampling processes, pose significant challenges for low-resource scenarios. This paper introduces Bellman Optimal Stepsize Straightening (BOSS) technique for distilling flow-matching generative models: it aims specifically for a few-step efficient image sampling while adhering to a computational budget constraint. First, this technique involves a dynamic programming algorithm that optimizes the stepsizes of the pretrained network. Then, it refines the velocity network to match the optimal step sizes, aiming to straighten the generation paths. Extensive experimental evaluations across image generation tasks demonstrate the efficacy of BOSS in terms of both resource utilization and image quality. Our results reveal that BOSS achieves substantial gains in efficiency while maintaining competitive sample quality, effectively bridging the gap between low-resource constraints and the demanding requirements of flow-matching generative models. Our paper also fortifies the responsible development of artificial intelligence, offering a more sustainable generative model that reduces computational costs and environmental footprints. Our code can be found at https://github.com/nguyenngocbaocmt02/BOSS.

Published as a conference paper at ICLR 2024
BELLMAN OPTIMAL STEPSIZE STRAIGHTENING OF FLOW-MATCHING MODELS
Bao Nguyen
Binh Nguyen
VinUniversity
National University of Singapore bi 2 @nus. edu. sg
bao. n.o.in22vimmm nuni. edu. vn
Viet Anh Nguyen
Chinese University of Hong Kong nguyenêse. cuhk. edu. hk
ABSTRACT
Flow matching is powerful framework for generating high-quality samples in various applications, especially image synthesis. However, the intensive com- putational demands of these models, especially during the finetuning process and sampling processes, pose significant challenges for low-resource scenarios. This paper introduces Bellman Optimal Stepsize Straightening (BOSS) technique for distilling flow-matching generative models: it aims specifically for a few- step efficient image sampling while adhering to a computational budget con- straint. First, this technique involves dynamic programming algorithm that optimizes the stepsizes of the pretrained network. Then, refines the veloc- ity network to match the optimal step sizes, aiming to straighten the genera- tion paths. Extensive experimental evaluations across image generation tasks demonstrate the efficacy of BOSS in terms of both resource utilization and im- age quality. Our results reveal that BOSS achieves substantial gains in efficiency while maintaining competitive sample quality, effectively bridging the gap be- tween low-resource constraints and the demanding requirements of flow-matching generative models. Our paper also fortifies the responsible development of ar- tificial intelligence, offering more sustainable generative model that reduces computational costs and environmental footprints. Our code can be found at https: /ttes://githann
02/BOS5..
INTRODUCTION
There have been impressive advancements deep generative models in recent years, which con- stitute an appealing set of approaches capable of approximating data distributions and generating high-quality samples, as showcased in influential works such as Ramesh al. (2022); Saharia al. (2022); Rombach et al. (2022). They are primarily driven by a category of time-dependent genera- tive models that utilize a predefined probability path, denoted as (Te hejo,111 This probability path process that interpolates between the initial noise distribution xo and the target data distribution æi- The training for these models can be broadly characterized as a regression task involving a neural network function Uy and target ideal velocity vz(x);
((0)= Ergj,, X.~a, [f(vs ((ooo(X,)) (X,, (x.)))
Here, the velocity network vo maps any input data time I. [0, 1) to vector-valued velocity quantity 2e t) and plays a crucial role in the generation of samples from the interpolation process through the relationship:
= + L' vo (o.,.)))dk
which is the solution of the ODE dX,, (x,, t)dt with the boundary condition Xı=0 Xo- A noteworthy class of algorithms that fits within this framework includes denoising diffusion mod- els (Ho et al., 2020; Sohl-Dickstein al., 2015; Song et al., 2020b) and the more recent flow-

---

Published as a conference paper at ICLR 2024
matching/rectified- flow models (Liu et al., 2022b; Lipman et al., 2022; Albergo & Vanden-Eijnden, 2022; Neklyudov et al., 2023).
The latter type of model extends the principles employed in training diffusion models to simulation- free continuous normalizing flows (CNF, Chen ct al. 2018). is particularly attractive because fixes the suboptimal alignment between noises and images of diffusion models by introducing straight trajectory formula connecting them. This leads to (empirically observed) faster training and inference time than diffusion models. The rectified flow framework Liu et al. (2022b) also includes a technique called reflow, which gradually rectifies the probability paths, It significantly reduces the number of function evaluations needed for sampling and thus belongs to the family of distillation methods.
However, the standard reflow technique proposed in Liu al. (2022b) requires a significant amount of computational budget: on a mime--nim dataset such as CIFAR-10 x32x32 pixel images), it takes at least an additional 300,000 retraining iterations of the pretrained velocity networks to reach FID (Fréchet Inception Distance, Heusel et al. 2017) of 4,85 for 1-step generation. The additional retraining time can reach approximately 200 days of A A100 GPU for distilling models for 1-step sampling on higher-dimensional scale datasets to achieve competitive FID, as stated in an extension of the rectified flow framework (Liu et al., 2023). Motivated to fix this problem, in this work, we aim to distill the Rectified Flow model while satisfying the following objective.
Given a pretrained flow-matching velocity Uo and target of K number of function evalua- tions (NFEs), how can we adapt Uo with a proper sampling schedule to generate high-fidelity images using only modest computational resource?
Contributions. We propose BOSS, the Bellman Optimal Stepsize Straightening method, to finetune pretrained flow-matching models. Our proposal includes two phases. The first phase seeks the optimal K-element sequence A for the initial model vg. The second phase utilizes A to retrain vo such that the retrained model to- performs better. With the proposed procedure, we straighten the velocity network with just about 10,000 retraining iterations while outperforming the standard reflow strategies regarding image quality. Quantitatively, our procedure consistently achieves lower FID in unconditional image generation with four different datasets. Furthermore, as the additional results in Appendix E show, the straightening procedure using -ow- Adaptation (LoRA) can finetune only 2% of the model's parameters, yet it performs competitively to that of full-rank finetuning.
Related works. Our approach is directly related to existing works on improving the sampling ef- ficiency of diffusion and flow-matching models with training-based algorithms. In Salimans & Ho (2021), the authors proposed an approach to enhance the sampling speed of unguided diffusion models through iterative distillation. This is extended to the case of classifier free guided diffusion models in Meng et al. (2023). In Wang et al. (2023), the authors propose a method leveraging re- inforcement learning to automatically search for an optimal sampling schedule for Diffusion Prob- abilistic Models (DPMs), addressing limitations in hand-crafted schedules and the assumption of uniformity across instances. Our work has a few common features and motivations with Watson al. (2021), which achieved significant speed-ups through dynamic programming and decomposed loss terms. However, their focus on individual Kullback-Leible divergence loss that neglects the cumulative information loss during sampling. This results in images with reduced overall quality, In contrast, we focus on minimizing the local truncation error during the sampling procedure, which improves the image quality consistently across all budgets of NFEs. Moreover, we propose a fine- tuning method that allows faster sampling with just a few NFEs. Recent work by Song et al. (2023) introduced a framework that learns a model capable of mapping any point at any time to the tra- jectory's s starting point, called the consistency model. After submitting this work, we discovered concurrent study by et al. (2023). This work pointed out that using uniform stepsizes is subopti- mal for diffusion model sampling and instead used evolutionary algorithms search for the optimal stepsizes and score network architectures, with the FID score being the optimized metric.
Within the context of Rectifed Flow/Flow Matching, Liu et al. (2022b) proposed a reflow method that uses retraining to straighten the probability sampling path. This results in low NFE sampling with favorable image quality. The recent work of Liu et al. (2023) takes this framework to larger scale, demonstrating impressive results on high- high-resolution image datasets. However, both rely on computational intensive retraining procedures, which we improve in our work.
2

---

Published as a conference paper at ICLR 2024
The other direction that aims to accelerate the sampling process of diffusion//low matching models is training-free samplers (Song et al., 2020a; Bao al., 2022; Liu et al., 2022a; Tachibana et al., 2021; Zhang & Chen, 2022; Karras et al., 2022; Lu et al., 2022: Zheng et al., 2023). Although required no additional training step, these works mainly relied on the properties of the SDF/probability flow ODE dynamics to propose heuristic solvers/diffusion noise schedulers. Therefore, verifying whether the proposed sampling stepsizes are optimal is hard.
BACKGROUND
Suppose we given a (pretrained) model vo (x(,,, t) with parameter 0, which is estimator of the function v from an ordinary differentiable equation (ODE) on the span t E [0, 1]:
dX, = v(X,, t)dt.
(I)
In generative modeling with diffusion/flow matching models, this dynamic system is called the probability flow ODE (Song et al., 2020b; Lipman et al., 2022). The estimator eo allows us to flow from the distribution (noises) to the distribution πi (real images) via the equation:
X1 = Xo L' vo (x,
(2)
where Xo To and X1 TI In the context of our problem, Xo is observable. Xi is only determined by the equation (2) which a deterministic process that for each Xo = Zo. there is only value X2 21 coupling with it through the following equation:
[ volae, tjdt.
To +
We are interested in the low-cost estimate of the integral So us va t)dt with respect to f over the interval [0, IJ. It is essential concern when calculating the velocity field is computationally ex- pensive. The amount times calling Up is defined as the number function evaluations (NFE).
First Order Sampling scheme. To solve for the integration that appears in the sampling equa- tion (2), is necessary to invoke numerical integrator that uses discretized time steps. Any numer- ical integration scheme will induce truncation errors, which can be quantified in two forms: first, when we have the value at the previous time step X,-8. the solver estimates the subsequent true value X, as X,, causing a local Iruncation error X, X,- These local errors accumulate over the number of intervals, eventually resulting in cumulative error known as the global truncation error. The most popular discretization scheme is perhaps Euler's method: given a budget K number of function evaluations (NFEs), the Euler uniform sampling computes the interval A 1/K, and the sample successively
# ()/K U0 (k- 1)4) x A
K,
(3)
with the initial condition =1 To- We denote this uniform sampling scheme by Eo (K(K. Fuler's method with uniform stepsizes A has local truncation error O(A²), and global truncation error O(A). One can generalize the Euler sampling with non-uniform intervals by dividing the time domain [0. 1] into unequal intervals with timestamps 0 = To successively
TK- = 1, and sample
Tx = x + Uo Co 1o Tk-1) x (Th - 1 Vk = 1, K,
(4)
with the initial condition on ~ Tp- The timestamps equivalently determine the stepsize T T for each sampling iteration. We denote this scheme by E((T), Tl. }K)). If the timestamps T
are equally spaced in [0, 1]. then we obtain the equivalence E(TTo, Th, ... TK }) = (K)). Since
the reflow procedure in Liu (2022b) deals exclusively with Euler's method for being the fastest
with fixed computational budget, we focus only on this method in our paper.
3 OPTIMAL SAMPLING STEPSIZES
The objective presented in Section can be defined more rigorously as follows. Given a pretrained model vo and insignificant value of K, we aim find the optimal value 0" and sequence A such
3

---

Published as a conference paper at ICLR 2024
that E(., A+) is reasonable estimate for the coupling sample of any This can be posed as an integer optimization problem of finding the best schedule for sampling. Given a fixed budget of K NFEs, we find schedule {To, Ti..... TK satisfying To Ti ... Tк and that the associated Fuler non-uniform sampling scheme T(T ....... TK }) has minimal sampling error for the pretrained velocity vo- In Section 3.1, describe estimate of the sampling error for any valid schedule. Section 3.2 presents an integer programming formulation to find the optimal stepsizes for sampling, and Section 3.3 provides a dynamic programming algorithm find the optimal schedule.
3.1 SAMPLING ERROR ESTIMATION
Given any two arbitrary timestamps 0 S 1 < La I 1, we are interested in estimating the local Fuler truncation error, i.c., measuring the discrepancy between the true value
X Xe, + 1/ uo (X,,
and the one-step Euler sampling value
x = Xt, + us (Xe, (2) (tx - -t))
where Xz, is sampled from the distribution Tr, which is induced by the initial distribution To of Xo and the ODE (1). This local truncation error can be formalized as
truncation = EX my [II S, Ojk
ue (X,, t)dt us va( tyx -t))tt [1)
Unfortunately, computing Gk
cruncction is computationally intensive because of both the expectation operator and the integration. We instead employ the following two simplifications:
1. We fix the possible choice of time-stamps: for sufficiently large number Kmmax, the anchoring timestamps (42 IU..,,,,.. Kmax with = k// Kmax, In doing so, we have restrained the space of all possible sampling schedules to the combinations of finite anchoring timestamps {4x Later in the numerical experiments, we choose mmax 100, leading the anchoring timestamps {0, 0.01, 0.02..... 0.999 1}.
2. We approximate the local truncation error jjk by a sample average estimator Cjk. that constructed as follows:
ctrnncatin for any two anchoring timestamps t, < tx
N
Cjk-
jk
where Cjx = =111/1 x1; co (1) x (ts - 1
(5)
where for sample i, the noise af drawn from no. and of, and w/ are extracted from Fuler uniform sampling path starting from zi. taken at time t, and t. respectively.
2 + w (m 4o) 24
Figure illustrates how the sampling error sj are calculated for the simple case with Kmmax" 5 (or equivalently, with a uniform interval A 0.2). First, noise z is drawn from To. and the blue curve depicts the nonlinear trajectory following the ODE (I). The piecewise linear trajectory is the path generated by the uniform Euler sampling with K'max NFEs, leading to the observed trajectory (#f)) )==.... Kmax. For a concrete example of computing c2s we mea- sure the difference between the value of a one- step Euler sampling from f2 to to with a step- size t5 t = 3A to obtain x. (ai x33,, and the observed value zτ.. Intuitively, we can view Cjk the difference between the Euler one-step and the Euler (k - j)-step uniform sampling between iy and 4--
x
i
00
va (a)..t)
(x) x
2
x
23
C
A
0.2
20
time
to
6
12
t3
t4
t5
1
Figure 1: An example with K Kmmax to illustrate the computation of the sampling error.
4

---

Published as a conference paper at ICLR 2024
CO Km
Coj
C1j
C02
C12
tu 0
1/
source node
C2Kmax
sink node image
noise
Figure 2: A network flow formulation to find the optimal sampling schedule for image generation. Time to = 0 represents noise, while Kymax = is the terminal data (images). Each discretized timestamp is represented by a node, with edges reflecting the one-dimensional flow of time from noise to image. The cost Cjk associated with each edge is the sampling enror estimate, measured by the average difference between the Euler one-step and the Euler (k - j))-ttp sampling between ty and tu, see Section 3.1.
One may observe that Gjk is only an approximation of the true local truncation error Atrnncatin because Cjk is computed based the Euler trajectory (red piecewise linear path in Figure 1), while the truncation error ctruncation should be computed based on the nonlinear trajectory of the ODE Cjk (blue curve in Figure 1). One downside of using Cjk is that for any two consecutive timestamps to and tj+1. we have Cj((++) = O. This downside can be mitigated by taking Kmmax sufficiently large. On the other hand, as K Kriaxx gets large, calculating all the values Cjt: is computationally intensive because there are, in total, Kmaax (K (Kmmx 1)/2 pair of timestamps whose errors are to computed. Nevertheless, demonstrate empirically in Section 5 that even when Cyk is computed using small number of samples, the resulting optimal schedule already demonstrates a superior performance vis-à-vis competing methods.
Cjk
3.2 INTEGER PROGRAMMING FORMULATION
As we now describe, finding the optimal sampling schedule can be formulated as network-flow. based problem (Ahuja et al., 1993). First, construct a graph of Kmma + nodes; each node represents one timestamp, see Figure 2. There is an edge connecting node t; to node t. if tj < tg. and this edge associated with sampling error cost Cjk, computed in Section 3. For target K NFEs, the optimal sampling schedule is path that traverses from the source node 4n to the sink node [gmax that is comprised of exactly K edges. This path can be recovered from the optimal solution of the problem
m
K
min
2jk2
1=0
k=j+1
Km
8.
==
Lk=jt1 2jk K
(6)
Kms
5
k=1 20k
2j Kun
Ok=0 žkj
Σ
-==++1
Kix
2jk
Vj E [1, K maxx mma
2jke E (0.1)
vO < < h: < K max
Above, Σjı E {0, is a binary decision variable, 2,a = if the path takes a one-step sampling from time to time 4. The objective function of (6) minimizes the path's accumulated sampling error, which approximates the global truncation error of the Euler sampling with the corresponding step sizes. The first constraint indicates that the path should consist of exactly K edges; the second constraint imposes that to and xmax are the source and sink nodes, respectively. Finally, the last set of constraints is the flow conservation each intermediary node between to and gaax.
5

---

Published as a conference paper at ICLR 2024
3.3 DYNAMIC PROGRAMMING ALGORITHM
While the integer programming problem (6) can be solved using commercial solvers such as GUROBI or using network flow algorithms (see Skiena (2008, $6) for an example), there are prac- tical cases in which we need to find optimal paths for multiple values of the budget K NFEs. A convenient way to address this computation is to leverage dynamic programming formulation, which successively builds up the error-to-go function at each node and for each number of remain- ing NFEs. To this end, for any timestamp by and any number of remaining NFEs & [1, Kmax], define the error-to-go function
K
K
K
min
Σ
Cjk2yb
k=j-1
K
k==+t1 2zH = k
K
S. t.
kG. =
(7)
Σ kk=j+1 Pjk
Kama
k Zkj = > k= 2jk
mm
Vj f K m=] vj S < v K max
3jk E {0,
The error-to-go VG, is the minimal sampling error accumulated from time 1 to the terminal time fgmax using exactly le NFEs. It is easy to see that the optimal value of problem (6) is equal to (00, K). At initialization, we set the base case
vG.1) = JjKmax VJ [_,K
The dynamic programming update step rolls backward following
4k K....,,maa... Kmax
k(j.k)
min j<jsKmux
OB + VG
(8)
The output of the dynamic programming algorithm is the error function V, and one can assess the Bellman optimal schedule by tracing the minimizing path following (8) for each value K of NFEs.
4 STRAIGHTENING FLOWS WITH BELLMAN STEPSIZE
Given the Bellman optimal stepsizes, we de- scribe a piecewise linear straightening of the sampling curve. The straightening procedure aims to re-align the velocity network vp to re- duce the accumulated sampling error at the ter- minal timestamps. For a fixed number of NFEs K, let (TO. ( } be the optimal timestamps found in Section 3 with To 0 and тк = 1, which corresponds to K stepsizes defined by Th Th- for k 1,...,K K. We now mod- ify the network weights to straighten the sam- pling path on each interval Tk41l- An in- tuitive explanation for the straightening pro- cedure is illustrated in Figure 3: here, sup- pose that K 2, and the optimal schedule is 1). For the sample i drawn in Figure 3, the Bellman sampling in- duces two-piece linear path z + 2 - x (dashed line). If the velocity vectors evaluated at inj and z align with the dashed line, then the Bellman optimal Euler sampling with K = 2 incurs zero loss compared to the Fuler uniform
322
}
D
i
time
to 0
t2 T1
to = 1
Figure 3: Continued example following Figure for straightening with K 2 NFEs, evaluated at time 40 = and time l2 = 0.4. Blue arrows are velocity vectors given by the pretrained model, and purple arrows following the dashed lines are the ideal straight path. The straightening proce- dure in Section 4 aims to align the blue arrows towards the purple arrows. Arrows illustrate di- rections and are not drawn with proper scale.
(To 0, = 0.4,

---

Published as a conference paper at ICLR 2024
sampling with K
5. This motivates the following alignment procedure:
K-
min E Xxowto
Σ
(x
Th))
Tx-
Tk+
7k
where X (
is obtained by the Euler uniform sampling with K NFEs of the initial condition Xa ~ To. Replacing the expectation with empirical paths obtained by E k Kmax), we have the sample averaging optimization problem for straightening
K
x²
min
I (1 Th)
T6+1 Tu
1.
Tk++1 7
K=0
We straighten the velocity model using a stochastic gradient descent algorithm to solve the above problem. In the main paper, we train all parameters of the pretrained model, whereas in Ap- pendix E, we employ Low-Rank Adaptations to reduce the number of trainable parameters while preserving the performance of the straightening process.
NUMERICAL, EXPERIMENTS
Settings. We evaluate our methods on unconditioned image generation tasks. In particular, we use the CIFAR-10 (Krizhevsky et al., 2009) and three high-resolution (256x256) datasets CelebA-HQ (Karras et al., 2018), 1)), LSUN-Church, LSUN-Beeroom (Yu al., 2015), and AFHQ-Cat. We take the checkpoints of pretrained velocity networks v from the official implementation' of Rectified Flow (Liu et al. 2022b), which is based on the U-Net architecture of DDPM+ (Song et al., 2020b). If not mentioned otherwise, we evaluate the sample schemes with NFE=(4, 6, 8). The quality of image samples is with Frechet inception distance (FID) score (Heusel et al., 2017).
Baselines. A comparison between Bellman optimal stepsizes and the conventional first/second order methods using uniform stepsizes is presented in Section 5.1. We also include adaptive strategy, the Runge-Kuue Kutta method of order 5(4) from SciPy (Virtanen et al., 2020). In Section 5.2, our fine- tuning procedure, Bellman Optimal Stepsize Straightening (BOSS), is compared with baselines including the Uniform-Reflow, and Distill-k-Reflow introduced in (Liu et al., 2022b).
5.1 IMPROVEMENTS OF FIRST ORDER AND SECOND ORDER SAMPLING SCHEME USING BELLMAN OPTIMAL STEPSIZE
First, we benchmark pretrained Euler samplers with uniform and Bellman optimal stepsizes, caleu- lated following Section 3.3. The quantitative results are demonstrated in Table 1. The FID is much lower for samples generated by Euler's 's method with the Bellman step. which shows that with Bell- man's optimal step size, the generated images are of much higher quality general. Specifically, for sampling on the three larger dimension datasets (256x256), Bellman steps can help drastically reduce FID compared the uniform step size. This trend is also reflected in Figure 5, which shows our qualitative results.
Table 1: FID (L) of Euler's sampling method with uniform stepsizes vs. Bellman optimal stepsizes on unconditional image generation task across different datasets.
4 NFE
NFE
8 NNFE
Dataset
Euler
Bellman
Euler
Bellman
Euler
Bellman
CIFAR-I0
51.95
47.57
25.69
23.35
16.82
15.74
CelebA-IIQ
158.95
92.03
127.01
72.54
109.42
49.80
LSUN-Church
106.94
80.91
53.85
45.09
34.74
33.22
AFHQ-Cat
68.95
45.54
61.50
36.15
56.96
33.94
84.35
61.60
39.19
35.35
32.15
25.80
We also empirically analyze the effect of Bellman optimal stepsizes on popular ODE solvers includ- ing Euler and Heun (second-order version). To avoid the confusion between schemes, we denote compared methods as follows:
httpps://gi thub. htwpbb.co/goobittttta tab/Recti FiedFlow/

---

Published as a conference paper at ICLR 2024
Uniform Euler and Uniform Heun are the conventional Euler and Heun methods that uniform sampling steps.
Bellman Euler and Bellman Heun are two variants of the above methods, but using our proposed Bellman step sizes.
RK45 is an adaptive strategy, the Runge-Kuutt Kutta method of order 5(4) from Scipy.
The quantitative results are displayed in Figure 4. is evident that employing Bellman optimal stepsizes can significantly improve FID scores (image quality) compared to using uniform stepsizes
across all pretrained models on four distinct datasets. The Bellman Heun methods approach the performance of RK45 with substantially fewer NFEs. To elaborate, the Bellman Heun method achieves approximately a 1% gap compared to RK45 with just 20 sampling steps and fully recovers the performance of RK45 with 50 steps.
100.0
- Unifum Ederr
+
Usioom Edlo
- Unioom Euler
Unifurm Heun
Unifum Houn
- Unioo Houn
FID (log
50.0
- Bellman Heun R545
- Belinan Buler
FID clogsscaee
100.0
melman Houn
melmman Eveer
FID (leg-scaal)
100.0
Iman Euter
Iman Houn
20.0
BC4S
H945
5.0
20.0
2.6
11.3
21.9
20
205
Number of Stops (log.sce scale)
Number of Stops (log scale)
Number Steps scals)
(a) CIFAR-10
(b) CelebA-IIQ
(c) AFIIQCaa
100.0
Untom Duker
100.0
Untum Buker
Untom heus
Untom heun
FID (lop-scale)
Belmmnn Edler
FID (log scale)
Beliman Euler
50.0
Delmmn Heus
RK45
50.0
Bellman Hiun
RK-15
11.4
11.6
20
50
206
10
50
209
Number of Steps (log-scale)
Number of Steps (log-scale)
LSUN--hurch
LOUN-Bedroom
Figure 4: The FID score of sampling methods with different numbers of function evaluations (step sizes). Images generated by samplers using Bellman stepsizes clearly show lower FID than conven- tional ones that use uniform step sizes. Note that Uniform Heun and Bellman Heun are second-order sampling methods that use twice the NFEs.
5.2 EFFECTS OF REFLOW WITH BELLMAN STEPSIZE
After calculating the Bellman optimal step size, we follow the procedure described Section to straighten the probability path. The results,
in Table 2, suggest that BOSS performs al- most equally with the reflow procedure on CIFAR-10 but markedly better on the other four higher- dimension datasets. This consistent with the visible improvements in sampled image quality observed in Figure 5.
Table 2: FID (L) of different retraining methods on the unconditional image generation task across different datasets. Distill-K K-Reflow a distillation technique that relies on the reflow the velocity network on discrete grid of K uniform stepsizes between and I, elaborated in Liu et al. (2022b).
Dataset CIFAR-10
Rstill-6-Reflow
BOSS-6
Uniform-Refloe
4.35
4.80
4.33
CelebA-HQ
34.56
18.67
43.57
LSUN-Church
34.52
17.43
40.45
AFHQ-Cat
46.24
26.10
51.24
LSUN Bedroom
41.17
18.45
45.13
8

---

Published as a conference paper at ICLR 2024
(a) Fuler (6 NFEs)
(b) Bellman (6 NFEs)
(c) BOSS (6 NFEs)
(d) RK45 (208 NFEs)
Figure 5: Qualitative results unconditional image generation task. From first to last row: CelebA- WWI.SSS---SSN Church/AAHHQ-C dataset. (a)-(b): Comparisons of Euler stepsizes
between uniform (a) and the Bellman optimal stepsizes (b); (c)-(d): Comparisons of BOSS retraining and Kunge-Kutta-45 sampling. Notice our proposed BOSS sampling has comparably similar visual quality RK45 while requiring only 6 NFEs, compared to 208 NFEs of RK45.
CONCLUSIONS
This paper proposed BOSS, the Bellman Optimal stepsize Straightening method, to adapt pretrained flow-matching models under low computational resource constraints. Our method consists of two phases: first, find optimal sampling stepsizes for the pretrained model, then straighten out the veloc- ity network on each interval of the sampling schedule. We demonstrate empirically that BOSS performs competitively in adapting pretrained models in the image generation task. Similar to training-basee samplers for diffusion and flow matching models, limitation of our method the additional training cost to output the optimal sample step sizes. There many potential extensions to our proposed framework to distill a guided velocity network, similar to Meng et al. (2023), or computationally cheaper algorithm for calculating the Bellman sampling step sizes.
Acknowledgments. Viet Anh Nguyen gratefully acknowledges the generous support from the CUHK's Improvement Competitiveness in Hiring New Faculties Funding Scheme and the CUHK's Direct Grant Project Number 4055191. The work of Binh Nguyen is supported by Sin- gapore's Ministry of Education grant A 00045595000-00.

---

Published as a conference paper at ICLR 2024
REFERENCES
R.K. Ahuja, T.I.. Magnanti, and J.B. Orlin. Network Flows: Theory, Algorithms, and Applications. Prentice Hall, 1993.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2022.
Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-DPM: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. arXiv preprint 3xiv:2201.06506 2022.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by 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. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for im- proved quality, stability, and variation. In International Conference on Learning Representations, 2018.
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.
Alex Krizhevsky et Learning multiple layers of features from tiny images. cs.toronto.edu, 2009.
Lijiang Li, Huixia Li, Xiawu Zheng, Jie Wu, Xuefeng Xiao, Rui Wang, Min Zheng, Xin Pan, Fei Chao, and Rongrong Ji. Autodiffusion: Training- free optimization of time steps and architec- tures for automated diffusion model acceleration. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 7105-7114, 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow match- ing for generative modeling. In The Eleventh International Conference on Learning Representa- tions, 2022.
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778, 2022a.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2022b.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. arXiv preprint MxXiv:2309.06380 2023.
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, 35:5775-5787, 2022.
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 Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learn- ing stochastic dynamics from samples. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Gaurav Parmar, Richard Zhang, and Jun Yan Zhu. On aliased resizing and surprising subtleties in GAN evaluation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
10

---

Published as a conference paper at ICLR 2024
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with CLIP latents. arXiv preprint arxiv:2204.0612,, 2022..
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Patlern Recognition, pp. 10684-10695, 2022.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, et Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Informa- tion Processing Systems, 35:36479-36494, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference Learning Representations, 2021.
Steven S. Skiena. The Algorithm Design Manual. Springer, 2008.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2266-2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2020a.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2020b.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. Proceed- ings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 3221 1-32252. PMLR, 23-29 Jul 2023.
Hideyuki Tachibana, Mocho Go, Muneyoshi Inahara, Yotaro Katayama, and Yotaro Watanabe. Quasi-TTalor samplers for diffusion generative models based on ideal derivatives. arXiv preprint arXiv:2112.13339, 2021.
Pauli Virtanen, Ralf Gommers, Travis E Oliphant, Matt Haberland, Tyler Reddy, David Courna- peau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, et al. SciPy 1.0: Fundamental algorithms for scientific computing in Python. Nature Methods, 17(3):261- 2020.
Yunke Wang, Xiyu Wang, Anh-Dung Dinh, Bo Du, and Charles Xu. Learning to schedule in diffu- sion probabilistic models. In Proceedings of the 29th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 2478 2488, 2023.
Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan, Learning to efficiently sam- ple from diffusion probabilistic models. arXiv preprint arXiv:2106.0388 2021.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: 130600353,,, 2015.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2022.
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. DPM-solver-v3: Improved diffusion ODE solver with empirical model statistics. In Thirry-seventh Conference on Neural Information Pro- cessing Systems, 2023.
11

---

Published as a conference paper at ICLR 2024
A
DETAILS OF EXPERIMENTS
We use the following checkpoints that are downloaded from the GitHub folder 2:
CIFAR-IO: at iteration 800,000;
CelebA-HQ: iteration 1,000,000;
LSUN-Church: at iteration 1,200,000;
L.SUN-Bedroom iteration 1,000,000;
AFHQ-Cat: at iteration 1,000,000.
The pretrained models are finetuned in 12,000 iterations. One iteration is the passing and back- propagation process for batch including 15 samples. Due to the similar cost of training between finetuning methods, we report the average GPU hours consumed on each pretrained model up to 12000 iterations, using NVIDIA RTX A5000.
. CIFAR-10: 3.56 training hours.
CelebA-HQ: 10.35 training hours.
LSUN-Church: 13.43 training hours.
LSUN-Bedroom: 14.23 training hours.
AFHQ-Cat: 9.30 training hours.
With this limited budget of resources, the proposed method, BOSS, achieves significantly better performance than other methods terms of FID score. The value N Equation (5) and K max are fixed at 100, and 100 in all experiments if not mentioned. This setup N 100 means only use one batch sampling to calculate the truncation errors between timestamps in equation (5). It is worth noting that this setup highlights the low-cost and rimitee-resourees requirement of our proposal.
FID Calculation. All the FID metrics of related works either cited from previous baselines or are calculated (when there are no such figures reported) based on the Clean-Fil (eennFFID paper (Parmar al., 2022), which unifies the FID calculation to make a fair comparison between papers. These four datasets were downloaded following the instructions from their original papers. We then create the stats file by the clean-fid project. The FID score calculated based on 50,000 generated images and the stats file.
B DESCRIPTION OF THE DYNAMIC PROGRAMMING ALGORITHM
This section presents the pseudocode in Algorithm for the practical implementation of the dynamic programming algorithm designed to determine the Bellman optimal stepsizes. The algorithm takes a matrix, denoted as as input, where Cjk is computed using Equation (5) and the specified number of function evaluations (NFEs). The most resource-intensive aspect of this code is the nested loop responsible for calculating x(j, k). incurring a time complexity of O((Imax)2 x K). The choice of max is crucial, aiming for the Euler sampling method with Kmaa stepsizes to precisely replicate the trajectory of the Ordinary Differential Equation (ODE). Typically, Kmax falls within the range of 100 to 1,000, ensuring accuracy. Given this range for Kmmax and the variable K ranging from 2 to 1,000, the algorithm executes within milliseconds in all scenarios.
Phttps:// thub. cow.gbobooo//ooboooocoooo
htttts////ggiibb...//o/oottt/t// tab//kct/ FiedFlow/
12

---

Published as a conference paper at ICLR 2024
Algorithm 1 Minimum Cost Path Computation
Input: Cost matrix c(j.k) k)
cj k O............................0 Kmax) the target NFEs K.
Set k(i, j) <- +00
40 < Kaax .0 j K
Set k(i, 1) c(i, Kmmax) v0 i K Kmax
for = 2 to K do
for 0 to Kmaax_
do
for = ++ to K -111 do
k(j, k) < min k), c(j. m(ii, 1)}
end for
end for end for
Initialize vi + [0]. w f- 0
for A K to do
for j = w to K mmax do
if n(w, k) c(w, j) k(j, 1) then Append j to ab and set w +- j break end if end for end for
Append K-max th.
return 1/, 10(0 K)
C EMPIRICAL ANALYSIS ABOUT BELLMAN OPTIMAL STEPSIZES
C.I A COMMON TREND IN BELLMAN OPTIMAL STEPSIZES FOR PRETRAINED MODELS ON DIFFERENT DATASETS
We plot the Bellman Optimal stepsizes in Figure 6. It shows common trend of sampling with small initial steps and then larger stepsizes for intermediate iterations. At K = 6, we can observe that the last step is smaller than the penultimate step, hinting that the sampling process aims to take smaller final steps to refine the output. This refining trend is more evident for K 8.
CIFAR10
CelebA .87
.28
.72
.13
K =
12
311
.35
.05 .17
.32
.46
K 4
H
.3
.17
.21
.22
.18
.03
.17
.26
.28
.18
K 6
H
.08
time
time
Figure 6: Optimal Bellman stepsizes for CIFAR-10 and CelebA-HQ K 4 and 6. One identify a common pattern of smaller steps at the beginning and the end of the sampling procedure
C.2 EMPIRICAL EVIDENCE FOR THE STEPSIZES TREND
This section aims to experimentally explain the trend of smaller stepsizes at the beginning and the end of the sampling procedure. Given velocity network vo, empirically estimate the curvature using the following procedure:
1. Sample N noise instances zf i.i.d. from a Gaussian distribution (we set N 1000).
2. Forward each noise instance using to obtain of for t 1, L,..., Kmax
3. Compute the local curvature for each trajectory:
2/+1 x
Curvf = =|12;
for = 2,..., max 1.
13

---

Published as a conference paper at ICLR 2024
4. Calculate the average curvature:
Curv Σ
Curv. == 2.... Kmaax 1.
Subsequently, we plot Curv, and investigate whether the Bellman stepsizes coincide with the straightness of the curve, demonstrated in Figure 7. The curvature of all pretrained rectified models is significant at timestamps close to zero (near the space of noises) and one (near the space of real images). This observation indicates variant stepsize trajectory, notably at initial and final timestamps. It can be intuitively explained that additional steps are required to match the high- curvature region accurately. Consequently, the outcome of our proposal is sensible, as it correctly identifies the high curvature levels at the beginning and end of the sampling trajectories.
16.-5
1e.11
Bellman stapsizes
2.9
Ballman stepsizes
1.4
stepsizss
Cary
Coffe
1.2-
Curts
1.5
1.0
0.8
1.5
0.6
0.4
0.2
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
(a) CIFAR-10
(b) CelebA-HQ
(c) AFHQ-Cat
In-2
1e-2
Bellman stepsizes
Beliman stepsices
Curv:
Curv:
1.5
1.0
6.5
0.0
01
0.0
0.2
0.4
0.5
0.8
0.0
0.2
0.4
0.5
0.5
1.0
(d) -.SU--Church
(e) L.UN-Bedrroom
Figure 7: Curvature measurement (blue curve) and Bellman optimal timestamps (black horizontal line) for different datasets. We observe that the timestamps are denser at regions with higher esti- mated curvature.
C.3 FINDING OPTIMAL STEPSIZES IS A LIGHT- WEIGHT PROCESS
In this section, elaborate further on the efficiency of our framework. This efficiency arises because we only need to pass a single batch of noise through the forward process to obtain the values for each intermediate timestamp. Subsequently, we calculate the local truncation error for any two timestamps. These local truncation errors have in total Krmax x (Kreax 1)/2, and they can efficiently stored without requiring large memory. The dynamic programming involved in this process is also time-efficient, as elaborated in the Appendix B. For added credibility, we provide a running time the entire stepsizes calculation process with NFE = 10, detailing the running time of each component across all our datasets in Table 3. All running times for the Nvidia A5000, an old-gen aid-generation GPU launched in April 2021. The whole process takes around 115 seconds to complete for the 256x256 datasets.
14

---

Published as a conference paper at ICLR 2024
Table 3: Running time (seconds) for different datasets
CIFAR-10
CelebA-HQ
LSUN-Church
BSUN-Bedroom
AFHQ-Cat
One batch forward
8,5977
47.7143
47.4024
48,4324
45,5489
Local truncation
10.066
67.9143
67.5273
65,9875
66.7532
Dynamie programming
0.0134
0,0138
0.0138
0.0139
0.0138
The whole process
18.6771
115.6424
114,9435
114,4338
112.3159
Table 4: Bellman Optimal Stepsizes for = 2 and
The total sum of stepsizes equals one.
K
CIFAR-10 CelebA-HI
[0.28, 0.72
[0.12, 0.22, 0.31,
00.13, 0.87
00.05, 0.17, 0.32, 0.46
LSUN-Church
[0.20, 0.80]
[0.08, 0.21, 0.34, 0.37
Table 5: Bellman Optimal Stepsizes for K = 6 and K = 8. The total sum of stepsizes equals one.
=
K
CIFAR-10 CelchA-HQ LSSUN-Currch
0.09, 0.13, 0.13, 0.21, 0.18)
00.07, 0.09, 0.12, 0.14 0.17.
0.03, 0.08, 0.17, 0.26, 0.28, 0.18
0..02, 0.05, 0.11, 0.17, 0.21 0.20, 2.2220,00100 0.16, 0.08
0.06, 0.15, 0.24, 0.26, 0.20, 0.09)
00.44, 0.15, 0.17, 0.18, 0.17, 0.13, 0.07
The Bellman steps being far from uniform also that the probability path of the pretrained models far from straight, and performing the straightening operation would be beneficial.
We report in this section the Bellman optimal stepsizes obtained in Section 3.3. We want focus the case K = 8 to see the common trend of sampling step sizes. Initially, the sampling process takes small step sizes, possibly for structural determination the images. The stepsizes become larger for the intermediate steps. The last two stepsizes show decreasing trend: the sampling process takes small stepsizes the end to refine and potentially make the final output less noisy.
D THE TRANSFER OF OPTIMAL STEPSIZES ACROSS DATASETS
To verify the generalization of optimized stepsizes, we transferred the optimized stepsizes from ASUN-Chured to the pretrained models on CelebA-HQ and tained with 4, 6, and NFEs for CelebA-HQ resulting from this transfer are presented in Table 6.
The FID scores ob-
Table FID scores for CelebA-HQ with different methods and NFEs
Method
NFEs
6 NFEs NFES
Uniform Euler
158.95
127.01
109.42
Bellman Fuler
92.03
72.54
49.80
Dellman-transfer
132.04
100.68
72.88
Uniform Euler uses uniform step sizes, Bellman Euler uses optimal stepsizes for CelebA-HQ, while Bellman-transfer uses the stepsizes taken from LSUN-Church. Bellman Buler still the optimal method. However, what is important here is that Bellman-transfer is better than Uniform Euler. This hints that there is a certain degree of transferability of the step sizes.² In the empirical realm, we can attribute this transferability to comparable curvature pattern exhibited by pretrained rectified models, as discussed in Section C.2.
"This empirical claim; we do not impose any theoretical claim.
15

---

Published as a conference paper at ICLR 2024
Table 7 is for the that the stepsizes have a certain degree of transferability. Nevertheless, optimizing the stepsizes using Bellman Euler would still obtain the best performance.
dataset. We observe the same trend here, empirically confirming
Table 7: FID scores for LSUN-Bedroom with different methods and NFEs
Method
NFE
6 NFE
8 NFE
Uniform Euler
84.35
39.19
32.15
Bellman Euler
61.60
35.35
25,80
Bellman- transfer
70.23
38.01
29.14
E LOW-RANK ADAPTATION FOR STRAIGHTENING
In this section, we expertiment with adding low-rank adaptation to the linear and convolutional layers of the velocity network vo For instance, the 1th linear layer represented by an m n matrix We is adapted to
W. = We+ + A.B,
where A, is x r matrix, and B, is an x matrix. The value r < min(m,, n) represents the rank of the adaptation. This adaptation is similarly applied to convolutional layers, with a slight adjustment: convolutional layer is first reshaped into a two-dimensional matrix before incorporat- ing the adaptation term. We keep all original parameters A, and B, matrices during the straightening process. We have four versions of straightening named LoRA-r, where r is chosen from the set {1, 4, 16, 64). Their FID scores on the CelebA-HQ dataset over training iterations are plotted in Figure 8, while the number of trainable parameters for each method is presented in Table 8. The FULL-RANK variant described in Section 4 is the method that finetuning all parameters of the original model. is noticeable that versions LoRA-4, LoRA-16 16, and LoRA-64 can almost match the FID score of the FULL-RANK version, which is 33.86. Specifically, 175,000 training iterations, LoRA-4, LoRA-16 and LoRA-64 achieve FID scores of 36.70, 34.52, and 34.20 respectively. Furthermore, LoRA-4 only finetunes 2% of the parameters of the original model but still achieves competitive results compared to the FULL-RANK version.
the models fixed and only update the
FID vs Number of Training Steps
120
LoRA-1
LoRA-4
100
LoRA-16
LoRA-64
FULL-RANK
FID
80
60
0
25
50
75
100
125
150
175
Number of training iterations (10²)
Figure 8: The FID score of straightening procedures along the training iterations on the CelebA-HQ dataset.
16

---

Published as a conference paper at ICLR 2024
Table The number of trainable parameters of straightening procedures compared the full-rank straightening on the CelebA-HQ dataset
Method
Number of Trainable parameters
Percentage on the FULL-RANK version (%)
LoRA-1
330,562
0.5
LoRA-4
1,322,248
2.02
LoRA-16
5,288,992
8.07
LoRA-64
21,155,968
32.26
FULL-RANK
65,574,549
100
F ADDITIONAL QUALITATIVE RESULTS
K 4
K =
K 15
Bellman
Euler
shut Muutterster stock
shut
(ours)
Uniform
K = 207
Euler
shute Chutersstoe
stock
Bellman Heun (ours)
shutoo Sassst
shut chuterstnc
RK45
LIBASTO
Uniform
shut
Henm
Figure 9:
Comparison between images generated from an identical noise with different sampling methods and the number of stepsizes the LSUN-Churd dataset.
17

---

Published as a conference paper at ICLR 2024
=
K =
1
Bellman Euler
(ours)
Uniform
K=
Euler
Bellman
Heun
RK45
(ours)
Uniform Heun
Figure 10: Comparison between images generated from an identical noise with different sampling methods and the number of stepsizes the L.SUN-Bedrod dataset.
18

---

Published as a conference paper at ICLR 2024
TTO
(a) BOSS (NFE = 6, FID = 17.3)
(b) RK45 (NFE = 203.4, FID = 11.4)
Figure 11: Comparative qualitative outcomes BOSS with NFE 6. The image on the right showcases the generated images referenced by RK45.
Figure 12: Uncurated images generated from the model finetuned by BOSS (NFE 10, FID=13.89)
19

---

Published as a conference paper at ICLR 2024
(a) Uniform stepsizes with 4 NEE
(b) Bellman optimal stepsizes with NEE
(c) BOSS with 2 NFE
(d) BOSS with 4 NFE
(e) BOSS with 10 NFE
(1) RK45 with 211.2 NFE
Figure 13: Comparative qualitative outcomes of BOSS with different NFEs on the LSUN-Bedroom dataset.
20

---

Published as a conference paper at ICLR 2024
(a) Uniform stepsizes with 4 NFE
(b) k-step reflow with 4 NFE
(c) Bellman optimal sizes with 4 NFE
(d) Uniform Reflow with 4 NFE
(c) RK45
(f) BOSS with 4 NFE
Figure 14: Samples from LSUN-Bedroom. All corresponding samples use the same initial noise.
21


Reviews:
[{"review_id": "KHH8rQEmAW", "reviewer": "Reviewer_TqqW", "summary": "The author introduces the Bellman Optimal Step-size Straightening (BOSS) technique, a method for distilling flow-matching generative models that enhances the efficiency of image sampling within computational budget constraints. BOSS utilizes dynamic programming to optimize step sizes in a pretrained network and refines the velocity network to straighten generation paths. In this paper, the proposed method has been extensively evaluated on image generation tasks, showing that it significantly improves resource efficiency and maintains high image quality. This approach serves to reconcile the intensive computational demands of flow-matching models with low-resource availability, contributing to the sustainable development of artificial intelligence by reducing computational expenses and environmental impacts.", "questions": "Could the authors extend their testing to include a broader range of datasets and perform a comprehensive comparative analysis with current state-of-the-art methods to better highlight the efficiency and quality improvements of the BOSS method? \nCan the paper provide a detailed comparison of resource requirements, such as memory usage, power consumption, or processing time, to evaluate the practicality of BOSS for users with limited computational resources? \nHow does the computational intensity of the dynamic programming algorithm and the network retraining requirement impact the method's applicability in resource-constrained environments?\n Additionally, has the scalability of BOSS been assessed in relation to increasing dataset sizes or complexity, and if not, would the authors consider evaluating this to enhance the paper's comprehensiveness?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The BOSS (Bellman Optimal Step-size Straightening) method presents an innovative two-phase approach for adapting pretrained flow-matching models. This paper illustrates that BOSS can straighten the velocity network with approximately 10,000 retraining iterations, which marks a significant improvement in efficiency compared to standard practices. Consistently, BOSS achieves lower FID scores in the task of unconditional image generation across a variety of datasets, suggesting superior image quality relative to competing methods. Moreover, the paper introduces a distinctive methodology for calculating optimal sampling step sizes through dynamic programming, thereby increasing the sampling process's efficiency.", "weaknesses": "Although the paper demonstrates significant improvements in image quality and efficiency, its testing is concentrated on specific datasets. A more comprehensive comparative analysis with current state-of-the-art methods would elucidate the advancements BOSS provides, especially in efficiency and quality. Additionally, while the paper addresses low-resource scenarios, it lacks a clear comparison of resource requirements such as memory usage, power consumption, or processing time, essential for evaluating BOSS's practicality for users with limited computational resources. Furthermore, despite the method's enhancements over existing approaches, the computational intensity of the dynamic programming algorithm and the network retraining requirement could limit its utility in resource-constrained settings. The scalability of BOSS, with respect to increasing dataset sizes or complexity, is also not addressed, and an evaluation of this aspect would greatly benefit the paper's comprehensiveness.", "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": "VMwsMdacGj", "reviewer": "Reviewer_pSht", "summary": "Flow matching is a powerful framework for generating high-quality samples, especially in image synthesis. However, the computational demands of these models pose challenges in low-resource scenarios. This paper introduces the Bellman Optimal Step-size Straightening (BOSS) technique for efficient image sampling within a computational budget. BOSS optimizes step sizes and refines the velocity network to improve generation paths. Experimental evaluations demonstrate BOSS's effectiveness in resource utilization and image quality. It provides a sustainable solution that reduces costs and environmental footprints.", "questions": "as above", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Not only in flow matching, but also in diffusion models, all endeavors are focused on predicting the efficiency of noise, score, and vector field. While most of them set the step size based on a heuristic principle (e.g., DDIM, DDPM), this paper explores generative models of ODE from a different perspective - the step size itself, rather than the \"direction\" of a specific step.", "weaknesses": "- The baseline is too weak, as it only compares quantitatively with the fixed-step size ODE solver euler, while it doesn't compare with adaptive step size ODE solvers such as dopri5 or rk45. This comparison should be included in Table 1.\n- The optimization needs to be conducted on a case-by-case basis. Additionally, the use of the dynamic programming algorithm using Gorubi may result in slow performance, rendering the method impractical. Optimization time is also not discussed in this paper.", "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": "xaD0l3jWvX", "reviewer": "Reviewer_GU1Q", "summary": "This paper proposed a new sampling strategy for flow-matching (flow-reflected) generative models. Formally, the authors analyze the sampling error estimation and propose to formulate the sampling process as a dynamic programming problem. Then the task is addressed by the bellman optimal schedule. Moreover, the authors proposed to re-align the velocity network along the accumulated sampling errors, which enjoys good performance.", "questions": "Although this paper enjoys interesting idea, it suffers from incomplete related works, experimental results, analysis, and methodology/implementation details, which largely limit the quality of this paper. I think this paper needs a major revision before the publication.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The authors clearly present the issue of sampling error in flow-matching generation.\n2. It is interesting and intuitive to formulate the sampling schedule as a dynamic programming problem. The solution is convincing.\n3. This paper is well-written with clear motivations.", "weaknesses": "1. Although the main idea of this paper is interesting, this paper gives me the initial impression of being incomplete, with this incompleteness of the presentation of both the methodology and experimental sections.\n2. For method Sec.3.3, the authors did not provide enough details about how to address the dynamic programming (DP) problem in practice. From the experiments, different datasets should share different sampling schedules. Moreover, how many samples are used to calculate the optimal scheduler? What will the optimal scheduler be like, if using different initial noises? The analysis of the generalization of the optimal scheduler selected by DP is very important.\n3. Missing necessary discussions about diffusion sampling schedules in related works (such as [1]). Moreover, the related works are too simple, and only include some matching-flow methods.\n3. Insufficient qualitative comparisons: the authors only provide qualitative comparisons of human faces.\n4. No details about the velocity network re-alignment training are provided.\n\n[1] Karras T, Aittala M, Aila T, et al. Elucidating the design space of diffusion-based generative models[J]. Advances in Neural Information Processing Systems, 2022, 35: 26565-26577.", "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 (poster)