Constant Acceleration Flow

Summary

The paper presents an extension to rectified flow by Liu et al., 2023, a method based on optimal transport that can match samples from two distributions. This class of methods may have different uses but is often employed for the efficient generation of new samples and for morphing one sample to another. As such rectified flow and the extension proposed in this paper may be seen as finding efficient trajectories for use during generation by diffusion models. The method solves the transport flow in a simple Euler approximation of the ODE describing the flow. The paper identifies crossing of flow trajectories as a major obstacle in generating efficient paths in rectified flow and proposes constant acceleration flow (CAF). The paper proposes Initial Velocity Conditioning to produce the initialize the acceleration field in solving the ODE. Rectified flow uses constant velocity to solve the ODE using a reflow procedure to optimize the paths in a second step. Similarly, the proposed CAF also uses reflow to improve the initial paths. The paper reports experimental evaluation of the model with synthetic distributions and with CIFAR-10. The experiments demonstrate improved flow between distributions for synthetic data and improved FID score on CIFAR-10. The paper also reports an increase in straightness of trajectories.

Strengths

The constant acceleration flow is novel to the best of my knowledge. The paper describes the Initial Velocity Conditioning and related steps to make the approach work on CiFAR-10. The approach is suitably summarized in pseudo-code in two algorithms. The synthetic experiments demonstrate the paper's motivation visually. The ODE framework for the transport and the forward Euler solver makes the steps in the procedure easy to follow. The analogy makes it easy to see that constant velocity and constant acceleration lead to straight trajectories is the initial coupling of the samples x_0 and x_1 is correct. A comparison with many alternative methods is provided in Table 1. The proposed CAF visually leads to improved quality and stability on CIFAR-10 over rectified flow.

Weaknesses

The paper presents arguably a substantial improvement of results obtainable with rectified flow in terms of quality at the price of increased computation. The paper does not attempt any argument why this is significant, i.e., has the improved quality the potential to lead to wider adaptation of rectified flow type methods. The paper has no real-world use case to motivate the work. The claim that experiments with CIFAR-10 at 32x32 pixels is a real-world example seems far-fetched. The initialization of the method and its impact on the result is not well explored. The initial velocity calculation requires sample x_1 drawn from the target distribution. Both, the paper under consideration and rectified flow employ a pre-trained generative model to obtain that sample. The quality of initial samples may impact the results of CAF. The authors should provide evidence to the contrary or provide an experimental evaluation of the influence. There is no theoretical justification why constant acceleration is superior to constant velocity flow generation. While the experimental evaluation appears clear, it remains unclear if this behavior is due to the tasks considered or holds more broadly. The paper claims reflow as a contribution but it has been proposed by rectified flow. The same is true for the measurement of straightness of flow trajectory with Normalized Flow Straightness Score (NFSS) in Eqn 12. The paper does not discuss how CIFAR-10 is used. One can speculate based on earlier work such as GLOW that unconditional generation and conditional generation may refer to use of class conditioning but a concise description of the experimental setup is required. The use of a pre-trained generative model should also be taken into account in the required computational effort during training. It may be application-dependent if such a model is available or if it need to be trained.

Questions

A discussion of the use of rectified flow on real-world tasks would enhance the justification for the proposed CAF. Can such examples be provided and discussed? Are there other options to provide the initial velocity field rather than using a pre-trained generative model? The paper is generally well written. The following are minor typos: l. 20 a huge computation burden -> a large computational burden l. 128 thet -> the l. 176 Experiment -> Experiments or Experimental Evaluation l. 188 128-dimensional dimensions

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations section should be moved from the appendix into the main paper as the computational impact of CAF over rectified flow is essential.

Summary

This work proposes Constant Acceleration Flow (CAF), which, instead of learning a velocity-based flow model like in Flow Matching, jointly trains an initial velocity model together with a constant time-dependent acceleration model. This framework aims to learn straighter paths and thus enable better results for few-step generation. Additionally, they propose to condition the acceleration model on the initial velocity and propose to improve the initial velocity using a reflow procedure. CAF is empirically validated on toy data and CIFAR-10, including ablations on the different elements of the proposed framework.

Strengths

- The Constant Acceleration Flow is a new formulation for achieving straighter trajectories in flow-based generative models derived based on the assumption of constant acceleration. - Good empirical results on CIFAR-10 - Includes important ablations on the velocity conditioning as well as on the magnitude of the initial velocity

Weaknesses

