**The paper has a clean motivation and an insightful problem formulation **[…] The paper provides massive-scale experiments i.e., 2M×2M coupling matrices are employed—a scale that has not yet appeared in the literature […]The proposed methods give substantial practical gains

Many thanks for your very encouraging comments!

The role of coupling size and entropy in training stability or convergence speed is not adequately explored. Authors state they did not ablate learning rates for scale (p. 8, l. 280).

We are happy to report that learning rates have been ablated for Imagenet-32 and Imagenet-64. It seems that the original choices reported in the literature for IFM (lr=1e-4) are suitable for OTFM at all $n$ and $𝜀$ scales.

While sharing plots would be more convenient, we cannot. Hence, here are two table describing FID along iterations for ImageNet-32. We do observe, for ImageNet-32, a decrease then increase of the FID metric along iterations, even for IFM and other settings --- this was mentioned in Table 1 of the supplementary. We do not observe this in our runs on ImageNet-64.

$n = 1048576, 𝜀= 0.03$, FID@DOPRI5

$n = 1048576, 𝜀= 0.03$, FID@4

While the paper claims the Sinkhorn step is relatively low-cost, 2M×2M couplings take 8+ GPUs and substantial memory (p. 5, l. 192–207), which will hold back adoption.

This is a valid point, although 8-GPU nodes are by now the norm. We observe that most config files shared nowadays to train even small FM models on ImageNet-32 assume access to a node.

In our response to Reviewer zUcn we have considered several ways to improve the picture in sections

• 1. Warmstart (to reduce overall the number of iterations),
• 2. Computing matchings in PCA Space (to alleviate per-iteration cost and lower memory requirements for couplings)
• 3. Precomputing Noise/Data Pairs to split preprocessing from training.

Experiments lack missing standard deviation/error bars, which would help to estimate variability across runs (Checklist point 7, p. 13).

As mentioned in our response to W4 of Reviewer 7qcL, we cannot realistically provide so many $n,𝜀$ results along with error bars, as this would make our compute budget explode.

We notice that other researchers have recently complained specifically about this requirement (Ben Recht, Standard error of what now?, blog post) because it is not realistic.

The paper has not investigated the usage of other variants of OT which are shown to be better in a mini-batch setting [1] [2].

Thanks for these references! We have added these to the draft. However, we observe that:

• These approaches explore partial or unbalanced OT, which are (for now) unrelated to our setting.
• The code provided for these approaches cannot scale as they both rely on POT and GeomLoss solvers, which are intrinsically either CPU or single-GPU.
• We also note that [2] claims, regarding Sinkhorn, in its introduction, “that storing an n × n matrix is unavoidable in this [Sinkhorn] approach, thus the memory capacity needs to match the matrix size”. Interestingly, we believe that this is incorrect, and this is one of the pillars of our contribution, as explained in L.200, see also our point 2. Computing matchings in PCA Space to Reviewer zUcn for more details on online recomputations of cost matrices, a technique that was leveraged first by GeomLoss.

As explained in our answer to point S1, Reviewer 52g8, we do not see, at the moment, an alternative solver to the Sinkhorn algorithm that can simultaneously

• be implemented to utilise efficiently / maximally all 8 GPU nodes, to scale to millions of points;
• return couplings that can be arbitrarily close to the solution of exact OT;

Of course, we hope that other researchers will be encouraged, after seeing our work, to target the 1 million point scale with other solvers / variants, and demonstrate their relevance to flow matching.

Many thanks for your detailed comments.

The method is meta-dataloader: training and generation coupling are separated, thus making the procedure modular and computationally realistic.

Many thanks for highlighting this point! We will do our best to underline this further in our revision.

The main weaknesses is undoubtedly the lack of any form of theoretical justification to support the experimental conclusions, though achieving this may be beyond the scope of this work

While our main focus remains on scaling-up Sinkhorn and comprehensive experiments to study its effect, we are also able, following your comment and that of Reviewer 7qcL to provide some theoretical intuition into the necessity of using large OT batch sizes. This result is provided as Proposition 1 in response to W3 of Reviewer 7qcL.

We highlight that this argument will only be shown in the appendix, as it is only meant to provide insights into the necessity of large batch size and is not a core contribution of our work.

It appears to me that the usage of Sinkhorn coupling as initialization fits well with using stochastic rather than deterministic interpolants. In particular, I would use brownian bridges instead of straight lines. Have the authors thought in this direction? How do they believe their conclusion would change considering stochastic interpolants?

We agree that this is a promising direction. At the moment we do not exploit stochasticity of the coupling, but rather simply sample from it (Step 4) to form pairs, assuming straight paths, in accordance with the flow matching framework. Incorporating this stochasticity would be more akin to switching from ODE flow models to SDE Schrödinger bridge type models. We might explore this, or others might use our implementation to do it themselves before.

In that vein, we have also tested the ability of OTFM to help training for modality translation (i.e. flowing from a modality to another, without having access to pairings, rather than starting from nois). This setup was typically considered by Schrödinger bridge models.

