Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling

We sincerely thank the reviewer for the valuable comments. We address the questions and clarify the issues accordingly as described in the following.

Q1: Concerns about the vast solution space.

A1: Thank you for the insightful comment. As explained in lines 304-307 and 912-917, our goal was not to find the globally optimal solution to discrete OT with LOOM-CFM, although that would certainly be ideal. Moreover, there are configurations of the initial pairings for which there are guarantees that the globally optimal assignment cannot be found with any of the minibatch-based approaches (see Figure 8). Thus, rather than directly outputting a globally optimal solution $\tau^*$, our method iteratively refines an initial estimate, providing an approximation that is consistently more accurate than those produced by prior minibatch-based methods and showing improved empirical metrics.

Besides this, the extensive solution space does not necessarily pose a concern, as long as local updates are effective. Indeed, LOOM-CFM can be viewed as a stochastic block-coordinate descent algorithm. Such algorithms typically operate in much larger continuous solution spaces, yet converge to an optimum under specific conditions. Another example is stochastic gradient descent, widely used in optimization, that does not necessarily guarantee an optimal solution but achieves strong performance in practice.

Finally, if we were to relax the optimization in Equation 7 by allowing soft assignments (i.e., permitting entries in $P_\tau$ to take any real value between 0 and 1, thereby expanding the search space even more), the problem would become a convex optimization that can be solved exactly with block-coordinate descent (as shown in, e.g., [1]). However, storing soft assignments would require $O(n^2)$ disk space, making this approach infeasible for larger datasets.

Q2: The training time of our model.

A2: We provided the details regarding the training time of LOOM-CFM in Appendix D, lines 968-976, as well as in Table 4. For fair comparisons, we ran our method for the same or fewer number of epochs than in the prior work. Our empirical results indicate that LOOM-CFM does not require more training time than previous approaches.

Q3: The meaning of the subscript $j$ in Equation 9.

A3: Thank you for spotting this! There is a typo in Equation 9: the subscript $i$ should have been used instead of $j$. We will update the paper accordingly.

[1] Xie, Y., Wang, Z. & Zhang, Z. Randomized Methods for Computing Optimal Transport Without Regularization and Their Convergence Analysis. J Sci Comput 100, 37 (2024). https://doi.org/10.1007/s10915-024-02570-w

We sincerely thank the reviewer for the positive comments on our work! We address the questions and clarify the issues accordingly as described in the following.

Q1: To what degree does LOOM-CFM accelerate model convergence compared to other coupling techniques?

A1: In our experiments, for fair comparison with baseline methods, we ran LOOM-CFM for the same or fewer training epochs than those used in prior work. For example, on CIFAR10, LOOM-CFM with 4 caches was trained for 1000 epochs. Notably, even within the first quarter of the training, LOOM-CFM achieved the following results that already outperform the baseline “w/o saving reassignments” (equivalent to BatchOT) for lower NFE values (see Table 5):

Results for LOOM-CFM, 4 caches, at 100k iterations (out of 391k)

On ImageNet-32(64) LOOM-CFM has been trained for only 200(100) epochs and outperforms BatchOT that was trained for 350(575) epochs (see Tables 2 and 4).

Q2: Relatively weaker empirical performance compared to EDM [1] and EDM2 [2].

A2: We did not include comparisons to [1] and [2] in section 4, as those would not be apple-to-apple, since the training settings in [1] and [2] are subtantially different from ours. To ensure apple-to-apple comparison with prior methods, we designed the experimental settings to closely match those used in baseline methods, including the same loss function, model architecture, hyperparameters, and comparable number of iterations. By isolating the data-to-noise couplings as the sole variable, we aimed to clearly highlight the advantages introduced by LOOM-CFM over previous approaches. However, we believe that LOOM-CFM’s performance could be further enhanced by combining it with additional training modifications from related works, such as those in [1] and [2], thanks to their modular nature. We illustrate this potential in our experiments by incorporating LOOM-CFM into such frameworks as Reflow and latent space training.

We sincerely thank the reviewer for the positive comments on our work! We address the questions and clarify the issues accordingly as described in the following.

Q1: Training time overhead and convergence analysis.

A1: The training time overhead of LOOM-CFM is addressed in Appendix, lines 968-976. In theory, LOOM-CFM does not introduce additional computational overhead compared to other minibatch OT methods, as it performs the same operations (locally solving the optimal transport). However, a minor I/O overhead arises from loading and saving the current assignments to disk during training, which slightly impacts the training time.

For the convergence analysis, we agree that theoretical convergence rate results, potentially depending on dataset and minibatch sizes, would strengthen the paper. Notably, if soft assignments were allowed (i.e., relaxing $P_\tau$ so its entries could take any value between 0 and 1), the problem in Equation 7 would become a convex optimization problem. In that case, a block-coordinate descent approach, with large enough blocks, could yield linear-time convergence (see [1] for details). However, in practice, this soft assignment approach is infeasible for large datasets due to the need for $O(n^2)$ storage and increased computational complexity. Instead, LOOM-CFM restricts updates to smaller block sizes and to permutation matrices and thus remains efficient. However, this complicates the analysis. We hypothesize that results from [1] could potentially help establish convergence bounds for LOOM-CFM, and we plan to explore this in future work.

