Statistical Analysis of the Sinkhorn Iterations for Two-Sample Schr\"{o}dinger Bridge Estimation

We thank the reviewer for reading our paper.

Accuracy and computational cost of Sinkhorn bridge

• Accuracy: The Sinkhorn bridge estimator performs well in terms of accuracy. In fact, [1] observed that it outperforms Schrödinger bridge solvers based on learned forward–backward dynamics under the experimental setting of [2]. For details, see Table 1 in [1].

• Computational cost: The Sinkhorn Bridge requires the coupling $\pi_{m,n}$ which can be obtained using Sinkhorn iterations. These iterations converge exponentially fast and have a per-iteration cost of $mn$. In the sampling phase, the Sinkhorn bridge applies Euler–Maruyama discretization, which incurs a cost of $O(n)$ per discretization step. Therefore, generating a new sample from the estimated target distribution costs $O(nK)$, where $K$ is the number of discretization steps. To reduce the inference cost, we also proposed a neural network-based approximation of the drift.

[1] Pooladian, A.-A. and Niles-Weed, J. "Plug-in estimation of Schrödinger bridges." arXiv:2408.11686, 2024.

[2] Gushchin, N., et al. Building the bridge of Schr ̈odinger: A continuous entropic optimal transport benchmark. NeurIPS, 2023.

Neural network based drift approximation

We note that Algorithm 1 performs online SGD to minimize $\mathcal{L}(b)$ (line 210), and this is feasible since $\mathcal{L}(b)$ is tractable. Specifically, the expectation in $\mathcal{L}(b)$ is taken with respect to a mixture of Brownian bridges defined by the transportation plan $\pi$, from which we can sample to perform SGD. As a result, there are approximation errors (due to model expressivity) and optimization errors (due to finite training iterations), but with sufficient network capacity and training time, the NN-based estimator will converge to the Sinkhorn bridge drift, which is the minimizer of $\mathcal{L}(b)$.

We thank the reviewer for carefully reading our paper.

Tradeoff between $n$ and $R$

Several recent works that established the faster convergence rate of $O(1/n)$ for EOT assumed either the boundedness of $R$ or Lipschitz continuous bounded cost (i.e, the bounded $R$ with Euclidean cost.) These results suggest a trade-off between $n$ and $R$. However, the optimal dependence on these constants remains unclear. Our bound exhibits exponential dependence on $R$, but we note that such dependence is common in the literature of the entropic transport [1,2]. In fact, the term $\exp(CR^2)$ in our bound is inherited from the existing bound on the Schrödinger potential functions [2]. More recently, [3] established the polynomial bound of the form $R^d$ for $L^2$-distance between potential functions. However, translating their result for EOT into statistical bounds for the Schrödinger bridge is nontrivial and is a future research direction.

[1] Rigollet, Philippe, and Austin J. Stromme. "On the sample complexity of entropic optimal transport." The Annals of Statistics, 2025.

[2] del Barrio, Eustasio, et al. "An improved central limit theorem and fast convergence rates for entropic transportation costs." SIAM Journal on Mathematics of Data Science, 2023.

[3] Stromme, A. J. "Minimum intrinsic dimension scaling for entropic optimal transport." In International Conference on Soft Methods in Probability and Statistics, 2024.

Dependence on (intrinsic) dimension

Our bound relies on the following decomposition: $\mathbb{E}[TV^2(P\_{\infty,\infty}^{\*,[0,\tau]},P\_{m,n}^{\*,[0,\tau]} )] \lesssim \mathbb{E}[KL(P\_{\infty,\infty}^{\*,[0,\tau]} | \bar{P}\_{\infty,n}^{[0,\tau]})] + \mathbb{E}[KL(\bar{P}\_{\infty,n}^{[0,\tau]} | P\_{m,n}^{\*, [0,\tau]}) ]$.

The first term admits an upper bound of the form: $\frac{R^2}{n (1-\tau)^{d_{\nu}+2} \varepsilon^{d_{\nu}}}$ (Proposition 4.3 in [4]) when the support of $\nu$ is a $d_{\nu}$-dimensional submanifold. However, for the second term, we have not obtained a bound that depends on the intrinsic dimension $d_{\nu}$. As a result, the overall bound depends on the ambient dimension $d$ rather than $d_\nu$.

[4] Pooladian, A.-A. and Niles-Weed, J. Plug-in estimation of Schrödinger bridges. arXiv:2408.11686, 2024.

Weakness 2: relationship between Algorithm 1 and plug-in estimator