"Reflow" for initial velocity: - The reflow procedure was proposed to straighten the paths of an existing generative model. The authors propose to use a "reflow" procedure for training CAF. However, they don't generate the couplings with a model trained with CAF but instead with a pre-trained EDM model, which is very different from how the reflow procedure was proposed. Additionally, it is unclear how dependent CAF is on this procedure and, thus, from the pre-trained generative model. In the CIFAR-10 ablation, the setting with constant acceleration and velocity conditioning but no "reflow" procedure (and thus no pre-trained generative model) is crucial to answer the question of whether CAF also works without a pre-trained model. If it does not, then I would argue that the proposed framework is more of a new distillation technique rather than a new generative model. This contextualization should also be more clearly explained in the main text. Missing discussion and contextualization of related work: - Acceleration Generative Modeling has been proposed in [1]. While CAF is based on Flow Matching and a constant velocity is chosen, [1] considers Bridge Matching in a stochastic optimal control framework with a changing velocity induced by the acceleration prediction. However, their trained model is also a parameterized acceleration prediction that takes as input the current data point $x_t$, current time $t$, and the current velocity $v_t$ effectively conditioning on the velocity. This framework is very related to the approach proposed by the authors, and the connection should be discussed in detail in the main paper. Additionally, an experimental comparison illustrating the differences between the two approaches could be beneficial. - Coupling preservation through initial velocity conditioning: The authors refer to the reflow approach as existing related work. However, the reflow approach was proposed to straighten paths, not preserve couplings. Moreover, the problem of coupling preservation has been thoroughly analyzed in [2], which proposes to use a source point conditioning, i.e. conditioning on $x_0$. How does Augmented Flow/Bridge Matching (Flow/Bridge Matching with $x_0$ conditioning) compare to CAF? As mentioned, velocity conditioning is also used in [1]. - A possible way to empirically compare to these competing methods would be to extend the ablation study on CIFAR-10 to include [1] and [2]. Minor Weaknesses: - Experiments are only conducted on CIFAR-10. Including at least one more dataset, e.g. 64x64 ImageNet/CelebA/etc, could strengthen the empirical results. [1] Tianrong Chen and Jiatao Gu and Laurent Dinh and Evangelos A. Theodorou and Joshua Susskind and Shuangfei Zhai. "Generative Modeling with Phase Stochastic Bridges". In ICLR 2024. [2] Valentin De Bortoli and Guan-Horng Liu and Tianrong Chen and Evangelos A. Theodorou and Weilie Nie. "Augmented Bridge Matching". In Arxiv 2023.

Questions

- As mentioned above, how dependent is CAF on the pre-trained model? Does it also work well without any pre-trained model? - Does the 2-Rectified Flow comparison for the CIFAR-10 experiments also use generated samples by EDM for its first training round? - Can the velocity and acceleration models be trained separately? E.g. first train the initial velocity model and then train the acceleration model?

Rating

4

Confidence

4

Soundness

2

Presentation

2

Contribution

3

Limitations

- Again, how dependent is CAF on the pre-trained model? - Training both an acceleration and a velocity prediction model effectively doubles the amount of total parameters. This could be mentioned as another limitation.

Summary

This paper develops a new flow based generative model whose vector field is constructed similar to rectified flows (straight path between source and target samples), but instead of using a path with a constant speed, uses a path with constant acceleration. This requires learning a neural network to parametrize the initial velocity field and another neural network to parametrize the acceleration field. To improve their model's performance, the authors also propose parametrizing the acceleration neural network using the initial velocity and also applying reflows to the initial velocity. The proposed method helps mitigate the flow crossing problem at a model level, which should in theory avoid excess reflows to learn non-crossed flows. The experiments section demonstrates the ability of the proposed method to learn high performing generative models on single step generation on the CIFAR-10 dataset.

Strengths

The idea is novel, solves a relevant problem and fits nicely into the current landscape of the flow based generative modeling research area. The paper was easy to read and the method itself is simple enough to understand and implement after a single pass over the paper. The main contribution of the paper, in my opinion, is the conditioning of the acceleration neural network on the initial velocity. This gives the final trained model a non-markov path measure, which is something that is of current interest to the field (Augmented Bridge Matching, De Bortoli et al., 2023).

Weaknesses

- The related work should include references to building diffusion models with non-markov path measures (see De Bortoli et al., 2023 and its related work). - For completeness, the paper should include proofs to show that the learned distribution is the same as the data distribution. This follows directly from interpreting Eq. 2 as a conditional flow matching objective (Lipman et al., 2023). - In the initial presentation of the vector field $a(x_t,t)$ at the start of section $4.1$, it would make more sense to drop the dependence on $t$ to emphasize that the ground truth acceleration field is constant in time. - The scope of the empirical evaluation is limited. The authors should have evaluated on more datasets than just CIFAR-10. There are other small scale image datasets, like ImageNet-32, that could have been evaluated.

Questions

- Is there a reason why algorithm 1, line 6, updates $\theta$ before computing the loss for $a_\phi$?

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations.

Rebuttal well-addressed my concerns

Thank you for the detailed rebuttal. Based on the explanations and especially the demonstrations of additional applications, I have revised my recommendation to weak accept.