It is also important to note that LOOM-CFM is guaranteed to yield a more globally optimal coupling than other minibatch OT methods already after the 2nd epoch (when assignments start to be shared across minibatches). Moreover, empirically, we find that LOOM-CFM converges quickly enough to achieve strong results within a reasonable number of iterations.

Q2: Potential influence of multiple noise caches on the probability flow.

A2: As discussed in Section 3.2, storing multiple noise caches and sampling from them is equivalent to duplicating data points. Since each data point is duplicated an identical number of times, the empirical data distribution, from which points are sampled during training, is unaffected. More precisely, before adding additional noise samples per data point, the empirical data distribution is defined as:

$

p_{\text{data}}(x) = \frac{1}{n}\sum_{i = 1}^n \delta(x - x_i).

$

After duplicating each data points $c$ times, we obtain:

$

p_{\text{data}, \text{c noise caches}}(x) = \frac{1}{c \cdot n} \sum_{i = 1}^n c \cdot \delta(x - x_i) = \frac{1}{n}\sum_{i = 1}^n \delta(x - x_i) = p_{\text{data}}(x).

$

The distribution of noise samples also remains the same, as they are still drawn from the same source distribution. Since probability flow depends only on the source and target distributions, it should not be affected by the use of multiple noise caches. However, using multiple noise caches does lead to different final weights in the optimized neural network. Empirically, we observe that this approach improves generalization for smaller datasets.

[1] Xie, Y., Wang, Z. & Zhang, Z. Randomized Methods for Computing Optimal Transport Without Regularization and Their Convergence Analysis. J Sci Comput 100, 37 (2024). https://doi.org/10.1007/s10915-024-02570-w

We sincerely thank the reviewer for the positive comments on our work! We address the questions and clarify the issues accordingly as described in the following.

Q1: How scalable is LOOM-CFM?

A1: First and foremost, LOOM-CFM consistently delivers superior couplings compared to any local minibatch OT method, regardless of dataset size. Moreover, this is achieved with negligible additional cost, making it an easy-to-implement and robust alternative for coupling noise-to-data in flow-based generative models.

Some preliminary results on how LOOM-CFM scales with increasing dataset sizes can be seen in our experiments where we vary the number of noise caches. In those experiments, the batch size is fixed, while the dataset grows twice in size (from 2 noise caches to 4). Figure 11 demonstrates that while the convergence of LOOM-CFM assignments slows slightly for larger datasets, all configurations successfully converge by the end of training.

Nonetheless, for an ultimate and precise conclusion on scalability of LOOM-CFM to datasets larger than ImageNet, further experiments on such large-scale datasets would be necessary, provided sufficient computational resources are available to handle such extensive training.

Q2: Refreshing the noise during training.

A2: Thank you for the thoughtful suggestion. Indeed, in the early stages of developing LOOM-CFM, we explored various methods to prevent overfitting to a fixed set of noise samples. Although none of these approaches yielded satisfactory outcomes, we summarize them here for clarity:

1. Completely refreshing the noise: We experimented with replacing all noise samples every $N$ epochs. While intuitive, this approach caused training instability since the newly assigned noise samples did not necessarily align with the previous ones. Consequently, the network’s targets could shift significantly when the noise was refreshed, impairing convergence.
2. Gradual noise injection: In this approach, we introduced a hyperparameter, $\phi$, to control noise refreshing. For each data point in a minibatch, the assigned noise was replaced with a new sample with probability $\phi$. While this method allowed for a smoother refresh, choosing an optimal $\phi$ proved challenging. For example, setting $\phi = 0.1$ led to approximately one-third of each minibatch being reshuffled, which weakened the effect of LOOM-CFM. Conversely, a lower value of $\phi=0.01$ was insufficient to prevent overfitting.
3. Interpolation between LOOM-CFM and independent coupling: For each data point, we coupled it with its assigned noise with probability $\phi$ and with freshly sampled noise with probability $1 - \phi$. Unlike approach 2, the new noise did not replace the cached noise. We found this technique to be quite effective and the results indeed interpolated the results of LOOM-CFM and the independent coupling. We also explored making $\phi$ a function of $t$, as the curvature of sampling paths depends on $t$ (as shown in Figure 6). Unfortunately, this hasn't lead to substantial improvements. Nevertheless, we found this approach interesting for its generality, as it can interpolate any pair of coupling techniques. So we leave further exploration as future work.

In contrast to these more complex methods, the approach in the paper is straightforward, as it artificially increases the dataset size and equalizes the settings for problems with small and large dataset sizes. We opted not to include all of these alternative methods in the final paper, as they diverge from the primary approach and cannot be regarded as simple ablations. However, if the reviewer believes that discussing these approaches would enhance the paper’s clarity or contribution, we would be open to incorporating a discussion of them.