Towards Hierarchical Rectified Flow

Authors: Yichi Zhang, Yici Yan, Alex Schwing, Zhizhen Zhao
Primary area: generative models
Keywords: Generative Model, Flow Matching, Rectified Flow

Abstract:
We formulate a hierarchical rectified flow to model data distributions. It hierarchically couples multiple ordinary differential equations (ODEs) and defines a time-differentiable stochastic process that generates a data distribution from a known source distribution. Each ODE resembles the ODE that is solved in a classic rectified flow, but differs in its domain, i.e., location, velocity, acceleration, etc. Unlike the classic rectified flow formulation, which formulates a single ODE in the location domain and only captures the expected velocity field (sufficient to capture a multi-modal data distribution), the hierarchical rectified flow formulation models the multi-modal random velocity field, acceleration field, etc., in their entirety. This more faithful modeling of the random velocity field enables integration paths to intersect when the underlying ODE is solved during data generation. Intersecting paths in turn lead to integration trajectories that are more straight than those obtained in the classic rectified flow formulation, where integration paths cannot intersect. This leads to modeling of data distributions with fewer neural function evaluations. We empirically verify this on synthetic 1D and 2D data as well as MNIST, CIFAR-10, and ImageNet-32 data. Our code is available at: https://riccizz.github.io/HRF/.

Decision: accept (poster)
Reviewer scores: 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730132968001

thread_type:
official_review

<thread_object>
creator: reviewer_4etx

<content>
summary:
Flows $p_t$ transform one distribution $p_0$ into another $p_1$ via a corresponding vector field $f(x_t,t)$ that satisfies its continuity equation. The flow is typically defined via conditional flows, for example $p_t(x|x_0, x_1)\sim t(x_1)+(1-t)x_0+\sigma\epsilon$, where one can take the limit $\sigma\rightarrow 0$. The corresponding flow is simply $p_t(x)=\int p_t(x|x_0, x_1) \pi(x_0, x_1)dx_0 dx_1$. The conditional vector fields $v_t(x_t|x_0,x_1)$ that satisfy the continuity equation of such flows are $x_1-x_0$, while the unconditional vector field is $f(x_t,t)=\int v_t(x_t|x_0,x_1) p_t(x_0, x_1|x_t) dx_0 dx_1=\int v_t(x_t) \pi_t(v_t|x_t) dv_t$.

This paper studies the possibility of using a sample of $\pi_t(v_t|x_t)$ at each integration step during data generation, instead of using the mean $f(x_t,t)$ which is akin to using stochastic gradient descent instead of gradient descent. . The result is a stochastic sampling process, which by using a derivation of $\pi_t(v_t|x_t)$ is proven to have marginals that coincide with the flow $p_t$.

The distribution $\pi_t(v_t|x_t)$ is itself modeled using the methodology of rectified flow matching, and the methodology is applied to synthetic and high dimensional image data.

strengths:
The paper is generally well written, even though a few parts could be improved.

Theorem 1 and 2 are interesting results, and provide insights about the conditional distribution of velocities at each current point $x_t$ and they enable using higher order stochastic sampling.

weaknesses:
My **main** concerns are the following:

1. The proposed method requires taking the position and the current velocity as input in order to predict the acceleration. The authors mention expanding the Resnet for their framework and increasing the amount of data processed. Can authors provide a detailed table comparing HRF and RF models, including parameter counts, training time per iteration, and memory usage across all experiments?

2. It is not clear if the NFE includes the steps required to sample the velocities. The paper should report both L and J in each experiment. Also the compute time for each step should be reported and compared with RF.

3. The results for CIFAR-10 unfortunately are not encouraging. Considering the results on MNIST and CIFAR-10 it seems that the proposed model does not scale as well. This is why it is important to also test the model on Imagenet.

4. The resulting model is not a diffeomorphism and as such we lose the ability to perform density estimation.

In my opinion the main contributions of the work are theoretical, and applying the proposed methodology in practice is challenging.

questions:
1. What are the parameter counts for both HRF and the baseline HF? Also what is the training time per iteration?

2. Is the reported $NFE=J\cdot L$? What is the time required per NFE in both HFR and RF?

3. The paper claims that the resulting paths are straight as the trajectories can intersect. However the trajectories are stochastic, and thus they are likely to increase overall length. I would appreciate a clarification from the authors regarding this matter.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730547417934

thread_type:
official_review

<thread_object>
creator: reviewer_cauw

<content>
summary:
The work enriches the concept of rectified flow for generating data from known base distribution to data distribution. Compared to the base approach, hierarchical rectified flow formulation models the multi-modal random velocity field and acceleration field, which leads to generating more interesting trajectories. The authors evaluate their approach on some low-dimensional use cases and small image benchmark data.

strengths:
***Clarity and Accessibility***: The paper is well-written and presented in a logical manner, making the concepts and methodology easy to follow. Complex ideas are explained in a way that is accessible to readers with varying levels of familiarity with flow-based generative models, facilitating broader understanding and engagement with the research.

***Clear and Impactful Motivation***: The authors provide a compelling motivation for developing hierarchical rectified flow models, addressing limitations in existing flow-based generative models. This work has the potential to significantly impact the field by introducing a more flexible and expressive approach, which could inspire further research and applications in generative modeling.

***Solid Theoretical Foundations***: The paper presents strong theoretical foundations that support the proposed hierarchical model. These well-motivated considerations provide a rigorous basis for understanding how the model improves upon previous approaches, making the proposed framework robust and trustworthy.

