Improving the Training of Rectified Flows

Summary

By retraining with Rectified flows, they straighten the ODE, allowing sampling for small number of steps. They propose using rectified flow as replacement to complex distillation methods such as consistency models. They reflow only once to straighten the path which makes it as efficient as distillation. They propose some propose using a different distribution on t and a LPIPS loss. They obtain better performance than the original rectified flow and competitive results with distillation. The authors provide geometric intuition to why one reflow should be enough. They provide well detailed ablation on each aspects of their improvements on multiple datasets. They combine Pseudo-Huber (differentiable Huber loss to get less sensitivity to outliers while being differentiable) with LPIPS (better distance metric than L2 for image quality). They take a weighted combination of both losses. If one was not generating images, we could envision using pseudo-huber with L2, so the method could still be useful in that case. The focus here is on images though. Edit: I apologize, I added some extra text to the summary that was not related. I just removed it.

Strengths

The experiments are on multiple datasets, focusing on the 1-2 NFE case which is arguably the most important problem in diffusion-flow models to solve. The results are quite impressive. Reflow seems more efficient that distillation and its especially much more clean and natural as a solution. The experiments and discussion is quite thorough.

Weaknesses

Honestly this is a near perfect in my view. Maybe, the only thing I could think of is the focus being only on images, so maybe not fully generalizable accross domains.

Questions

The only thing I would like added is the difference in memory from the first flow and then the reflow only because of the use of the LPIPS. Basically, I would like to know if LPIPS massively improve the memory cost over L2 and if so it would be important to mention it as a limitation. The authors already mentioned the 15% speed reduction (which is quite minor), it would be great to know about memory.

Rating

10

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors have well addressed the limitations. I just asked that they talk about the memory cost of LPIPS too.

Summary

The paper targets efficient training of a class of flow models called Rectified Flow (RF) trained using flow matching objective. The paper has two broad contributions: (1) justification of 2-RF (‘reflow’-ed once) being close to optimal, (2) and using those findings to improve training of 2-RF. Authors argued that, in practical scenario, the pairs generated by optimal 1-RF model is ‘crossing-free’, i.e. the stochastic interpolating paths rarely cross each other. They provided an intuitive and empirical evidence of the same. They claimed this to be a motivating factor for two improvements of 2-RF model training — a new time-step distribution $p_t$ & the a new distance measure for the regression task of diffusion objective. Good empirical performance is shown in terms of FID on several datasets and models. An ablation study is also done for the relevant part of the finding.

Strengths

- The problem statement chosen and the arguments provided are quite credible. It is well known that, for RF to work well, one requires several reflows. Reflows are expensive, hence the contributions of the paper (if credible) can be significant. - A good intuitive analysis is done by the authors to justify their argument of 2-RF being near-optimal.

Weaknesses