We tried this with the cat→wild image translation for AFHQ-64. We follow the evaluation protocol as mentioned in [De Bortoli et al. 24]. We use their network architecture, increasing the number of channels in the U-Net to 192; train the network for 400k iterations with 512 effective batch size.

As noted by Bortoli et al., visual inspection is best way to assess the quality of results. While the pictures we obtain look good, we are not allowed to share these images this year, and might want instead to add them in the appendix.

As a result we can only report metrics (LPIPS, MSD and FID). As highlighted by Bortoli et al., FID is not really meaningful. as they are computed on a very small set of points (~4.5K training set of the wild domain and ~450 validation set of cat → wild transferred images). They show, however, a very clear failure of IFM vs. OTFM.

These preliminary results should be compared to the values reported in Figure 19 of [De Bortoli et al. 24].

[ V. De Bortoli, I. Korshunova, A. Mnih, A. Doucet, Schrödinger Bridge Flow for Unpaired Data Translation, Neurips 2024]

I understand the general recommendation is to have a renormalized entropy around 0.2. What are the recommendations for $n$ ?

Our recommendation at the moment is to choose a $n$ as large as possible, within what is available in the preprocessing compute budget, since in our generation experiments show so far that larger $n$ always improves on all metrics, notably faster generation.

Additionally, our implementation is now fully efficient as we have transitioned to a metaloader (See our point 3. Precomputing Noise/Data Pairs shared with Reviewer 52g8) in order to split entirely NN training from pairing effort, making this computation a one-off effort.

In general, I find the plots tell little about the influence on $n$ . What was the choice of for the plots in Figure 1 and 2 ?

We apologise for not being more clear in the legends of Figs. 1 & 2 (please look at 10, 12, 16, 18 in supplementary). In all our figures, $n$ is color coded (yellow = large). We will add the mention n=2048, n=16384, ... in all legends. As can be seen, $n$ impacts performance in all settings.

Many thanks for taking the time to review our paper, and providing several comments.

The paper is clearly-organized

Thanks!

W1. The contribution of the paper is limited. The main innovative point is dropping norms and focusing on the dot-product between points in the Sinkhorn algorithm for better stability, which is a narrow view to some extent. The overall contribution is mostly empirical.

We agree that our contributions are driven by empirical concerns, more precisely by the goal to scaling up $n, d$ to unprecedented regimes, while tracking convergence rigorously (small $\tau$, L.213), and staying as arbitrarily close to sharp couplings (small $𝜀$).

In that sense, our paper is informed by the “bitter lesson” (Rich Sutton) applied to OT computations, and is trimmed to the minimum to work in this case.

While our findings may seem simple in retrospect, from Lemma 2, to $\mathcal{E}$, the adoption of the dot product cost in L. 190, and rescaling by std , we are confident that these modifications will be impactful and become the norm when running Sinkhorn.

In addition to this, we have incorporated additional tricks (see 1. Warmstart , 2. Computing matchings in PCA Space in our response to Reviewer zUcn) and demonstrated that they speed up computations significantly in our setting.

Some experimental details like the types of GPU could be moved to the appendix

Hardware specifications only take a few characters in our draft. We prefer to be sure the reader is aware of these aspects, as they matter practically speaking, and can help get a clearer picture. Are there other details you would like to see moved?

W2.The paper states that larger batch size is needed to cover the bias that cannot be traded off with more iterations on the flow matching loss. On the other hand, Figure 1 and Figure 2 show that larger batch size suffers from longer training time.

Indeed, using larger $n$ and decreasing 𝜀 incurs a larger compute effort. We ask, however, to consider the following aspects:

• Coupling more carefully noise and images is not, strictly speaking, a training effort, it is a data-preprocessing effort, independent of model training (L.142) as highlighted in L. 141 and expanded in item 3. Precomputing Noise/Data Pairs shared with Reviewer 52g8.
• The cost to compute such couplings is parameterised by $n, d, 𝜀$ and is independent of model parameters, which can run anywhere from billions to hundred of millions, and will always dominate overall training cost.
• We believe, following discussions on Warmstarting and PCA, that one may continue improving on the efficiency of computing couplings.
• A larger effort spent on better couplings is always associated in our image experiments with a higher quality when inference time compute is constrained. This is usually desirable.

W3. It would be beneficial to include some theoretical complexity analysis like [1] that includes the sample complexity to further explain the “necessity of large batch size”.

Thanks for this great suggestion. We will gladly include [1] to the 5 references we have provided in L.145-150. We will expand this section and include any other suggestions you may have.

Following your comment, we can add (in response to your comment that this would be beneficial) to our appendix a simple mathematical argument, similar to [Chewi et al. 2024, Thm. 2.14], that justifies the use of large batch sizes. We emphasize that this argument is only meant to provide insights into the necessity of large batch size and is not a core contribution of our work. Using this insight to establish a tight characterization of the curvature of flow models as a function of the OT batch size would require significantly more work, which can be an interesting research direction.