We note that Algorithm 1 performs online SGD to minimize $\mathcal{L}(b)$ (line 210), and this is feasible since $\mathcal{L}(b)$ is tractable. Specifically, the expectation in $\mathcal{L}(b)$ taken with respect to a mixture of Brownian bridges defined by the transportation plan $\pi$, from which we can sample to perform SGD. As a result, Algorithm 1 converges to the plug-in estimator, and no additional statistical error is incurred if the neural network is sufficiently large. It is also important to note that $N$ denotes the mini-batch size of SGD and is unrelated to the definition of the objective $\mathcal{L}(b)$.

Minor typo

Thank you for noticing the typo. It has been corrected accordingly.

We thank the reviewer for engaging with us in the discussion.

2: Statistical error in Algorithm 1

We would first like to clarify an important point that may have been misunderstood in the review. It seems the reviewer may have interpreted Algorithm 1 as optimizing an empirical approximation to $\mathcal{L}(b_\theta)$, possibly because we introduce the empirical approximation $L(\theta)$ in Algorithm 1. Also we stated "SGD can be efficiently applied to minimize $L(\theta)=\mathcal{L}(b_\theta)$" in the submission. This is actually our typo, which is correctly "SGD can be efficiently applied to minimize $\mathcal{L}(b_\theta)$"

We emphasize that Algorithm 1 does not optimize empirical loss $L(\theta)$. Instead, $L(\theta)$ is used solely to define a stochastic gradient at each iteration, and the mini-batch that defines $L(\theta)$ is resampled every time. Therefore, Algorithm 1 performs SGD directly on the true objective $\mathcal{L}(b_\theta)$, and convergence to a stationary point is guaranteed even when using $N=1$, where $N$ is simply the mini-batch size in SGD.

To make this clearer, let us abstract the setting using a general form: suppose the objective function is $F(\theta) = E_Z[ l(\theta, Z)]$, where $Z$ follows a continuous distribution. Then, Algorithm 1 with $N=1$ reduces to the standard online SGD update $\theta_{k+1} = \theta_k - \eta \nabla_\theta l(\theta_k, Z_k)$ where $\\{Z_k\\}_k$ are i.i.d. samples. The convergence of such SGD with small or diminishing step-size to a stationary point of $F(\theta)$ is well-established in the literature. For example, see Theorems 4.8–4.10 and related corollaries in the following reference. (Actually, the objective $F$ in these results covers the expected loss. See Eq. (4.1) in [5].)

[5] L Bottou, F.E. Curtis, and J. Nocedal. Optimization Methods for Large-Scale Machine Learning https://arxiv.org/pdf/1606.04838

As for concerns regarding the approximation error due to neural networks and global convergence, we believe these are minor. Most neural network-based diffusion models and Schrödinger bridge methods face similar challenges. That said, we note that the objective $\mathcal{L}(b)$ is (linearly) convex in $b$, so global convergence can be ensured using recent results from mean-field theory [6, 7, 8].

We would be happy to address any further concerns during the discussion period.

[6] L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. NeurIPS, 2018

[7] A. Nitanda, D. Wu, and T. Suzuki. Convex Analysis of the Mean Field Langevin Dynamics. AISTATS, 2022.

[8] L. Chizat. Mean-field Langevin dynamics: exponential convergence and annealing. TMLR, 2022.

1: Tradeoff between $R$ and $n$

Thank you for the suggestion. Following your advice, we will include a comment to address this point in the final version.

We thank the reviewer for reading our paper.

Relevance to NeurIPS community

Schrödinger bridges (SB) have attracted significant attention at NeurIPS as a key research theme at the intersection of diffusion models and optimal transport. Notably, several influential papers on SB have been published at NeurIPS, including Diffusion Schrödinger Bridge (NeurIPS 2021) [1] and Diffusion Schrödinger Bridge Matching (NeurIPS 2023) [2]. Our work reinforces the theoretical foundations that underpin this growing line of research.

On the applied side, SB techniques have already been utilized in a wide range of real-world tasks. These include image restoration and transformation using I$^2$SB [3] (e.g., inpainting and super-resolution), unpaired image translation via UNSB [4] and SB Flow [5], domain adaptation through distribution alignment using maximum-likelihood SB solvers [6], and trajectory inference for single-cell RNA-seq data in the life sciences [6, 7].

Taken together, the theoretical contributions of this work strengthen the foundation of SB-based methods and enhance our understanding of their practical applicability, making the results broadly relevant to the NeurIPS community. We will include the above related works in the final verstion.

[1] De Bortoli, Valentin, et al. "Diffusion Schrödinger bridge with applications to score-based generative modeling." NeurIPS, 2021.

[2] Shi, Yuyang, et al. "Diffusion Schrödinger bridge matching." NeurIPS, 2023.

[3] Liu, Guan-Horng, et al. "I$^ 2$SB: Image-to-Image Schrödinger Bridge." arXiv:2302.05872, 2023.

[4] Kim, Beomsu, et al. "Unpaired image-to-image translation via neural Schrödinger bridge." arXiv:2305.15086, 2023.