I like the overall outcome of the paper in terms of empirical performance. I also like the problem statement. However, some concerns remain. Following points explain the issues and points to some more questions in question section. ### Major concerns - This is a major concern for me. The paper has two parts: (1) justifying that 2-RF is near optimal & providing intuitive and empirical evidence; (2) proposing two new measures for training improvement. I think, (1) & (2) both are overall correct on their own. I just don’t think that (1) is the right motivator for (2). - Continuing the above point, I felt that the rationale behind both ‘improvements’ are weakly connected to the empirical observations (in section 3). After all, what the authors ultimately proposed are ad-hoc and exists in the literature. (L162) “.. *focusing on the tasks where training loss is high ..*” is a known technique which authors admitted themselves. Using non-L2 loss is also not unheard of [1]. More importantly, I don’t think (not) using L2 loss has anything to do with your observation is section 3 (as argued in L195). The L2 loss has its origin in score-matching which yields theoretical benefits, but I suspect it is not necessary. Q1 is related to this. - Let’s talk about the observation of section 3 itself (i.e. 2-RF being optimal). Authors must be clear about whether they are making a theoretical assertion (L64-65) or just an empirical observation (they used words like “under *realistic* setting”, “*rarely* intersect each other”). If you are providing a *guarantee*, you must provide better formal proofs. To clarify, I think the argument provided in section 3 is indeed correct and seems reasonable. But if it is a guarantee, empirical evidence is not enough. Q2 is a related question. [1] “Improving Diffusion Models's Data-Corruption Resistance using Scheduled Pseudo-Huber Loss”, Kharpov et al., arXiv 2403.16728. ### Presentations - The notation used in section 2 & 3 are confusing sometimes. There are three pairs of notation — $(\mathbf{x}, \mathbf{z})$, $(\mathbf{x}_0, \mathbf{z}_0)$ \& $(\mathbf{x}_1, \mathbf{z}_1)$. I am confused about which is what. I would recommend authors to follow the same notation as the original RF paper by Liu et al. - Notations like $\mathbf{z}\_0 \sim p\_{\mathbf{x}}$ (L86) are very confusing. - The factor $\frac{t}{1-t}$ in Fig 2(a) (at the top left) should be $\frac{1-t}{t}$, right ? - Fig. 2 caption says $\mathbf{z}'' = \mathbf{z} + (\mathbf{x}' - \mathbf{x}'')$ — where is the factor $\frac{1-t}{t}$ ? Did you assume $t = 0.5$ ? - Again, it is hard to parse notations like $\mathbf{x}, \mathbf{x}', \mathbf{x}''$ or $\mathbf{z}, \mathbf{z}', \mathbf{z}''$. - Eq.4: Can you please denote the suffix of the $\mathbb{E}_{??}\left[ \cdot \right]$ properly ? It is hard to read otherwise. ### Results/Experiments Result section is okay-ish. The following are some comments/questions. - Table.1: Need clarification: ‘Base (A)’ is a 2-RF, and the other ones are written as “(A) + <something>” — does that mean they are 3-RF (meaning 2 reflows) ? - Section 5.2 seems totally unnecessary. That has nothing to do with the core contributions of the paper. - Section 5.3 is also very much unnecessary. I even doubt its correctness. You seem to be proposing a new sampler/solver with a very weak (intuitive) motivation. Designing a solver require a lot more than that. And then “*.. detailed analysis is provided in appendix E*” — appendix E barely has any details ! Also, obvious question, why are FIDs going up (fig.4) with higher NFE ? Does that even make sense ? - Fig.5(b): The inverted noise norm distribution still looks quite different (higher variance) from the true noise. Just having the norm closer to the truth isn’t necessarily making a good case for your method.

Questions

1. Authors said “*any other premetric .. deviates from the posterior expectation*” — but is it really true ? Can you mathematically show why ? 2. Is it possible to have trajectory crossings *at all* when samples are from $p_{xz}^2$ ? If no, can you prove it formally ? 3. L99: “.. use a specific non-linear interpolation” — what is that exactly ? I thought the interpolation is still linear. 4. Eq.4: How did you decompose the loss — can you show the steps ? And what is $\bar{\mathcal{L}}$ ? 5. L65: “*training loss of 2-RF has zero lower bound*” — please clarify: Isn’t it true that *any* L2 loss has zero lower bound ? How does it matter whether trajectories cross or not ? Even the FM loss with independent coupling $p_{\mathbf{xz}}^1$ has zero lower bound — no ? 6. I think it is unclear which quantity the authors are arguing to be zero when trajectories do not cross. The term ‘curvature’ was used some times (L124) — what does that mean ? Can you write this object in mathematical terms ? Just curious, what is the equation that needs to be proved if one wants to show that no trajectories from 1-RF coupling cross each other (ideal case) ? (Related to Q2 above).

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

Some limitation are mentioned, which are reasonable.

Summary

This paper mainly improves the training of rectified flows empirically, making it comparable to the distillation method in terms of performance with fewer steps.

Strengths

1. This article provides a comprehensive analysis of a single return, and there is a clear motivation for improvement. 2. The authors mainly improved the training of the whole network from three aspects: time-step sampling, loss function, and initialization from the diffusion model. 3. The improved rectified flows method can be compared to the sota distillation method, and it also supports operations such as inversion, which may be important in applications such as image editing or translation.

