Generalized Schrödinger Bridge Matching

1. Remarks on global convergence

• We kindly note the significance of monotonic non-increase in the training objective, which implies training stability, particularly in the context of learning-based methods. This is the same type of analysis done for standard alternating optimization algorithms such as Expectation-maximization (EM) which only has this type of (local) convergence guarantees. Nevertheless, we do agree that convergence analysis could be improved for certain special cases of the state cost. For example, when the state cost $V(x)$ penalizes the deviation from a compact manifold (similar to the manifold cost defined in Eq 13), GSB should degenerate to Riemannian SB [1] given a sufficiently large $V$, and from which existing convergence analysis could be adopted.

• A more rigorous statement could be attained by re-examining existing proofs of global convergence for DSBM [2], i.e., when $V=0$. Below, we provide a proof sketch to illustrate the concept. Specifically, the global convergence of DSBM hinges on two steps: first, casting the l2 norm $\mathbb{E}[\int\|u_t\|^2]dt$ as the KL divergence between a controlled process and a reference process $\mathbb{Q}$, which, in such cases, is simply a Brownian motion; and secondly, showing that $\mathbb{Q}$ satisfies regularity conditions (see Prop 5 in [2]) such that the fixed point of the learning algorithm coincides with the unique solution.

• Now, it has been shown that the objective in Eq 3 can also formulated as a KL divergence [3], except that its reference process is now reweighed with the state cost: $\tilde{\mathbb{Q}} \propto \mathbb{Q} \exp(\int_0^1 V_t(X_t) dt)$. Notice that when $V:=0$, $\tilde{\mathbb{Q}}$ degenerates to $\mathbb{Q}$ as expected. This suggests that if we can identify a family of $V$ such that their corresponding $\tilde{\mathbb{Q}}$ obey similar regularity conditions, convergence analytic from [2] may be carried over for our GSBM. This is an interesting theoretical question albeit requiring additional works, nevertheless we thank the reviewer for raising these comments.

2. Differentiability of the state cost $V$

• We note that the differentiability of $V$ is required only when adopting spline parameterization (Alg 3) for solving the conditional stochastic optimal control problem (CondSOC; Eq 6). That being said, one may consider alternative solvers to CondSOC solver that, albeit less efficient, provide an approximate solution to Eq 6 yet without having to assume differentiability of $V$. This includes (i) solving Eq 6 directly with path integral as in our Prop 4, (ii) enforcing the HJB PDE in Eq 18 with an PINN loss [4,5], or (iii) expanding Eq 18 with nonlinear Feynman Kac lemma and then applying temporal difference loss, as initially proposed in [6].

• When $V$ is differentiable and known in priori---which is a realistic setup for a wide range of general distribution matching problems as demonstrated in our experiments, gradient-based solvers generally enjoy superior scalability and computational complexity, hence leveraging by our GSBM.

3. Other cases when CondSOC (Eq 6) has analytic solutions

• We've identified two other cases, apart from the configuration mentioned by the reviewer, when Eq 6 has analytic solutions. First, as briefly sketched in 1., for a sufficiently large $V$ that penalizes deviation from a compact manifold and $\sigma :=0$, Eq 6 accounts to solving the geodesic path, which admits analytic solutions for manifolds such as sphere, tori, hyperbolic ball, and SPD matrices. Secondly, for quadratic $V$ and $\sigma :=0$, Eq 6 reduces to a linear quadratic regulator with two end-point constraints. The solution relates to a Lyapunov differential equation that admits analytic solution. We detail the derivation in Appendix E (Page 25) for completeness. For general state costs considered in Sec 3.2, solving Eq 6 necessitates solving its necessary condition in Eq 18 (see Page 16). Hence, as long as Eq 18 has analytic solution for some given $V$, so does Eq 6. This is an interesting question, and we thank the reviewer for raising it.