Assumption 1: The data distribution $\mu_1$ is such that if $X$ and $X'$ are independently drawn from $\mu_1$ then, \mathbb{P}[\Vert X - X\' \Vert \geq t] \leq Ct^r, for all $t > 0$ and some constants $C,r > 0$.

This assumption holds for example if $\mu_1$ is supported on a manifold of dimension $r$ with bounded density w.r.t. the volume measure on the manifold. Specifically, $r$ can be seen as the effective dimension of the support of $\mu_1$, and for structured data distributions we expect $r \ll d$.

Proposition 1: Given two probability measures $\mu_0, \mu_1 \in \mathcal{P}_2(\mathbb{R}^d)$ where $\mu_1$ satisfies Assumption 1, suppose $\boldsymbol{X}_0, \boldsymbol{X}_1 \sim \mu_0^{\otimes n} \otimes \mu_1^{\otimes n}$, i.e. they represent $n$ i.i.d. samples from $\mu_0$ and $\mu_1$ respectively. Let $\hat{\pi}(\boldsymbol{X}_0,\boldsymbol{X}_1)$ denote any coupling supported on $\boldsymbol{X}_0$ and $\boldsymbol{X}_1$, including (entropic) optimal transport. Then, $\mathbb{E}\_{X_0,X_1 \sim \hat{\pi}(\boldsymbol{X}_0,\boldsymbol{X}_1), \boldsymbol{X}_0, \boldsymbol{X}_1}[\operatorname{Var}(X_1 | X_0)] \geq cn^{-2/r}$ for some constant $c > 0$ depending on $C$ and $r$.

In Proposition 1, we look at the variance (sum of coordinate variances) of data conditioned on the noise it is paired to. Note that for IFM, this is simply the variance of the data distribution, as the noise/data coupling is independent. Ideally, with OT couplings we would always couple noise to a unique data sample, thus this variance would be zero and the flow trajectory would be straight. The above lower bound shows the necessity of large batch-size $n$ due to the slow rate of $n^{-2/r}$, which motivates us to scale Sinkhorn to as large $n$ as possible. While we do not include it for brevity, one can obtain matching upper bounds through standard law of large numbers arguments [Chewi et al. 2024, §2.5.1].

W4. Multiple independent runs of the experiments could be done to obtain the confidence interval of the results.

Thanks for this suggestion. Unfortunately, at this time, most papers in flow matching or large scale generative modeling do not include error bars when using image generation, e.g. Table 5 in [Tong et al.], Table 1 & 2, Figure 3 in [Pooladian].This is because each run takes anywhere between 4 days to 2 weeks.

To make things worse, in our case we report dozens of $n,𝜀$ variants. Producing error bars would multiply by 5 our already very substantial compute budget.

Additionally, the randomness in coupling computations is independent from that resulting from neural network parameter initialization, increasing the challenge further (to get error bars for FM metrics, one would need to rerun multiple seeds for couplings $\times$ multiple seeds for NN training)

This being said, we observe qualitatively fairly stable results, as long as we reuse similar config files. For instance, we ran the learning rate ablations (in answer to Reviewer 3ETv) **** using different NN seeds, and obtained similar results. We sometimes observed variability in FID computations across experiments, due to the batch of 50k images that is sampled. We always ensured, however, that the multiple runs done within a single experiment use the same batch, but were not always able to guarantee this across experiments.

Q1. To balance between training time and mini-batch size $n$ , is there any criterion on how to choose the proper $n$ when the total training time is limited?

Compute involves three independent terms:

• pre-processing, to select pairs of noise / data → this is where OTFM happens.
• training time: pick FM model, optimizer, learning rate, epochs, etc.. fed with pairs from above.
• inference time: ODE integrator, # of steps.

Any rule to select $n$ must depend on how these compute budgets are prioritised, as with LLMs training.

Nowadays, pre-processing steps are allocated a high budget, training budget is even higher, and inference budget is assumed to be limited. A message of our paper is that $n$ should be as large as possible (although probably not larger than dataset size, a lesson learned with the relatively toy-ish CIFAR-10 experiment).

In practice, it is easy to estimate the required time to run a single Sinkhorn computation: Sinkhorn time is very stable across batches for a given $n,𝜀$ choice (see e.g. Fig. 8, top, in supplementary). The practitioner can then fairly easily calibrate their training setup with a dry run, and adjusting $n,𝜀$ according to the preprocessing compute they wish to spend, and store these results as per item 3. Precomputing Noise/Data Pairs shared with Reviewer 52g8.

Q2. Is it possible to give a theoretical complexity analysis like [1] to show the effectiveness of the modification of Sinkhorn algorithm?

The modifications we provide are not theoretically motivated. For instance,

• the std rule is only designed to help set 𝜀 adaptively, robustly and consistently across many setups.
• removing norms does not change the convergence analysis of Sinkhorn, as described in Remark 4.12 [Peyré / Cuturi]. It can be interpreted as a way to bypass completely norm information in dual scalings, whereas convergence theory of Sinkhorn focuses on lower-bounding the improvement from one iteration to the other.