***Intuitive Experimental Demonstrations***: The inclusion of low-dimensional toy experiments adds an important educational aspect to the study, offering readers an intuitive way to grasp the dynamics of the approach. These experiments clarify how the model handles multi-modal distributions and complex trajectories, offering insights that make the approach more transparent and easier to analyze.

weaknesses:
***Ambiguity in Extended Hierarchical Approach***: Although the acceleration-based approach and its motivation are clear and well-justified, the transition to the extended hierarchical flow model lacks clarity. Specifically, while the training objective for the acceleration-based approach is defined by Equation (8), the relationship to the hierarchical model’s training objective, outlined in Equation (10), is not thoroughly explained. The conceptual progression and the structural specifics of the higher-dimensional source distribution (D-dimensional) are left somewhat ambiguous. Additional details on how this source distribution is constructed and how it interfaces with the hierarchical model would clarify the extension.

***Scalability and Computational Complexity***: While the hierarchical rectified flow method is innovative, its computational demands are significant, particularly when scaling to high-dimensional data. Even when considering latent models, the approach appears to be resource-intensive and potentially impractical for high-dimensional images due to its complex sampling procedure. This aspect could hinder its use in real-world scenarios where efficient sampling and rapid training times are crucial. The authors should provide a more detailed analysis of the computational overhead, comparing training and sampling times with other flow-based generative models to highlight both the benefits and trade-offs of their approach.

***Limited Scope of Experiments***: The experimental evaluation primarily uses small, low-dimensional image datasets, which limits the generalizability and relevance of the results. Given current advancements in generative modeling, such datasets do not fully showcase the potential of the hierarchical approach. Adding experiments on higher-dimensional or more complex data—perhaps by incorporating latent diffusion architectures—would better demonstrate the method’s scalability and practical value. More comprehensive experimentation would help readers assess how well the model performs in scenarios closer to those encountered in modern applications of generative models.

questions:
Please refer to the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730619791114

thread_type:
official_review

<thread_object>
creator: reviewer_test

<content>
summary:
This paper introduces a novel framework, hierarchical rectified flow, to model data distributions. It addresses the limitations of conventional rectified flow (RF), which only captures mean velocity, by incorporating distributional information in the velocity space. The objective, equivalent to acceleration matching, derives target acceleration from the data and prior velocity distributions. The framework can be extended beyond acceleration, creating a hierarchical generation model. The proposed approach offers benefits such as easier path integration and improved results with fewer NFEs.

strengths:
The strengths of the paper are listed below:
- Clear motivation with well-written, easy-to-follow presentation.  
- Experimental results that partially support the theoretical claims.

weaknesses:
The weaknesses of the paper are listed below:
- While the paper's motivation is sound, my main concern lies in the practicality and application of the proposed approach. The generation framework demands significantly more model calls compared to the conventional RF framework, specifically NFEs multiplied by the number of discretizations in the velocity space. Although RF only targets mean velocity, it delivers strong empirical results with far greater efficiency than the proposed method. Additionally, RF can be rectified multiple times for better results. I recommend that the authors conduct a comprehensive comparison between the rectified RF and the proposed method. This inefficiency could lead to significantly limited applicability, especially for larger datasets like ImageNet. Could the authors include a comparison of generation performance and training/inference time between the rectified RF and the proposed method on the existing datasets, or ideally on ImageNet?

- The inefficiency becomes even more pronounced with increased depth; for example, a D-depth model requires a substantially larger architecture compared to a one-depth model.  Does the forgetting problem of the optimal trajectory occur with increased depth, necessitating an expanded architecture? Could the authors provide an ablation study on how model performance and computational requirements change as depth increases?

- The paper lacks a comprehensive comparison of efficiency, including training and inference time, against RF and rectified RF. The paper should include details of the neural network architecture to ensure reproducibility in toy settings. The paper should include more details of the neural network architecture and experimental settings to ensure reproducibility in toy scenarios.

- On the theoretical side, does increasing depth make it harder for the model to converge during training? It would be helpful to include empirical evidence of convergence rates for models with different depths or a discussion of any theoretical bounds on convergence that might exist when using larger depths.

- Minor typo on line 156: it should be \pi_0. Please review and correct all typos.

Note: The review has been edited to make it more actionable for the authors, as suggested by the PCs.

questions:
Please see the Weaknesses section for my concerns and questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732092699854

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Changes

comment:
We thank all reviewers for their time and feedback. We are excited to see that our work was assessed as well-written (R1, R2, R3), with clear and impactful motivation, solid theoretical foundation, intuitive experimental demonstration (R2), and interesting theorems (R3). We updated the paper and appendix to answer reviewer questions. Additions include: 1) the density estimation in Appendix D, 2) more details on the relationship between the objectives given in Eq. (8) and Eq. (10) in Appendix E, 3) details on the models in Appendix F, 4) ablation studies in Appendix G, 5) adopting minibatch optimal transport into HRF in Appendix H.2, and 6) additional results on MNIST and CIFAR-10 in Appendix H.3. We also clarified that **our comparison is fair** because the NFEs reported in the paper are the total NFEs, i.e., the product of the number of integration steps at all HRF levels (L349 of the revised paper and updated figure axis). We answer questions for each reviewer individually and point to the corresponding paper and appendix additions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes the Hierarchical Rectified Flow that allows to model multiple ODEs instead of a single ODE in the regular Rectified Flow. It allows to better represent multi-modal distributions, e.g. for physical simulations. The method also allows to reduce the number of neural function evaluations and considerably speed up the forward pass. The paper has solid theoretical foundations, supported by experiments. Among weaknesses, the experiments primarily use small, low-dimensional image datasets, and it is not clear how the model will scale to larger modalities.