Weaknesses

1. The main weakness of this paper is that all the improvements are incremental and empirical. In addition, many of the improved techniques overlap with existing diffusion improvements, such as Pseudo-Huber loss [1] and Initialization with pre-trained diffusion models [2]. 2. The comparison in Section 5.2 may be somewhat unfair. Although the reflow requires a forward pass for each iteration, the reflow needs to run ODE in advance to generate noise-sample pairs, while the CD does not need to simulate, so the number of forward passes should also be considered to generate noise-sample pairs. [1] Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. [2] Liu X, Zhang X, Ma J, et al. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation.

Questions

See Weaknesses.

Rating

4

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

The authors briefly discuss their limitations.

Summary

The method introduces an one-stage training of rectified flows, mitigating the costly process of multi-iteration training of the former model. Particularly, the authors propose a U-shaped timestep distribution for sampling and modified LPIPS-Huber loss. The method demonstrates superior FID scores for 1-NFE setting on CIFAR-10 and ImageNet $64\times 64$, outperforming other distillation methods like consistency distillation.

Strengths

- Provide an analysis of reflow algorithm and shows that one-stage reflow are sufficient for nearly straight solution trajectory - Improve training of rectified flow via introducing U-shaped time distribution and new loss function, LPIPS-Huber.

Weaknesses

- Have the authors tested with more than two-NFE like 100 to directly compare with 2-rectified flow? Does model still get consistent improvement? - What is the advantage of U-shape time distribution over lognormal distribution? As it is constrained to already hold a straight solution paths like in 2-rectified flow which cannot directly apply to 1-rectified flow. The intuition of this part does not clearly present. - Have the authors tested LPIPS-Huber-1/t (config F) for 1-rectified flow? - New sampler seems to have counter effect when many-NFE (>6) is used. Could the authors explain this behaviour? Overall, the novelty seems to be limited as the method mainly involves engineering works in using U-shape distribution and LPIPS-Huber loss (a combination of huber loss and LPIPS loss) to improve model performance under few-NFE settings. Besides, though the performance shows a notable gain compared to RectifiedFlow and ConsistencyModels, it still falls behinds ImprovedConsistencyModels and CTM. However, I appreciate the findings in paper and these proposed techniques can be useful for the community.

Questions

NA

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

This is included in conclusion.

response

The fact that memory is almost the same with LPIPS is very good. I leave my score as it is.

Gap between real data and synthetic data

Thank you to the authors for the detailed answers. Another concern of mine is the mismatch between real data distribution and synthetic data distribution. Since the method is solely trained on synthetic data, generated by the pretrained model (similar to RectifiedFlow), does the model suffer from this issue? Additionally, the performance of the method still lags behind iCT. I believe that incorporating real data could address this problem and lead to a better FID score. If possible, I would appreciate it if the authors could conduct experiments on a small dataset like CIFAR-10 to test this hypothesis.

Response to author's rebuttal

I thank the authors for providing a response to my queries. Overall, I would say, the rebuttal is, in most parts, quite well written and convincing. Majority of my technical doubt were clarified, and the authors did so quite well. However, I would say the following points: 1. One of my core concerns were the link between the observations and proposed solution. The "_using pre-metric_" part was well explained in the rebuttal and I think I now vaguely understand what they meant. The other one, the "_timestep distribution_" part, I still think is weakly connected. 2. The other concern was about how "formal" the idea is. The authors admitted that it is more intuitive than formal, which isn't quite ideal. However, I do recognize that not everything can be easily formalized. But then you run into the issue of the observation being not true in every possible scenario. I guess that's a fair trade-off. 3. I would still advise against the new "update rule". I think the paper is better off without section 5.3. Showing a new update rule being better only for a one region of the NFE space does NOT make it correct. Prefer correctness over "content" in the paper. Do more study on that and consider submitting separately in future. That's my opinion. Anyway, I think overall the authors made their case well. I am increasing the score. I hope these discussions will reflect in the final version of the paper in some way. Thank you.