Align Your Flow: Scaling Continuous-Time Flow Map Distillation

Dear reviewer dC5b,

We are glad you found our paper well written and our ideas novel. Thank you for highlighting our state-of-the-art empirical results, our extensive baseline comparisons as well as our proofs. Below we will address your questions.

• Insufficient explanation in Sec 3.4: Thank you for pointing out the ambiguities in Section 3.4. We will make sure to clarify this point in the final version of the paper. Please do not hesitate to let us know if you see other areas of the paper where the presentation can be improved.

• AYF-LMD not being competitive in image generation: We appreciate the feedback. Although we did not find AYF-LMD to be competitive in image generation tasks, we did find that it was superior in our low-dimensional 2D settings in the Appendix. As such, the choice of which objective proves to be most beneficial for different applications may differ and depend on the dataset and data type (in the future, our framework may be applied in other domains, such as molecule or 3D point cloud generation, and it is not clear which objective will be better there). This is why we opted to include it in the original paper. This loss is also one of the things that differentiates flow maps from standard consistency models as it is only useful for the former. Nonetheless, we will consider your suggestion and may rearrange the content of the paper in the final version of the paper.

• Additional Text-to-Image Evaluations: We appreciate your comments on our text-to-image experiments. The goal of our text-to-image experiments was mostly to demonstrate the scalability of our distillation approach to expressive and complex image generation tasks. An important reason why we decided on FLUX.1-dev is that AYF assumes a pretrained continuous-time diffusion or flow model, whereas SD1.5 and SDXL were trained with discrete timesteps. Therefore, we cannot directly use these models without additional complications. However, to the best of our knowledge, the only previous text-to-image consistency models are TCD and LCM, which were trained on SDXL/SD1.5, and it seemed appropriate to nonetheless compare to these works (and we outperform them significantly). We believe our AYF is the first continuous-time consistency or flowmap framework that has been successfully demonstrated at the text-to-image scale.

  However, we agree that including additional comparisons between our few-step flow maps against an adversarially trained baseline (FLUX.1-schnell) as well as the original teacher model (FLUX.1-dev) is appropriate. As we are not able to share any images during this review process, we have conducted 3 additional pairwise user studies to compare the different approaches in terms of image fidelity and text alignment. We will include additional qualitative comparisons in the final version of the paper.

  We generate 260 images from each model, using 30 steps for FLUX.1-dev and 4 steps for AYF and FLUX.1-Schnell (this follows the model's default sampling settings described on Huggingface). We then perform pairwise comparisons between each two methods. See the results below (numbers indicate preference votes on image pairs generated by the two methods with the same text prompt):

  FLUX.1-dev vs. AYF (ours)	FLUX.1-dev	AYF (ours)	Tie
  Fidelity	266	229	25
  Alignment	262	238	20

  FLUX.1-dev vs. FLUX.1-schnell	FLUX.1-dev	FLUX.1-schnell	Tie
  Fidelity	261	238	21
  Alignment	253	242	25

  FLUX.1-schnell vs. AYF (ours)	FLUX.1-schnell	AYF (ours)	Tie
  Fidelity	246	249	25
  Alignment	273	223	20

  As we can see from the results, both types of distillation cause some quality degradation compared to the base diffusion model. This is expected, but the base model's better quality comes at a much higher sampling cost (30 steps). However, when comparing our AYF model against FLUX.1-schnell, both methods achieve roughly the same number of votes in terms of image fidelity. However, the FLUX.1-schnell model appears to exhibit better text-alignment, which we hypothesize may be due to the limited number of rendered text examples in our finetuning dataset (this could be easily addressed in the future). We would like to emphasize that our method’s results are achieved by a brief LoRA-based fine-tuning of the base model, whereas FLUX.1-schnell relies on complex adversarial objectives and full fine-tuning of the model[1]. Therefore, we consider our results as highly successful and very promising, achieving such strong performance without any adversarial objectives at all here. However, future work could further boost performance with an additional adversarial fine-tuning phase, which we observed to significantly improve quality without loss of diversity in our ImageNet experiments.

• MeanFlow models: Thank you for pointing out the concurrent MeanFlows work. We agree that this work is highly relevant to ours. The paper explores a loss similar to our AYF-EMD loss to train flow maps from scratch and achieves strong results. We will definitely include it in the discussion for the final version of the paper.

We hope we were able to address your questions and concerns. We will include the above results and discussions in the final version of the paper. Thank you.

References:

1. Sauer, A., Boesel, F., Dockhorn, T., Blattmann, A., Esser, P., & Rombach, R. (2024, December). Fast high-resolution image synthesis with latent adversarial diffusion distillation. In SIGGRAPH Asia 2024 Conference Papers (pp. 1-11).

Dear reviewer mGSb,

We appreciate your feedback and thank you for highlighting our strong empirical results, our detailed training algorithms, as well as our ablation studies and comprehensive comparisons. Below, we will address your comments and concerns.

• Paper not easy to follow: Please do not hesitate to kindly point to areas of the paper with ambiguities, and we will do our best to explain and to clarify in the final version of the paper. Overall, we will carefully analyze the paper and improve the writing where we believe the presentation can be improved.

• ImageNet 256: As our paper is directly related to continuous-time consistency models [1], we followed their experiments and chose ImageNet64 and ImageNet512 to perform our evaluations on. Furthermore, we would also like to point out that many papers in the field report results primarily on ImageNet64 and very small datasets such as CIFAR10. In our case, we wanted to show the scalability of our approach, so we extended our evaluations to ImageNet512 and text-to-image models. Finally, considering that we have extensive evaluations on both ImageNet64 and ImageNet512, there are no additional insights to be gained from training on the same dataset at yet another resolution. And on ImageNet64 and ImageNet512 we already compare to many dozens of baselines (see Tables. 1, 2, as well as 4 and 5 in Appendix). Last but not least, training models on ImageNet to state-of-the-art performance requires substantial GPU resources and can in some cases take a substantial amount of time (hence, this seems beyond the scope of the rebuttal). Therefore, we opted to focus our resources on ImageNet64, ImageNet512 and text-to-image experiments.

• Comparing AYF with training-free diffusion acceleration methods: Please note that our approach is about diffusion distillation which is not a training-free approach. Such distillation approaches are inherently different from training-free acceleration techniques of diffusion models, such as improved ODE solvers like DPM-Solver-3. Specifically, distillation techniques are able to dramatically accelerate the inference of these models by reducing their sampling steps all the way down to 1-4 steps. Training-free approaches such as DPM-Solver-3, on the other hand, are much easier to plug-and-play but still require at least around 10-15 steps for reasonable results. To illustrate this, below we compare our AYF models against DPM-Solver [2] and UniPC [3], which are two of the most widely used fast ODE solver approaches. Note that we chose these approaches because they are easy to add in a plug-and-play manner, unlike methods like DPM-Solver-3 (not 3M), which requires computing model-specific statistics and are harder to use in practice. | Sampling steps | 1 | 2 | 4 | 8 | 16 | 32 | 64 | |------------------------|---------|---------|---------|---------|--------|--------|--------| | EDM2 + Heun (default) | 566.137 | 453.135 | 388.007 | 142.896 | 9.324 | 1.675 | 1.402 | | EDM2 + DPM-Solver-2M | 290.284 | 287.783 | 100.359 | 14.204 | 2.189 | 1.485 | 1.388 | | EDM2 + DPM-Solver-3M | 290.284 | 287.783 | 100.359 | 21.476 | 2.735 | 1.414 | 1.379 | | EDM2 + UniPC-2M | 290.284 | 287.783 | 101.163 | 12.862 | 2.072 | 1.427 | 1.376 | | EDM2 + UniPC-3M | 290.284 | 287.783 | 101.163 | 30.016 | 4.399 | 1.396 | 1.369 | | AYF | 3.32 | 1.87 | 1.70 | 1.69 | 1.72 | 1.73 | 1.75 |

  We see that while these training-free methods help the base model reach near-optimal performance with fewer steps, they still need around 32 steps or more for good results. In contrast, our AYF models significantly reduce the step count, reaching strong FID scores in just 1-4 steps. Note that at 1 or 2 sampling steps, these ODE multistep methods behave identically and they only start to differ when the number of steps exceeds 3. This explains why several entries in the table have the same values. We will add these results as well as a discussion on this, including references to solvers like DPM-Solver and UniPC, in the final version of the paper.

• FLUX vs SDXL: We appreciate your insight on this matter and acknowledge that the base model from which we distill has a significant effect on the final performance after distillation. We would like to point out that the focus of our text-to-image experiments was mostly to demonstrate the scalability of our approach to expressive and complex image generation tasks. An important reason why we decided on FLUX.1-dev is that AYF assumes a pretrained continuous-time diffusion or flow model, whereas SD1.5 and SDXL were trained with discrete timesteps. Therefore, we cannot directly use these models without additional complications. However, the only previous text-to-image consistency models are TCD and LCM, which were trained on SDXL/SD1.5, and it seemed appropriate to nonetheless compare to these works (and we outperform them significantly). We believe our AYF is the first continuous-time consistency or flowmap framework that has been successfully demonstrated at the text-to-image scale. We would also like to point to our extensive experiments on ImageNet64 and ImageNet512, as well as our ablation studies in the Appendix, which unambiguously demonstrate the advantages of our framework over previous discrete-time techniques, which TCD and LCM rely on.

  Finally, please note that we now performed additional comparisons between our 4-step AYF text-to-image flowmap and the teacher FLUX.1-dev model as well as the adversarially distilled FLUX.1-schnell model. Please see our response to Reviewer dC5b for details.

We hope we were able to address your questions and concerns. We will include the above results and discussions in the final version of the paper. If you found our explanations convincing, we would like to kindly ask you to consider raising your score accordingly. Thank you.

References:

1. Lu, C., & Song, Y. (2024). Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.11081.
2. Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. (2025). Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. Machine Intelligence Research, 1-22.
3. Zhao, W., Bai, L., Rao, Y., Zhou, J., & Lu, J. (2023). Unipc: A unified predictor-corrector framework for fast sampling of diffusion models. Advances in Neural Information Processing Systems, 36, 49842-49869.

Dear reviewer jQYj,

Thank you for your feedback. We are glad you found our empirical results compelling, highlighting the practical value of our proposed methods. We are also delighted that you appreciate our visualizations, the mathematical framing, and our analysis of the flaws of consistency models.

Below, we will address your concerns.

• XXL runs: We primarily wanted to show via our experiments the benefits of using flow maps with 4 steps compared to consistency models which excel in 1-2 steps. Simply put, one might wonder given great 1- and 2- step results, why even bother with 4-steps? As such, we aimed to show that smaller 4-step models can be shown to be competitive with, and even outperform, larger 1- or 2-step models. Specifically, we are demonstrating that through distilling small but autoguided models into efficient high-performance flow maps, these small but expressive AYF flow map models can outperform the large baselines. This is validated in our extensive evaluations. In particular, because we use small architectures, the wall clock sampling time of our 4-step models is faster than the sampling time of previous 1- or 2-step sCM [1] models using very large networks. This is why we focused only on the small EDM checkpoints. Moreover, training XXL models to state-of-the-art performance would be extremely costly and slow, potentially taking a week or more, and beyond the scope of the rebuttal. We thank the reviewer for their suggestion, though, and will consider adding EDM-XXL results in the final version.

• AYF-EMD vs. AYF-LMD computation requirements: We would like to point out that our AYF-EMD and AYF-LMD objectives both require essentially the same computations (computing base velocity + computing a JVP operation), but these computations are performed in different orders.

  Specifically, for AYF-EMD, we first compute the base velocity $v_t = v(x_t, t)$ (with guidance) at the current position. We must then run the flowmap $f(x_t, t, s)$ and compute its total time derivative with respect to $t$, i.e. $\frac{\mathrm{d} f(x_t, t, s)}{\mathrm{d}t}$. To do this, we use the JVP operation in pytorch using the base velocity. The output of this JVP operation is used to obtain our regression target (see Algorithm 1 in the appendix). Note that we do not backpropagate through the regression target, which means we do not have to worry about higher-order derivatives.

  On the other hand, for AYF-LMD, we must first run the flowmap $f(x_t, t, s)$ and compute its time derivative with respect to $s$. Similar to before, we do this by using the JVP operation in pytorch. Next, we compute the base velocity (with guidance) at that new point $v_s = v(f(x_t, t, s), s)$, and use that velocity along with the time derivative of s to obtain our regression loss, which is also detached (i.e. also here no backpropagation through the JVP).

  As such, both methods need to compute the base velocity once, and perform a JVP operation. Therefore, they have similar computations and exhibit identical training speed.

• Why destabilization when $r \to 1$? The original sCM [1] paper introduces several training stabilization techniques, one of which is called adaptive double normalization. It applies a normalization to the scale and shift parameters in each AdaLN layer. In our experiments, using this normalization for both the flow map and the teacher flow-matching models results in a significant quality deterioration. As such, we opted to not include this technique in our architecture, and kept the network structure almost identical to the one from EDM2. However, this implied that we needed another means of regularization to stabilize training, leading to our proposal of clamping $r=0.99$. Importantly, unlike sCM, we make the (novel) connection between this clamping operation and a simple regularization added to the original objective, to avoid it seeming like an ad-hoc solution (see Eq. (5) in the paper and below). Without any regularization (i.e. setting $r=1$ and not using adaptive double normalization) training diverges. Note that it is standard practice in this field to leverage such regularizations. As we are unable to share the training curve as image during the rebuttal, we list the 2-step FID scores calculated from samples generated by the distilled model after different amounts of training:

  Training images seen	4M	8M	12M	16M	20M	24M	28M
  AYF without clamping	4.91	2.75	2.53	2.46	2.72	4.03	51.04
  AYF with clamping r=0.99	4.92	2.91	2.65	2.70	2.69	2.44	2.25

  In these experiments, both runs reach their maximum $r$ values at 20M images, after which the non-clamped version starts to exhibit unstable training dynamics, as seen by the diverging FIDs. In contrast, the clamped version continues to improve and exhibits smooth training curves.

• Text-to-Image LoRA rank ablations: Thank you for pointing this out. In our original experiments, we opted to use rank=64 to follow prior works such as TCD [2], which used the same rank. To properly evaluate the impact of the LoRA rank, we now trained additional LoRA models with ranks 16, 32, and 128. Qualitatively, all variants produced nearly identical outputs when given the same noise input. As we are unable to share images during the rebuttal, we instead ran a user study. For each rank, we generated 260 images and asked human raters to evaluate the outputs based on fidelity and text alignment (vote for “best” images or say "Tie"). The preference vote counts are summarized below:

  	Rank=16	Rank=32	Rank=64	Rank=128	Tie
  Fidelity	48	59	61	65	24
  Alignment	53	60	66	56	22

  The results suggest that all LoRA ranks perform comparably, with rank 16 showing slightly lower fidelity. During training, we also observed that lower-rank LoRAs converge faster, while higher-rank ones require more iterations. Our ablation shows that the results are generally not very sensitive to the LoRA rank. We will add this to the final version of the paper.

• Typos: Thank you for spotting and pointing out the typos. We will make sure to fix those in the final version of the paper.

• Better flowmap discussion: We will include a more in-depth explanation and discussion around prior works that propose learning flowmaps in the final version of the paper. However, we would like to kindly point to both section 4 (related work) in the main text as well as to the discussions of related methods in Appendices C.3, C.5, D.1 and D.2.

• AYF-LMD being inferior: As the reviewer points out, we introduce two objectives for training flowmaps, AYF-EMD and AYF-LMD, but in our image experiments we only end up using the former. However, we also show in the appendix that for smaller low-dimensional datasets, such as 2D points, it is the AYF-LMD loss that performs significantly better. As such, the choice of which objective to use in practice becomes entirely dependent on the dataset and even data type (keeping in mind that future work may apply our methods to data beyond images; e.g. molecules, 3D point clouds, etc.).

  As for why exactly the AYF-LMD objective leads to overly smoothened outputs in image generation, our main point is our empirical evidence (see Table 5 in the appendix). Intuitively, we believe it to be due to compounding errors in the objective itself. This objective essentially regresses $f(x_t, t, s)$ towards a target that can be obtained by computing $f(x_t, t, s + \epsilon)$ and following the ODE flow from $s + \epsilon$ to $s$. However, if $f(x_t, t, s + \epsilon)$ has high error and ends up far from the true $x_{s + \epsilon}$, then following the flow from $s + \epsilon$ to $s$ may end up pushing the target away from $x_s$, leading to an error accumulation problem. Note that this is not the case for AYF-EMD as we first follow the teacher ODE from $x_t$ to $x_{t - \epsilon}$ with this operation introducing no additional errors. We will add these intuitions to the final version of the paper.

• we prove for the first time: We have not seen the problem of error accumulation in consistency models being mathematically proven anywhere else. Hence, to the best of our knowledge, this is a novel contribution to the field. Note that consistency models have become an important research direction and model class. Therefore, we believe that demonstrating and studying these flaws analytically represents an important, new insight.

We hope we were able to address your questions and concerns. To summarize: we introduced two novel objectives for learning flow maps and showed that each excels in different settings (AYF-EMD for high-dimensional data, and AYF-LMD for low-dimensional tasks). Building on prior work, we also proposed training techniques that help stabilize gradients when using pretrained checkpoints and ensure stable flow map training. This allowed us to achieve state-of-the-art performance in extensive experiments and scale our AYF to the text-to-image generation level. Finally, we now extended these text-to-image experiments by ablating over different LoRA ranks, finding improvements up to rank 64, with diminishing returns beyond that.

We will include the above results and discussions in the final version of the paper. If you found our explanations convincing, we would like to kindly ask you to consider raising your score, accordingly. Thank you.

References:

1. Lu, C., & Song, Y. (2024). Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.11081.
2. Zheng, J., Hu, M., Fan, Z., Wang, C., Ding, C., Tao, D., & Cham, T. (2024). Trajectory Consistency Distillation. ArXiv, abs/2402.19159.

Dear reviewer Rb7o,

We are glad you found our paper well written, highlighted our model’s performance, and enjoyed our proof of the limitations of consistency models. We appreciate your feedback and address your comments below.

• EDM2 time embeddings are unstable, theoretical analysis: Note that, as mentioned in the paper, EDM2-style models embed their timestep t by turning it into the logarithm of the inverse signal-to-noise ratio $\log \sigma_t = \log (\frac{t}{1-t})$ first, before passing the embedding on to the neural network. In such cases, when the time derivative with respect to $t$ is computed in these models, analytically the result ends up looking like: $\partial_t v(x_t, t) = \partial_t NN(x_t, \log \sigma_t) = \frac{\partial}{\partial \log \sigma_t} NN(x_t, \log \sigma_t) \times (\frac{1}{\sigma_t}) \times \left(\frac{1}{(1 - t)^2}\right) = \frac{\partial}{\partial \log \sigma_t} NN(x_t, \log \sigma_t) \times \left(\frac{1}{t \times (1 - t)}\right)$ As we can see, when $t \to 0$ or $t \to 1$, the second term in this equation goes to $\infty$, which causes the time derivative to blow up, leading to instability. This is the reason we propose changing the time embedding from using $\log \sigma_t$ to $t$ instead.

• Using EDM2 checkpoints by stabilizing time embeddings: To make this change, a short supervised finetuning stage is performed where we align the time embedding layers as described in Section 3.4. Optionally, we follow this with a brief full-model finetuning stage, aligning the student’s output with the teacher’s to minimize bias. In practice, we found that aligning only the time embedding yields an L2 loss of approximately 0.0011 between student and teacher outputs, which improves marginally to 0.0010 with full finetuning. To quantitatively measure the effect of this finetuning, we compare the original EDM2 checkpoints and our finetuned flow-matching checkpoints using FID scores at varying numbers of function evaluations (NFEs, corresponding to sampling steps). Results are shown below: | NFE | 4 | 8 | 16 | 32 | 64 | 128 | |----------------|---------|---------|---------|---------|---------|---------| | EDM2 | 388.007 | 142.896 | 9.324 | 1.675 | 1.402 | 1.370 | | FM-Finetuned | 192.377 | 46.82 | 7.892 | 2.370 | 1.453 | 1.381 |

  As shown, our finetuned checkpoints achieve FID scores close to the original EDM2 models, with any differences falling within a reasonable margin likely due to differing timestep schedules used during sampling. In short, we observe no significant degradation in performance.

• Stable embeddings are important for smooth training: We would also like to reiterate the importance of the timestep embedding finetuning. In addition to the analytical derivation above, here we demonstrate this empirically. To this end, we ran an additional ablation, doing AYF distillation while comparing two teacher models: one using EDM2-style embeddings and the other using our stabilized variant. The EDM2-style embedding led to frequent instabilities mid-training that ended with the run diverging entirely, whereas the stabilized variant was robust. As we are unable to share the training curve as an image during the rebuttal, we list the 2-step FID scores of model samples generated after different training iterations below: | Training images seen | 4M | 8M | 12M | 16M | 20M | 24M | 28M | 32M | |-----------------------------------|------|------|------|------|------|------|------|-------| | Teacher with EDM2-style embedding | 8.64 | 5.04 | 4.91 | 5.43 | 5.38 | 5.83 | 10.33| 20.54 | | Teacher with FM-style embeddings | 4.92 | 2.91 | 2.65 | 2.70 | 2.69 | 2.44 | 2.25 | 2.11 |

  We see that the distilled model from the original EDM2 teacher with the original embeddings does not reach low FID scores and eventually diverges (after around 28M to 32M seen training images). These results will be included in the final version of the paper.

• Novelty in Section 3.4: As the reviewer mentioned, some techniques in Section 3.4 are adapted from sCM [1], which is the closest prior work to ours. Building on existing methods like this is standard practice, and in our case, we actually found several of sCM’s tricks to be unhelpful or even harmful, so we left them out. That said, there are also new contributions in this section. In particular, the connection between tangent warmup and regularization, as well as the idea of finetuning the time embedding layers to enable the use of pretrained checkpoints, are both novel and were not used in sCM. This is why we highlighted these new contributions in this section in the main text.

We hope we were able to address your questions. In particular, we hope that we could address all concerns regarding the adaptation of the time embeddings. We will include the above results and discussions in the final version of the paper. If you found our explanations convincing, we would like to kindly ask you to consider raising your score accordingly. Thank you.

References:

1. Lu, C., & Song, Y. (2024). Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.11081.