[5] De Bortoli, V., Korshunova, I., Mnih, A., & Doucet, A. Schrödinger bridge flow for unpaired data translation. NeurIPS, 2024

[6] Vargas, Francisco, et al. "Solving Schrödinger bridges via maximum likelihood." Entropy, 2021.

[7] Chizat, Lénaïc, et al. "Trajectory inference via mean-field Langevin in path space." NeurIPS, 2022.

We thank the reviewer for carefully reading our paper.

The role of epsilon

Revisiting the proof, we improved the exponential factor in equation (14) to $\exp\bigl(R^2(60+\varepsilon^{-1})\bigr),$ and the following is written with that improvement in mind (to be reflected in the final version).

The numerically unstable parts are the Sinkhorn Algorithm and the drift $b\_{m,n}^{(k)}$ with the instability being especially pronounced in the drift. This stems from the term appearing in equation (13)

where the denominator contains $(1-t)\varepsilon$; as $t$ approaches $1$, this induces sharp scaling. However, since the expression is in a softmax form, standard stabilization by subtracting the maximum of the exponent’s argument can be applied to ensure numerical stability.

(i) To verify the empirical behavior, we ran the experiment for finding $\varepsilon$ that achieves target error $\delta=1.0$ under the same setup as in “Experiments (Statistical convergence)”, for given $m,n$. As predicted, the error drops sharply as $\varepsilon$ increases, then transitions to a regime of slower, asymptotic improvement—i.e., a steep improvement in $m\wedge n$ followed by a gradual change. The observed values were:

(ii) This phenomenon is supported theoretically: from equation (14), the $\varepsilon$ required to achieve target error $\delta$ satisfies

This indicates that when $m\wedge n$ is moderately large, error reduction proceeds rapidly, and thereafter improves only logarithmically.

Experiments and d

We conducted experiments to verify the dependence of the Sinkhorn bridge estimator on the intrinsic dimensionality. The experimental setup is as follows: the ambient dimension $d = 50$, the regularization strength $\varepsilon = 0.5$, and endpoint $\tau = 0.9$. We sampled the source distribution $\mu$ uniformly from the unit cube $[0,1]^d$, and the target distribution $\nu$ uniformly from the $d_\nu$‑dimensional unit sphere embedded in the ambient space $\mathbb{R}^d$: $\\{ x \in \mathbb{R}^{d_\nu} | \||x\||_2 = 1\\} \subset \mathbb{R}^d$.

By varying $d_\nu$, we examined whether the convergence behavior of the estimation error with respect to sample sizes depends on the manifold dimension. For $d_\nu \in \\{5,\,10,\,20,\,30\\}$, we compared the time‑integrated mean squared error $\int_{0}^{0.9} \mathrm{MSE}_{\mathrm{sample}}(m,n,t)\,\mathrm{d}t,$ and observed that the error tends to increase as the manifold dimension grows, suggesting room for improvement.

Stopping and guidance

Considering the decomposition of TV-distance, the number of iterations $k$ satisfying the following inequality is sufficient: $\mathbb{E}[ TV^2(P\_{\infty,\infty}^{\*,[0,\tau]}, P\_{m,n}^{\*,[0,\tau]})] \geq \mathbb{E}\bigl[TV^2(P\_{m,n}^{\*,[0,\tau]}, P\_{m,n}^{k,[0,\tau]})]$. This implies $k \geq C\frac{\log\bigl(A/n + B/(m\wedge n)\bigr)} {\log\bigl(\tanh(R^2/\varepsilon)\bigr)}$ where $A = \frac{1}{(1-\tau)^{d+1} \varepsilon^{d-1}R^2}$ and $B = \frac{C_{d}R^{5(d+1)}}{(1-\tau)\tau\varepsilon} e^{R^2\bigl(60 + \varepsilon^{-1}\bigr)}$. Since $\log(\tanh(R^2/\varepsilon))<0$, the numerator $\log\bigl(A/n + B/(m\wedge n)\bigr)$ must also be negative for the right-hand side to be meaningful. This in turn implies that $m \wedge n \approx \max\{A, B\}$ samples are needed.

This prediction is also confirmed by Figure 3. In particular, $b_{m,n}^{(1)}$ obtained after a single Sinkhorn iteration, already transports points to their nearest neighbors and captures the bulk of $\nu_n$.

The only remaining discrepancy with the optimal drift $b_{m,n}^*$ lies in the mismatch of the marginal distributions, similar to the case of entropic optimal transport. As a result, a sufficiently large $k$ becomes necessary only when the empirical distributions $\mu_m$ and $\nu_n$ are already close to the true distributions $\mu$ and $\nu$ i.e., when $m$ and $n$ are sufficiently large.