Q.3 In the final row of Figure 1, why are there 6 different curves when only have 5 choices? The colors and the marks of the curves do not match the legend (there are two blue lines, and the colors are mismatched).

We apologise for this mistake and will redraw this plot. The Lowest curve can be dropped and corresponds to the setting $n=256$ that we decided to remove to improve legibility. See also our response to W1 of Reviewer zUcn

It would be beneficial to add the theoretical complexity analysis that includes the sample complexity $n$. Criterions for selecting proper $n$ given a limited time budget remains to be established.

Thanks, as mentioned above, we will discuss [1] and add the discussion above.

Many thanks for your detailed review and your insights.

We also found a few unfortunate typos (residual $m$ cardinalities in Alg.1, wrong transpose on $\mathbf{Y}$ in L.185, also missing a "+", etc...), and apologize for this. We do, however, stand by the maths in our paper.

various OT solvers[…] including the Hungarian algorithm

Let us clarify that we never use the Hungarian algorithm, we only use Sinkhorn.

S1. according to the authors, GPU computations of this scale[…]have never been reported in the literature.

Thanks for raising this point. Indeed, our computations are unprecedented for two reasons:

Scale. The hardness of an OT problem hinges on many parameters, typically: $d$ (dimension); $n$ (# of points); constraint tolerance (e.g. $\tau$ in Alg. 1); intended closeness to the optimal assignment/cost (e.g. 𝜀).

To our knowledge, no paper has pushed computations of Sinkhorn to this extent, e.g. for ImageNet-32, $n=524288, d=3072$ while tracking precisely constraints ($\tau= 10^{-3}$ in 1-norm) and sharpness (small 𝜀). The largest setup we know of is the computation of a few coupling ($n=60k, d=1k$) in [Kassraie & al. 2024, Fig.6].

While low-rank approximate OT solvers have been used recently to compute fairly large couplings (see [Halmos, 2025], [Klein & al. 2025]), these reformulations are not convex, and offer no guarantees to solve the OT problem. Using them in large scale OT-FM could be an interesting follow-up work.

Repeats. We compute large couplings many times: e.g. for ImageNet-32, with $n=524288$ points, we do so ~850 times, to cover 430k training steps.

P. Kassraie & al., Progressive Entropic Optimal Transport Solvers, Neurips 2024

Halmos & al., Hierarchical Refinement: Optimal Transport to Infinity and Beyond, ICML 2025

This is a salient finding — if one does not want to tune coupling hyperparameters in flow-matching, one can just use a better ODE solver for generation.

This finding is of little use in the context of FM/diffusion.

• The number of ODE steps is an inference-time parameter: in practice, for an end-user, better quality = more steps = longer wait.
• Computing larger/better couplings is a preprocessing effort.

Because inference and preprocessing are done by different actors on different hardware (e.g. consumer devices/lower grade GPU vs. GPU servers with interconnect), one cannot easily trade off one for the other. When models are shared with the public, inference cost will always dominate training cost. This is why we need fast flowing models.

One stated contribution […] this continuum is fairly widely understood.

We agree, this continuum is well understood in OT papers, but we stress

• in L.69 that it has not been leveraged enough in the OT-FM literature.
• Pooladian / Tong have not ablated 𝜀 in their studies (L.70) and emphasized instead a distinction between Hungarian / EOT / IFM.
• In L.73, we claim that these methods can be unified by varying 𝜀 in Sinkhorn, as long as convergence 𝝉/max_iters is rigorously tracked (L.213)
• In L.79 we facilitate this comparison across tasks and scales using $\mathcal{E}$.

many of the apparent trends reported on plots […] are actually marginal-to-nonexistent when one examines the y-axis limits.

numerical results are presented in a (I hope unintentionally) deceiving manner […] the y-axes on the problematic plots do not start at zero[…] trends would not be significant were the y-axes started at zero.

all need to start at 0[…] Please fix

We acknowledge your concern that FID or BPD sometimes vary in a narrow y-range in our plots, but this is usual in the diffusion/FM literature. Adding a 0 would be highly unusual and clutter plots. Instead we follow the practice of very highly cited papers in the literature (and more generally ML) to let matplotlib choose the y-axis bounds:

• Lipman & al, Flow Matching for Generative Modelling FID: Fig. 5 & 7
• Song & al., Score-based Generative Modeling through Stochastic Differential Equations FID: Fig. 10
• Liu & al., Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow FID: Fig. 8; Straightness (our 𝜅): Fig. 9
• Nichol & al., Improved Denoising Diffusion Probabilistic Models NLL/BPD: Fig. 15 & 16; FID: Fig. 17

While adding a 0 in our y-axes would be problematic, we commit, however, to

• provide more tables (as done in this rebuttal),
• highlight in the captions the narrow BPD or $\mathcal{R}$ ranges
• remove Fig. 11 & 13, as they just present the data in Fig. 1 & 2 without the 🔻 IFM / green▲ ground-truth lower bound.

The numerical results […] do not paint a clear picture […]exception is perhaps the Korotin benchmark.

For the Piecewise quadratic benchmark (Fig. 1)

• very large $n$ (1M, 2M) improves on small $n$ (16k) and IFM on all metrics.
• The curvature 𝜅 for large $n$ is 25x to 4x times smaller than IFM
• The reconstruction loss $\mathcal{R}$ of large $n$ OTFM is always 10~20% better than IFM.
• BPDs are comparable, but this is usual: BPD values tend to cluster when comparing similar methods (see e.g. Table 7 in [Papamakarios & al. Masked Autoregressive Flow for Density Estimation])

We agree that CIFAR-10 results lack significance, because CIFAR-10 is too small (the 50k database size is smaller than our larger batch sizes $n$). We will move CIFAR-10 to the appendix ([Pooladian et al.] skipped it altogether).

For ImageNet-32 and ImageNet-64, all FID metrics for small NFE are significantly better than IFM, and 𝜅 improves as well (Figs. 16 & 18, Tables 1 & 2). This is visually striking in Fig. 17 and Fig. 19.

Even though it is perhaps disappointing […] inform effective FM practice.

Given your lowest possible rating of 1 to the significance & originality of our work, we felt encouraged to see that it still managed to trigger so many open questions.

Our findings are, however, clear: using large $n$ / small 𝜀 Sinkhorn couplings always improved (or left unchanged) FM metrics in all problems we considered.

The magnitude of these gains depends on data, as you mention, but also on model capacity & training method. Predicting the magnitude of these gains will necessitate a "scaling laws" approach that we leave for follow-up work.

What is the rationale for scaling the entropy level by the std of the cost-matrix[…] numerical results comparing this choice[…]

As an example, consider a $n\times n$ cost matrix $C_1$ with values distributed in $[0,1]$. Consider the shifted cost $C_2 = C_1 + 100$.

• Mathematically, the EOT solution for $C_1$ is the same as that for $C_2$. Indeed, because $P$ is a bistochastic matrix, $\arg\min_P\langle P, C_2 = C_1 + 100\rangle - 𝜀H(P) = \arg\min_P \langle P, C_1 \rangle + 100 - 𝜀H(P) = \arg\min_P\langle P, C_1 \rangle - 𝜀H(P).$
• Hence, EOT solutions are invariant to the addition of a constant in the cost when using the same 𝜀 (a particular case of L.187 with a global shift)
• Hence, a data dependent rule to set 𝜀 (or its scale) should retain this shift-invariant property.
• The mean, median or max rules introduced in [Cuturi ’13] are not shift-invariant. They were introduced to avoid underflow in the kernel $K=\exp(-C/𝜀)$. Most solvers used nowadays avoid this problem by relying on log-sum-exp computations (L.109).
• Yet, the 𝜀 selected with mean, median or max rules would be 100x larger for $C_2$ compared to $C_1$, despite the fact that these problems are highly similar.

To scale 𝜀 in a shift-invariant way, we propose to quantify dispersion with std. While std is not the only way to achieve this, it is cheap and robust.

We do observe qualitatively that std scaling results in lower variance of # Sinkhorn iterations / renormalised entropy across experiments, compared to mean. We will clarify in the appendix.

changing notation to W_2

We will remove the Wasserstein notation entirely, and only introduce OT couplings.

On lines 115-117 […] Eq. (1) is the continuum OT problem, and therefore its solution is a coupling, not a matrix.[…]

Initially, we instantiated Eq. (1) with $\mu_n=\tfrac1n\sum_i \delta_{x_i}$ and $\nu_n=\tfrac1n\sum_j \delta_{y_j}$. We removed this when shortening the paper. We will revert and clarify.

the random variable $T$[…] Shouldn't specifically be uniform[…]?

Time is not necessarily sampled uniformly, see e.g. paragraph above Eq. 4.27 in [Flow Matching Guide and Code]. We will clarify.

On line 169, it is stated […] $\mathcal{E}$ does not actually tell us whether we are close to the optimal assignment, just whether we are close to any 1-1 assignment.

Thanks for the opportunity to clarify.

While we define $\mathcal{E}(P)$ for any coupling (L. 167), we write in L.168 “$\mathcal{E}(\mathbf{P}^𝜀)$ provides a simple measure of the proximity of $\mathbf{P}^𝜀$ to an optimal assignment matrix”. We evaluate $\mathcal{E}$ exclusively on Sinkhorn EOT solutions $\mathbf{P}^𝜀$ in the regularization path.

On that path:

• $\mathbf{P}^𝜀$ interpolates between $\mathbf{P}^\star$ (the optimal permutation, assuming it is unique w.h.p.) and $\mathbf{1}_{n\times n}/n^2$.
• $H(\mathbf{P}^𝜀)$ strictly increases as 𝜀 grows (e.g. Eq. 4.5 in [Peyré and Cuturi’19]), and is s.t. $\log(n)\leq H(\mathbf{P}^𝜀)\leq 2 \log n$.
• Therefore, if $\mathcal{E}(\mathbf{P}^𝜀)\approx 0 \Rightarrow H(\mathbf{P}^𝜀) \approx H(\mathbf{P}^\star) \Rightarrow\mathbf{P}^𝜀 \approx \mathbf{P}^\star$.

Hence, qualitatively speaking, $\mathcal{E}$ is a cheap proxy to quantify closeness of $\mathbf{P}^𝜀$ to $\mathbf{P}^\star$ without having to compute $\mathbf{P}^\star$ itself.

citation styles used are inconsistent and out-of-step with NeurIPS guidelines.[…]

We follow the Neurips guidelines (Section 4.1) and use the natbib \citeauthor macro to highlight famous names/results, to facilitate reading. We can remove this.

Following the comment of Reviewer zUcn on small OT batch size $n$, we had committed to include small $n$ experimental results in the draft. This is also a topic that appeared in your latest response to our rebuttal, hence we take this opportunity to share our latest results.

We are happy to report preliminary experiments with small $n=256$ on ImageNet-32, which confirm that while using small OT batch sizes $n=256$ can offer some FID gains over IFM for small NFE, it may also perform worse relative to IFM for larger NFE/Adaptive solvers. On the other hand $n=256$ is always significantly worse than large $n$.

The following table shows our experiment with small OT batch size $n=256$ (i.e. $32$ images per device) on ImageNet-32, and compares it with IFM and a larger OT batch size. As seen from the table,

• large OT batch size $n= 524288$ is substantially better than small $n=256$ for all metrics.
• For large NFE, we even see that IFM edges above $n=256$ regardless of 𝜀.

These comparison are made at the 180K step checkpoint (out of total $\approx 430K$ schedule) because our runs with $n=256$ batch size haven’t finished yet. We note, however, that measures like FID@NFE=4/8 have already plateaued and will not improve at later checkpoints (see also our response to Reviewer 3ETv showing that ImageNet32 metrics tend to plateau from 150k iterations, and might even slightly increase again after ~250k iterations)

We are not able to show generated images (as in Fig. 17) per the rebuttal policy this year, but these images come as expected (notably blurrier at $n=256$ for small NFE compared to large $n$).

Many thanks again to you (and Reviewer zUcn) for suggesting this experiment, we agree that it completes our message by also looking at what happens at the "left end of the range", not just at very large $n$.

Many thanks for your encouragements and detailed review.

They leverage the setup of [1] to keep entropy regularization parameter $\epsilon$ in the range [0,1] instead of $[0, \infty]$

Let us clarify this:

• The renormalised entropy $\mathcal{E}(\mathbf{P}^𝜀)$ is a metric for EOT solutions taking values in $[0,1]$ (L.176). This metric assesses whether the solution $\mathbf{P}^𝜀$ lies close to the optimal assignment (when $\approx 0_+$) or the independent coupling (when $\approx 1_-$), see also our answer to Rev. 52g8.
• We still need to set 𝜀: we set it as a multiple of the std of the cost (L. 166). Choosing that multiple in $[0.003, 0.3]$ we got a sufficiently wide coverage of $\mathcal{E}$ in $[0,1]$ in experiments.

The paper is very well written, with a nice introduction to the topics and literature review

Thanks!

W1. The authors considered a minimum batch size of 2048, it would be important to know the effect of mini-batch OT in the range below that […]

This is a great point.

We used $n=256$ early on in the piecewise affine benchmark and observed fairly poor results for $\mathcal{R}$ and 𝜅 (we were not tracking BPD at that time), see table below for $d=32/256$, which should be compared with other $n$'s in the left/right columns of Fig. 1 (preferably retrieved as Fig. 10 in supplementary)

Performance drops compared to $n\geq 2048$. More surprisingly, the reconstruction metric $\mathcal{R}$ for $d=256$ ($\approx 25$) is much worse than IFM ($\approx 21.6$) or OTFM baselines ($\leq \approx 21.4$).

• for small 𝜀, we hypothetize that this is due to the statistical bias of exact OT for large $d$ / small $n$.
• For large 𝜀, one would expect results to be, in principle, more similar to IFM (they both result in independent sampling). This difference is likely due to the way we implemented Alg. 1, more specifically Step 3.
  • When using IFM, we set $\mathbf{P}=\mathbf{I}_n/n$ (L.136), i.e. the arbitrary identity assignment;
  • For OTFM we use $\mathbf{P}^𝜀$, which looks like $\mathbf{1}_{n \times n}/n^2$ as 𝜀 grows.
  • Both procedures are roughly equivalent for large $n$ but differ for very small $n$ because OTFM would be less sample efficient, as it picks samples with replacement and likely duplicates some data points while dropping others (stratified vs. non-stratified sampling).

Generally, using very small $n$ was challenging in our earlier implementation, as it resulted in coupling sizes that were potentially smaller than the default batch size needed to train FM models.

We will rerun some small $n$ experiments for ImageNet32/64 and add this discussion to the appendix.

W2. Could the authors confirm if Imagenet was done in an unconditional setup?

Yes, all results in the paper & supplementary were unconditional.

W2. […] conditional set up is however, of great importance to the community. Would the authors be able to comment/provide evidence of how these results change when considering, for instance, class labels?

[Chemseddine 2023] and others have applied OT-FM to the class-conditional setup, augmenting the feature space with a one-hot class label. We extended our ImageNet32 experiments to conditioned generation over 1000 classes. The following FID results are for 𝜀=0.1/I-FM after 180k steps, and renormalized entropy is consistently ~0.1.

Imagenet 32: conditional generation FID

When testing conditional generation, we plan to move to latent-space representations. We prioritized for now unconditional / pixel space generation because the community wants to see generation methods operate reasonably well in pixel space (performance is harder to judge visually when rendered using high-quality latent spaces).

[ Chemseddine J, et al.,Conditional Wasserstein distances with applications in Bayesian OT flow matching ]

The authors mentioned that the results for the main paper were not complete, could you please provide the new results?

Indeed, the main paper results were incomplete at time of submission, as some of our jobs (e.g. Imagenet32/64) took a few days longer than expected. We apologise for this.

The complete results were submitted in the supplementary file, 1 week after the main paper deadline. They are in the .zip file at the top of this page.

Speeding up Sinkhorn

Since submission, we have progressed on our agenda to scale up Sinkhorn for OTFM to unprecendented scales. We report on the following changes, that we plan to add to our draft.

1. Warmstart

We leverage warmstarting for Sinkhorn. Rather than always reinitialize $\mathbf{f}, \mathbf{g}$ as $\mathbf{0}_n$ in L.1 in Alg.1, we extrapolate the solution obtained when coupling the previous batches of (noise, data) to initialize $\mathbf{f}, \mathbf{g}$ for the next batches.

More precisely, for a pair of optimal $n$-dimensional vectors $\mathbf{f}, \mathbf{g}$ outputted at Line 7 of Alg.1, corresponding to $n$ noise $(\mathbf{x}_i)$ and $n$ data $(\mathbf{y}_j)$ vectors, we follow [Eq.9, Thornton & Cuturi 2022] to initialize $\mathbf{f}, \mathbf{g}$ for a new batch $\mathbf{x}'_i, \mathbf{y}’_j$ as $\mathbf{f}^{\text{new-init}}$, computed adaptively as

$\mathbf{f}^\text{new}_i = 𝜀 \log \tfrac{1}{n} + \min_𝜀(-\langle\mathbf{x}'_i,\mathbf{y}_j\rangle - \mathbf{g}_j), \text{ and setting } \mathbf{g}^{\text{new-init}}=\mathbf{0}_n$

For ImageNet-32, we observe substantial ~1.7 x speedups for large $n$ as reported below. We report runtime in seconds, we have checked that using iterations results in the same trends.

Imagenet 32, Average Sinkhorn time (in seconds) to solve Alg. 1

Same as above, using warmstarting

Please note that warmstarting is not an approximation, it is simply a way to leverage previous solutions to solve Alg. 1 faster.

2. Computing matchings in PCA Space

To fit computations for large $n$ in memory, our implementation reinstantiates the cost $\mathbf{C}=-\mathbf{X}^T\mathbf{Y}$ (L.190) block by block at each Sinkhorn iteration (L. 201), as proposed originally in the GeomLoss package, later in OTT-JAX.This is needed because at $n=10^6$, storing $\mathbf{C}$ is infeasible, even across GPUs (L.204): A 1M x 1M cost matrix of float32 would require 4 Tb of memory. As a consequence, the compute cost per Sinkhorn iteration is $O(d n^2)$.

When both $n$ and $d$ are huge, we propose to alleviate this by projecting both noise and data onto the top $k$ PCA subspace of data to consider the cost $\mathbf{C}=-\Sigma\mathbf{X}^T(\Sigma\mathbf{Y})$, where $\Sigma$ is $\mathbb{R}^{k\times d}$. This results in per iteration cost of $O(k n^2+ knd)$. Note that FM training (Step. 5-8 in Alg.2) still happens in the original space, since PCA is only used to compute pairings.

When applying PCA to ImageNet-64,

• with $n=131,072$ and $𝜀=0.1$, we see up to 10x speedup over full dimension, with no impact on $\mathcal{E}$ nor downstream FM metrics.
• we can increase $n=1,048,576$ using $𝜀=0.01$ (rightmost cols), which was not feasible with full dimensionality, and achieve even better FID for low NFE.

We use the entire dataset to estimate FID statistics for this table, hence the numbers are more accurate than the original draft.

3. Precomputing Noise/Data Pairs

Our implementation, to be added to OTT-JAX upon publication, can now precompute pairs of noise / data, as claimed in L.143: We can now decouple Steps 1~4 of Alg. 2 from FM training (Steps 5-8).

• As $n$ data points are retrieved from a dataloader DL, and $n$ Gaussian noises are resampled, the outputs of Steps 1-4 (Sinkhorn + categorical sampling) are accumulated and buffered in a new augmented dataloader DL~.
• To avoid storing noise vectors, we generate each noise vector using $n$ random integer rng keys; Rather than store pairs of large vectors in DL~, we accumulate the output of Step 4 as pairs of data identifier id_{i} seed rng_{𝜎(i)} and in DL~.
• When training FM (Step 5-8), we load pairs of indices from DL~. For each rng, id pair retrieved from DL~, the corresponding data vector is retrieved and the noise vector is regenerated using the rng.

We use this approach to ablate any hyperparameter of FM training, e.g. learning rate as discussed in the answer to Reviewer 3ETv, avoiding Sinkhorn recomputations.

[J. Thornton and M. Cuturi, Rethinking Initialization of the Sinkhorn Algorithm, AISTATS23]

Many thanks for acknowledging our rebuttal. Many thanks for your encouragement and for your detailed recommendations.

Adding small $n=256$ results: We agree that adding these small $n=256$ results (which result in $n=32$ datapoints per device when running on a 8 GPU node) in the paper would be beneficial for the community.

Initially, we did not report these numbers because they were out of range for $\mathcal{R}$ in the $d=256$ column of our piecewise affine synthetic benchmark (see our Table above). We commit to:

• Report numbers for all dimensions and benchmarks in Fig. 10 / 12, and for $d=256$ do in the same way we report the IFM🔻when it is out of range in the y-axis (in Fig. 10), with an arrow and numerical values, or, possibly, use a broken y-axis (we toyed with this idea but gave up at the time due to its complexity, we will try again).
• Shorten the discussion above on sampling with/without replacement and add it to the main text
• Run $n=256$ for ImageNet 32 and ImageNet64.

In the interest of providing additional perspectives on this specific point around ultra-small $n$: we believe that very small $n$ scales not only yields worse FM metrics, but also brings additional challenges in our implementation which is tailored for larger scales:

• Small $n$ incurs a lot of compute overhead: to take a bit of an extreme argument, while running 8192 times $n=256$ small different Sinkhorn problems is much faster on paper than running a single $n=2,097,152$ problem, there is significant time wasted on GPU with our implementation when reinstantiating these problems, leading to irregular / inefficient GPU utilization, even when considering arguably the most compute intensive setting for 𝜀=0.003

Average GPU Utilization, Piecewise Affine Benchmark

• Some of our speed-up improvement, such as warmstarting and PCA, won't work as they explicitly leverage a better quality of the dual solution / dimensionality reduction, obtained as $n$ grows.
• For instance, if we provide the two tables above [Imagenet 32, Average Sinkhorn time (in seconds) to solve Alg. 1] and [Same as above, using warmstarting] in a single speed-up view, and include the smaller $n=2048$ (we initially omitted that column above because we are missing a point), we see that warmstarting is often detrimental for the smallest $n$.

ImageNet-32, speedup obtained when using warmstarting over naive implementation (>1 = better)

• This degradation for smaller $n$ is intuitive: because the sample size $n$ is small, the dual potentials estimated by Sinkhorn do not generalize well to unseen data (specially for smallest 𝜀) and fail to provide a useful initialization for the next batch.

We are extremely grateful for your time reading these our draft and our comments. We are thankful for your recommendations. Thanks again!

The authors

We are very grateful that you read our 2nd response, as well as for your score increase.

Following your comment on small OT batch size $n$ and our commitment to include such experiments in the paper, we are happy to report our preliminary experiments with small $n=256$ on ImageNet-32.

These preliminary experiments confirm that while using small OT batch sizes $n=256$ can offer some FID gains over IFM for the smallest NFEs, it may also perform worse relative to IFM for larger NFE/Adaptive solvers. On the other hand $n=256$ is always significantly worse than large $n$.

The following table shows our experiment with small OT batch size $n=256$ (i.e. $32$ images per device) on ImageNet-32, and compares it with IFM and a larger OT batch size. As seen from the table,

• large OT batch size $n= 524288$ is substantially better than small $n=256$ for all metrics.
• For large NFE, we even see that IFM edges above $n=256$ regardless of 𝜀.

Please note that we make the comparison at the 180K step checkpoint (out of $\approx 430K$ planned) since our runs with $n=256$ batch size haven’t finished yet. However, measures like FID@NFE=4/8 have already plateaued and will not improve at later checkpoints (see also our response to Reviewer 3ETv showing that ImageNet32 metrics tend to plateau from 150k iterations, and might even slightly increase again after ~250k iterations)

We are seeing very similar trends in ImageNet64 (if not worse), but we need to wait longer to see the runs conclude.

Many thanks again to you (and Reviewer 52g8) for suggesting this experiment, we agree that it completes our message by also looking at what happens at the "left end of the range", not just at very large $n$.