[1] Riemannian diffusion SB (https://arxiv.org/abs/2207.03024)
[2] Diffusion SB matching (https://arxiv.org/abs/2303.16852)
[3] Linearly-solvable MDP (https://homes.cs.washington.edu/~todorov/papers/TodorovNIPS06.pdf)
[4] Physics-informed neural networks (https://arxiv.org/abs/1711.10561)
[5] Transport, variational inference & diffusions (https://arxiv.org/abs/2307.01050v1)
[6] Deep generalized SB (https://arxiv.org/abs/2209.09893)

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.

3. Strategy for initialization of spline optimization (Alg 3)

• In practice, we find that the spline optimization (Alg 3) is quite robust to hyper-parameters. In the two tables below, we report the optimized objective values (Eq 6a) by initializing Alg 3 with variations in how the controlled points are distributed across time steps (regular (first table) vs irregular (second table)) and different numbers of controlled points for parameterizing the mean $\mu_t$ and standard deviation $\gamma_t$ splines (see Eq 9). The initial objective value is much larger relatively, around 18000. Hence, it is evident that across all configurations, the optimization converges to similar values. Furthermore, we observe stable convergence in all cases. We highlight these results as a consequence of how our GSBM decomposes the GSB problem (see Sec 3.1), which yields a variational problem (Prop 2 & Eq 6) that can be optimized more simply and steadily (since it is constrained by two end-points instead of two distributions).

  	Num of controlled points in $\mu_t$ = 5	10	15	20
  Num of controlled points in $\gamma_t$ = 10	448.71	449.80	439.11	439.92
  20	454.71	452.88	439.81	439.92
  30	459.50	456.13	444.10	443.14
  40	462.24	459.23	446.05	446.34

  	5	10	15	20
  10	449.37	449.74	439.39	440.04
  20	449.22	454.84	440.91	441.28
  30	455.23	454.64	443.15	443.64
  40	464.27	454.97	446.19	445.31

1. Remarks on global convergence

• We kindly note the significance of monotonic non-increase in the training objective, which implies training stability, particularly in the context of learning-based methods. This is the same type of analysis done for standard alternating optimization algorithms such as Expectation-maximization (EM) which only has this type of (local) convergence guarantees. Nevertheless, we do agree that convergence analysis could be improved for certain special cases of the state cost. For example, when the state cost $V(x)$ penalizes the deviation from a compact manifold (similar to the manifold cost defined in Eq 13), GSB should degenerate to Riemannian SB [1] given a sufficiently large $V$, and from which existing convergence analysis could be adopted.

• A more rigorous statement could be attained by re-examining existing proofs of global convergence for DSBM [2], i.e., when $V=0$. Below, we provide a proof sketch to illustrate the concept. Specifically, the global convergence of DSBM hinges on two steps: first, casting the l2 norm $\mathbb{E}[\int\|u_t\|^2]dt$ as the KL divergence between a controlled process and a reference process $\mathbb{Q}$, which, in such cases, is simply a Brownian motion; and secondly, showing that $\mathbb{Q}$ satisfies regularity conditions (see Prop 5 in [2]) such that the fixed point of the learning algorithm coincides with the unique solution.

• Now, it has been shown that the objective in Eq 3 can also formulated as a KL divergence [3], except that its reference process is now reweighed with the state cost: $\tilde{\mathbb{Q}} \propto \mathbb{Q} \exp(\int_0^1 V_t(X_t) dt)$. Notice that when $V:=0$, $\tilde{\mathbb{Q}}$ degenerates to $\mathbb{Q}$ as expected. This suggests that if we can identify a family of $V$ such that their corresponding $\tilde{\mathbb{Q}}$ obey similar regularity conditions, convergence analytic from [2] may be carried over for our GSBM. This is an interesting theoretical question albeit requiring additional works, nevertheless we thank the reviewer for raising these comments.

2. Differentiability of the state cost $V$

• We note that the differentiability of $V$ is required only when adopting spline parameterization (Alg 3) for solving the conditional stochastic optimal control problem (CondSOC; Eq 6). That being said, one may consider alternative solvers to CondSOC solver that, albeit less efficient, provide an approximate solution to Eq 6 yet without having to assume differentiability of $V$. This includes (i) solving Eq 6 directly with path integral as in our Prop 4, (ii) enforcing the HJB PDE in Eq 18 with an PINN loss [4,5], or (iii) expanding Eq 18 with nonlinear Feynman Kac lemma and then applying temporal difference loss, as initially proposed in [6].

• When $V$ is differentiable and known in priori---which is a realistic setup for a wide range of general distribution matching problems as demonstrated in our experiments, gradient-based solvers generally enjoy superior scalability and computational complexity, hence leveraging by our GSBM.

[1] Riemannian diffusion SB (https://arxiv.org/abs/2207.03024)
[2] Diffusion SB matching (https://arxiv.org/abs/2303.16852)
[3] Linearly-solvable MDP (https://homes.cs.washington.edu/~todorov/papers/TodorovNIPS06.pdf)
[4] Physics-informed neural networks (https://arxiv.org/abs/1711.10561)
[5] Transport, variational inference & diffusions (https://arxiv.org/abs/2307.01050v1)
[6] Deep generalized SB (https://arxiv.org/abs/2209.09893)

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.

1. Summary on the new presentation

1.1 Clarification between GSB problem and GSBM algorithm

• We first thank the reviewer for the valuable and detailed comments. We agree that the relation between the practical instantiation of GSBM algorithm (Alg 5) and GSB problem itself (Eq 3) can be stated more clearly and could've led to unintentional confusion. We value academic papers as pivotal means for communication and place high standards on clear and rigorous expression (as recognized by Reviewer 7Jx7, fcZ9). Misleading readers is the last thing we wish to inadvertently propagate. We agree with the reviewer that some of our claims, as explained in the discussion below, are only based on empirical findings, not theoretical guarantees, and should be distinguished as such. We aim to clarify this in the paper.

• To restate the relationship between GSBM and GSB, and as the reviewer has already recognized: The GSBM algorithm provides an approximate solution to the GSB problem. The problem itself has great potential, albeit relatively unexplored, for tackling more general distribution matching tasks and, as shown in this work, benefits equally from recent advances in diffusion models. Meanwhile, the contribution of the proposed GSBM algorithm is to provide an approximate solution to GSB that, compared to prior deep learning-based methods, scales much more favorably to higher dimensional and more realistic setups considered in machine learning (such as images and point cloud data).

• The "approximation" of GSBM to GSB comes from two places. First is the Gaussian path approximation to Eq 6 for tractability. The second approximation, which is shared across all matching algorithms, is the fact that it may be practically unlikely to reach the unique minimizer given prescribed $p_t$ (i.e., line 1 in Alg 5 may not fully converge). This, as pointed out by the reviewer, could lead to violation of the boundary constraint. We note that we set the Gaussian path to satisfy the boundary constraints in Eq 6b by construction, so we took the reviewer's comment to mean that we cannot satisfy the distributional boundary constraints in Eq 2, which we agree with. In practice, we still found our alternating optimization approach to perform much more favorably compared to some existing deep learning-based approaches (e.g., DeepGSB [1], NLSB [2]) in approximating solutions to the same problem.

• Nevertheless, we agree that both aforementioned approximations could be stated more clearly for improved clarification. In the revision, we revise abstract, Sec 1, Sec 3.2, and conclusion. Notable changes are marked in dark red. We aim to re-emphasize that modern learning algorithms (e.g., DeepGSB [1], DSBM [4]) provide approximate solutions to GSB/SB problems, and our GSBM algorithm is no exception. Compared to prior methods [1,2,3], approximate solutions provided by "matching" algorithms (like our GSBM) remain much closer to the feasible set throughout training, and the conditions for them to preserve the exact marginals (thereby obeying Eq 6b) are noted in Footnote 2. However, when discussing such comparisons, we will make it more clear that this is based on intuitions regarding the algorithmic framework and our empirical findings, rather than a theoretical guarantee. Due to page limit, Table 1 is moved to Appendix A.

• Finally, we note that the matching algorithms in this work can be divided into 3 classes. Algorithms such as Bridge Matching, Action Matching, and, as brought up by the reviewer, Flow Matching, do not provide a means to OT/SB (as stated in our introduction, we categorize such methods as those which prescribe the probability path and does not optimize it). In contrast, algorithms such as Rectified Flow and DSBM aim to solve---approximately---OT/EOT with l2 cost via an additional outer loop similar to lines 4-5 in Alg 5. Our GSBM offers an intuitive explanation of such the second class of "optimality-enhanced" methods within the context of an alternating optimization and generalizes them beyond kinetic energy to include nontrivial state costs. We summarize the above discussions in Fig 9 on Page 14.

1.2 Reference & other adjustments

• In the revision, the original references on SB are added on Page 1. Other suggested references are added on Pages 1 & 3 and distributed in the current way due to page limit, as we're not allowed to add a new page to the main paper during rebuttal this year. We kindly note that [Liu 2022] & [Albergo 2023] were already cited in the initial submission but indeed should have been included in the reviewer's suggested places. The citep typo in Sec 3.1 is now fixed.

• We restate the matching interpretation in Sec 2 & 3.1, and add additional comments on OT with general costs on Page 2. We thank the reviewer for the meticulous reading in strengthening the paper's rigor.

[1] Deep generalized SB
[2] Neural Lagrangian SB
[3] Diffusion SB
[4] Diffusion SB matching

1.3 Discussion on feasibility

• We do agree that feasibility is often traded off with optimality, and approximate solutions by learning-based methods may inevitably encounter mismatches with feasibility constraints. However, we believe that the approach employed in learning algorithms to address feasibility are crucial factors that influence scalability, particularly when learning-based methods are applied to tackle the same problem. On one hand, methods that simply relax the hard constraints (e.g., [2]) either necessitate running adjoint methods, thereby greatly hindering scalability, and require designing some terminal cost (e.g., through adversarial training [5]), which complicates the training & tuning (see our discussion on this in Appendix A). On the other hand, methods like [1,3] are able to deal with hard constraints but still face scalability issues, as feasibility is enforced by iterative projections in path space, which requires caching trajectories at all times. Compared to the aforementioned algorithms, matching algorithms also tackle hard constraints yet are much more favorable in their practicality.

• Given the aforementioned reasoning, we compare GSBM to DeepGSB [1] in Table 1, which is now moved to Table 6 in Appendix A, as both methods aim to approximate the solutions to GSB (Eq 3) without relaxing Eq 2 (i.e., when there exists data, we should make use of it). On top of this same basis, this table aims to characterize their algorithmic distinction, essentially between matching vs non-matching methods. In the revision, we update the table and add additional comments in Sec 1 to avoid future confusion. We thank the reviewer for the comments.

2. Non-Gaussian conditional path

• An instance where the conditional density $p_{t|0,1}$ may deviate from Gaussian occurs when the solution to Eq 6 exhibits multi-modality. In such a case, Alg 3 converges to only one of the modalities. This, however, could potentially be addressed by extending Alg 3 with a mixture of Gaussian paths or simulating model coupling with multiple seeds (provided multi-modality is learned in Stage 1). We leave them for exploration in future work. We also note that the emergence of multi-modal solutions depends on the critical value of $\sigma$, as demonstrated in Fig 8. Throughout our experiments, we found that despite approximating the conditional density $p_{t|0,1}$ as uni-modal, the density $p_t$, obtained by marginalizing over $p_{0,1}$, can still exhibit multi-modality.

3. Clarification on coupling

• We first emphasize that Fig 1 is NOT a learning result, but rather an illustration of how Alg 3 functions. As Alg 3 optimizes only the intermediate samples $X_t$ given fixed coupling $(x_0,x_1)$, the initial and final empirical densities are identical by construction. More precisely, Fig 1 was generated as follows: (i) sample $(x_0,x_1) \sim \mu \otimes \nu$, (ii) initialize the mean & std of $X_t|x_0, x_1$ with Brownian bridge, (iii) call Alg 3 for optimized $X_t^\prime|x_0, x_1$, and, finally, (iv) plot $X_t|x_0, x_1$ and $X_t^\prime|x_0, x_1$ given the same set of $(x_0,x_1)$.

• In practice, we do observe different couplings with nontrivial state costs $V(x)$. This is demonstrated in Fig 16 in Appendix E, where enabling $V(x)$ encourages stronger mixing between samples, resulting in different couplings compared to those induced without $V(x)$. As for the image couplings (Fig 5, two columns annotated with "Epoch 8"), although they may appear similar at first glance, certain couplings exhibit variations in color (e.g., first sample) or distinct semantics (eyes in the third sample). We acknowledge that further improvements could be achieved by designing a different $V$ (which itself is an intriguing open question beyond the scope of this work) or considering a larger dataset.

• For unsupervised image translation, most prior works rely on qualitative comparison as in our Fig 5,6,15. Quantitative metrics are less common, and we are only aware of human evaluation on online crowd-sourcing platform [6]. This is infeasible due to the limited rebuttal period. Nevertheless, in the table below, we attempt to quantify semantically meaningful "interpolations" via FID values (lower the better). Specifically, we provide the FID scores between our GSBM vs DSBM [4] at $t$ = {0.25, 0.5, 0.75}, considering the statistics of the dog, combined dog + cat, and cat datasets, respectively. The result indicates that our GSBM is much closer to the data distributions (with ground-truth semantics) compared to DSBM.

  	t = 0.25	0.5	0.75
  DSBM [4]	196.18	248.23	259.55
  GSBM (ours)	89.66	122.85	131.03

[5] Alternating population & control NNs to solve high-dim MFG (https://arxiv.org/abs/2002.10113)
[6] Guided image synthesis & editing with SDEs (https://arxiv.org/abs/2108.01073)

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.

1. Computational analysis

• Here, we provide further computational analysis in addition to Table 5, which is included in the initial submission, showing that our GSBM algorithm runs efficiently & outperforms prior methods. Specifically, GSBM (Alg 5) can be decomposed into 3 components: (i) the matching algorithm given the probability path $p_t$ (line 1 in Alg 5), (ii) simulating the coupling (line 4), and, finally, (iii) solving the CondSOC (Eq 6) using Gaussian spline (lines 5-6). As the first two components are inherited from prior matching algorithms for SB & OT, they exhibit identical computational complexity as in training or sampling from (denoising) diffusion models.

• On the other hand, the third component (i.e., Alg 3) is introduced as a result of our alternating optimization scheme in GSBM, and we agree that an expanded computational analysis would be beneficial. Below, we report the runtime and memory complexity of Alg 3 w.r.t. varying sample sizes. These two tables correspond, respectively, to crowd navigation (dimension $d$=2) and image transfer ($d$=3x64x64=12288), serving as examples of low and high dimensionality, as suggested by the reviewer. All values are measured on a V100 32G GPU.

  sample sizes (2D)	128	256	512	1024
  runtime (sec/itr)	0.01	0.01	0.01	0.01
  memory (GB)	0.799	0.825	0.861	0.907

  sample sizes (image)	8	16	32	64
  runtime (sec/itr)	0.01	0.03	0.06	0.06
  memory (GB)	2.383	4.315	6.849	11.393

• For low dimensional problem (first table), the runtime stays constant w.r.t. sample sizes, as the GPU hasn't reached its maximum utility. For high dimensional case (second table), the runtime grows roughly linearly, which is also expected. We emphasize that in both cases, the optimization in Alg 3 exhibits fast convergence, terminated in at most 10 seconds for both 2D and images. This is attributable to the simplified boundary constraints and the significantly fewer parameters in splines compared to, e.g., the network that parametrizes the SDEs. Using image transfer as an illustration, the count of spline parameters is only 2-9% in comparison to the UNet.

• Regarding memory complexity, both tables indicate linear growth w.r.t. sample sizes. It is noteworthy that, given the rapid convergence of Alg 3, one might contemplate reducing the sample size for a more favorable memory consumption. Such reduction does not impact the solution, as the splines are constructed independently for each pair of $(x_0,x_1)$. For completeness, the table below outlines the memory consumption of other relevant baselines, highlighting that our GSBM scales to much higher dimensions, as also shown in Table 6. All values are measured on a V100 32G GPU in units of GB. Note that NLSB incurs out-of-memory error beyond 100 dimensions.

  	crowd navigation ($d$=2)	lidar ($d$=3)	opinion ($d$=1000)	AFHQ ($d$=12288)
  NLSB [1]	2.3	2.5	n/a	n/a
  DeepGSB [2]	2.8	3.1	22.3	n/a
  GSBM (ours)	0.97	1.1	16.2	30.1

2. Clarification on image data

• We emphasize that our GSBM algorithm does operate on image data within the full pixel space at resolution 64 (thus the problem dimension is $d$=3x64x64). The latent space appearing in Sec 4.2 served solely as a means to construct a state cost $V(x)$ that takes into account some semantic information for image data. As $V(x)$ is defined as a L2 regression loss between two images (see Eq 14), the latent space was never used to reduce the problem dimension. We hope the reviewer will recognize this distinction. Our generalized construction allows us to make use of nonlinear state costs, and we chose to showcase this by adapting a latent space. It may also be possible to consider other perceptual distances, such as LPIPS [3]. In the revision, we add clarification in Sec 4.2 (marked dark red above Eq 14) to avoid future confusion, and we thank the reviewer for bringing up the topic.

[1] Neural Lagrangian SB (https://arxiv.org/abs/2204.04853)
[2] Deep generalized SB (https://arxiv.org/abs/2209.09893)
[3] Unreasonable effectiveness of deep features as perceptual metric (https://arxiv.org/abs/1801.03924)

3. Strategy for initializing the controlled points in spline optimization (Alg 3)

• In practice, we find that the spline optimization (Alg 3) is quite robust to hyper-parameters. The two tables below provide a quantitative demonstration of this. We report the optimized objective values (Eq 6a) by initializing Alg 3 with variations in how the controlled points are distributed across time steps (regular (first table) vs irregular (second table) and different numbers of controlled points for parameterizing the mean $\mu_t$ and standard deviation $\gamma_t$ splines (see Eq 9). The initial objective value is much larger relatively, around 18000. Hence, it is evident that across all configurations, the optimization converges to similar values. Furthermore, we observe stable convergence in all cases. We highlight these results as a consequence of how our GSBM decomposes the GSB problem (see Sec 3.1), which yields a variational problem (Prop 2 & Eq 6) that can be optimized more simply and steadily (since it is constrained by two end-points instead of two distributions).

  	Num of controlled points in $\mu_t$ = 5	10	15	20
  Num of controlled points in $\gamma_t$ = 10	448.71	449.80	439.11	439.92
  20	454.71	452.88	439.81	439.92
  30	459.50	456.13	444.10	443.14
  40	462.24	459.23	446.05	446.34

  	5	10	15	20
  10	449.37	449.74	439.39	440.04
  20	449.22	454.84	440.91	441.28
  30	455.23	454.64	443.15	443.64
  40	464.27	454.97	446.19	445.31

4. Gaussian path approximation

• This is an interesting question, and we thank the reviewer for raising it. We do acknowledge that Gaussian path approximation to the CondSOC (Eq 6) may be different from the actual minimizer necessary to establish the convergence analysis in Thm 5. Specifically, Thm 5 & 6 assume ideal conditions when Stages 1 & 2 converge to the actual global minimum. Although these ideal conditions may not always be met in our practical instantiations of GSBM, we emphasize that they offer valuable insights that validate our alternating optimization procedure.

• That being said, one might explore alternative solvers for CondSOC that, albeit are less efficient, eliminate the need of Gaussian pat approximation in Stage 2. This can include (i) solving Eq 6 directly with path integral as in our Prop 4, (ii) enforcing the HJB PDE in Eq 18 with an PINN loss [4,5], or (iii) expanding Eq 18 with nonlinear Feynman Kac lemma and then applying temporal difference loss, as initially proposed in [2]. In practice, we find that Gaussian path approximation within GSBM offers much better tradeoffs between efficiency/scalability and optimality, as compared to other deep learning-based approaches.

• With regards to other scenarios in which Eq 6 admits closed-form expressions, we've identified two cases. First, for a sufficiently large $V$ that penalizes deviation from a compact manifold and $\sigma :=0$, Eq 6 amounts to solving the geodesic path, which admits analytic solutions for manifolds such as sphere, tori, hyperbolic ball, and SPD matrices. Secondly, for quadratic $V$ and $\sigma :=0$, Eq 6 reduces to a linear quadratic regulator with two end-point constraints. As we detail in Appendix E (Page 25), its solution relates to a Lyapunov differential equation that admits analytic solutions. For general state costs considered in Sec 3.2, solving Eq 6 necessitates solving its necessary condition in Eq 18 (see Page 16). Hence, as long as Eq 18 has analytic solution for some given $V$, so does Eq 6.

[4] Physics-informed neural networks (https://arxiv.org/abs/1711.10561)
[5] Transport, variational inference and diffusions (https://arxiv.org/abs/2307.01050v1)

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.