Feedback Schrödinger Bridge Matching

4. Dynamic Formulation

Our approach reinterprets the static modified EOT problem in Eq. (8) of the manuscript through the prism of entropic interpolation [6], the solution of dynamical OT formulation corresponds to solving ‘local’ versions of optimal transport problems between infinitesimally close distributions [7]. Starting from the EOT problem, from $t_0=0$ and $t_N=1$, we can consider any intermediate point $t_1$ and find the optimal paths $p_1^\star$, $p_2^\star$ between the endpoints in the two newly formed subproblems (i.e., from $t_0$ to $t_1$, and $t_1$ to $t_N$). The union of the optimal subpaths forms the optimal policy of the initial problem, and the converse is also true. This can be extended for any number of intermediate points (which is equivalent to taking $\delta t\rightarrow 0), leading us to to consider ‘local’ versions of optimal transport problems between infinitesimally close distributions [7]. We have also provided a sketch of proof for Theorem 3.2 to enhance readibility.

5. Contradiction in the Hamiltonian

We would like to thank the reviewer, for identifying these contradictions attributed to typos and error in the notation in the proof of Proposition A.1. More specifically, the Lagrange multiplier in Eq. (45) is not the same as the vector field $a_t$ of the Hamiltonian $H(x_t, a_t, t)$, but rather for Lagrange multiplier $s_t$ it holds that $a_t = \nabla s_t$. We have corrected the notation errors and typos and further clarified the proof of the Proposition A.1.

6. Advantages of current FSBM framework

We emphasize that the dynamic formulation of our FSBM framework offers significant advantages, in terms of faster training and improved generalization. These benefits, stemming from leveraging the guidance of coupled pairs, were observed primarily through extensive empirical analysis. The faster training is intuitively explained by the use of aligned pair information to guide the matching of non-aligned data. While there is no theoretical evidence to verify our empirical findings yet, we aim to address this in future work by deriving theoretical bounds for our FSBM, proving that the utilization of keypoints results in a transport map which is better capable at generalizing.

References

1. Xiang Gu, Yucheng Yang, Wei Zeng, Jian Sun, and Zongben Xu. Keypoint-guided optimal transport, 2023b
2. Xiang Gu, Liwei Yang, Jian Sun, and Zongben Xu. Optimal transport-guided conditional score-based diffusion model. Advances in Neural Information Processing Systems, 36:36540–36552, 2023a
3. Chi-Heng Lin, Mehdi Azabou, and Eva L Dyer. Making transport more robust and interpretable by moving data through a small number of anchor points. Proceedings of machine learning research, 139:6631, 2021
4. Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A Theodorou, and Ricky TQ Chen. Generalized Schrödinger bridge matching. In International Conference on Learning Representations, 2024.
5. Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024
6. Ivan Gentil, Christian Léonard, and Luigia Ripani. About the analogy between optimal transport and minimal entropy. In Annales de la Faculté des sciences de Toulouse: Mathématiques, volume 26, pp. 569–600, 2017.
7. Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.

We would like to thank the reviewer for their positive remarks, as well as their useful comments and further suggestions for improving our work. In the following we address each of the concerns and suggestions of the reviewer.

1. Motivation and Meaning of guidance function $G$

• The key contribution of our work was the derivation of a more general dynamic formulation. Our novelty lies in the fact that our analysis starts from a static regularized EOT problem (Eq. (8)), and derive a dynamic formulation (Eq. (9)). Importantly, we highlight that our methodology can be adapted to any class of differentiable regularizing functions, not just guidance regularizing function. Our approach bridges between more general static and dynamic problems, enabling us to harness the flexibility of adding regularization terms in static EOT formulation, and then acquire a more complete description of the transport and leverage recent advancements in matching frameworks.
• The intuition behind the utilization of the guidance function $G$ in our approach is to capture structural information from the source distribution and propagate them throughout the entire trajectory. Importantly, the conditioning on the unpaired sample $x_0$ acts as an anchor. Our guidance function $G$ groups the unpaired data into clusters around the KPs in the source distribution, and penalize the deviation of the unpaired samples from the trajectory that their assigned keypoint sample follows, preserving structural and distance information in the dynamic transport map. This is performed by evaluating the distances $d(X_0, x_0^i)$ which corresponds to the distance of an uncoupled sample $X_0$ to the $i^\text{th}$ keypoint $x_0^i$, and similarly the distance $d(X_t, x_t^i)$ for every time step $t\in [0,1].$ This forces unpaired samples to preserve their relative distance from the closest keypoint sample in the source distribution at each time-step.
• Regularization using a guidance function has been implemented in previous works within the realm of semi-supervised Optimal Transport. These approaches leverage the guidance function $G$ to encode information about the relationship between unpaired data and keypoints. For example, in [1] the authors investigated incorporating various distinct $G$ functions into the Optimal Transport framework as regularization terms. They proposed a relation-preserving keypoint-guided model that steers OT matching by preserving both the correspondence between keypoint pairs and the relationships of each data point to the keypoints. While the utilization of keypoints [1] is similar to our methodology, the nature of the guidance function differs. In our approach, the guidance function acts in a distance preserving manner. More specifically, $G$ creates clusters of unpaired points around their closest KeyPoint in the source distribution, and penalize the unpaired samples when they deviate from the trajectory that their assigned keypoint sample follows. This contributes to preserving structural information in the dynamic transport map. A pertinent passage has been added in the appendix of the manuscript presenting more examples and illuminating the motivation behind the regularization with the guidance function $G.$

2. Stochastic Nature of Keypoint coupling

The guidance function is constructed deterministically with respect to the keypoints, similar to approaches seen in prior works [1, 2, 3]. However, this does not compromise the overall stochastic nature of the keypoints, as they are solutions of the static Schrodinger Bridge, which is inherently stochastic. In other words, sampling the keypoints from the static SB ensures that randomness is a fundamental component of the solution, rendering them valid for our guidance term. In practice, for our experiments we acquired the optimal pairings by sampling from the prior stochastic frameworks that approximate the Schrodinger Bridge (e.g., GSBM [4], LOSBM[5]). Given the inherent stochasticity of the static SB, there may be extreme instances where high stochasticity result in the keypoints coupling not providing with informative guidance. In such cases, we can easily reduce the influence of the guidance by adjusting its weighting. A pertinent comment on the stochasticity of the coupling of the keypoints has been added in the methodology section of the paper (line 183).

3. Update of Figure 3

We would like to thank the reviewer for their suggestion to improve Figure 3, in order to better explain the effect of $G$ on the guidance of the trajectories. We substituted the qualitative Figure 3 with a more rigorous representation of snapshots of the vector field of $\nabla G$ at times $t = \{0, \ 0.33T, \ 0.67T, \ T\}$. In the new Figure, it is evident that the gradient of the guidance function pushes the non-coupled data towards the trajectory of the KeyPoints.

We thank the reviewer for the reply. In the following points, we address the additional concerns and questions.

4. Dynamic Formulation

• The sketch of the proof starts with only two optimal subpaths to build the intuition that it is feasible and rigorous to start from the static EOT and to add intermediate points breaking down the initial problem into subproblems., whose optimal paths can combine to yield the optimal path of the initial problem. Subsequently, we mention in the sketch of proof that this can be extended into any number of intermediate paths i.e., $N$. As $N\rightarrow \infty$ , this leads us to to consider ‘local’ versions of optimal transport problems between infinitesimally close distributions.
• In our proof, we start from the EOT problem, from $t_0=0$ and $t_N=1$, directly consider $N$ distinct intermediate points at time steps $t_i$ for $i \in [0, N]$, discretizing the intermediate path into $N$ subpaths. It is highlighted that this discretization is an intermediate step to derive the dynamic expression from the static regularized EOT, by taking $N\rightarrow \infty$. This step of discretizing and considering infinitesimally many subpaths discretization is integral at recasting static OT problems into the dynamic expressions that have been widely celebrated recently [1, 2]. Therefore, all dynamic SB frameworks and formulations (including our FSBM) can be interpreted as finding $N$ optimal intermediate subpaths $p^\star_{t_i, t_{i+1}}, \ i \in [0, N-1]$ for $N\rightarrow\infty$ from the prism of dynamic EOT, as it holds that the union of these infinitesimally close and infinitesimally small optimal subpaths, is equivalent to finding the continuous optimal marginal path $p_t^\star$.
• Finally, having obtained the dynamic formulation we can integrate between the two endpoints, i.e., from $t=0$ to $t=1$, and, given the stochastic dynamics governing the path $X_t$, any function $f(X_t)$ satisfies the Ito’s Lemma.

We sincerely thank the reviewer for their comment and appreciate their acknowledgment that our responses have effectively addressed most of their concerns.

Optimization of intermediate path $p_{t|0,1}$

It is highlighted that we do not need to obtain $p_{t|0,1}$ explicitly. Conversely, for our analysis, we need to be able to:

1. efficiently sample $X_t\sim p_{t|0,1}$
2. compute the optimal drift from $p_{t|0,1}$: $u_{t|0,1}$ which is sufficient for computing the following matching loss.

While, in general, it is correct that $p_{t|0,1}$ does not admit closed-form, but if we are willing to restrict the solution space to Gaussian paths $p_{t|0,1}\approx \mathcal{N}(I_{t|0,1}, \sigma_t \mathbf{I}d)$, as is the standard practice in state-of-the-art frameworks [1], then we can achieve both (1) and (2) very efficiently. In fact, we can easily derive a closed-form solution for the optimal conditioned drift in a simulation-free manner (see Eq. (11)). In our future work, we aim to go beyond Gaussian paths, for instance, investigate stein-variational paths for multi-modal exploration.

Lastly, the iterative process to optimize with respect to the intermediate path is described in the next steps. Let $p_{t|0,1}^i$ be the initial interpolating path between the two marginals. We sample $X_t\sim p_{t|0,1}^{i}$, and evaluate the guidance function $G(X_t, x_0)$ and the conditional drift from Eq. (11). Subsequently, we evaluate the objective in Eq. (10), and take a gradient step to obtain the updated intermediate path $p_{t|0,1}^{i+1}$, and use it to sample in the next iteration. This process is repeated until convergence, where $p_{t|0,1}^\star$ is obtained. It follows that from $p_{t|0,1}^\star$, the corresponding intermediate drift is the optimal conditional drift $u_{t|0,1}^\star$ .

References

1. Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A Theodorou, and Ricky TQ Chen. Generalized Schrödinger bridge matching. In International Conference on Learning Representations, 2024.

We truly appreciate the encouraging comments of the reviewer, as well as their useful remarks. In the following we address the points raised by the reviewer.

1. Further Simplifications and Explanations of theoretical Results

In the updated version of the manuscript, we have refined our explanations on key points of our framework to enhance readability. Specifically, we have elaborated on how our guidance function works in our dynamic formulation. Additionally, we have included a Figure representing the vector field of the gradient of the guidance function better capturing the effect of the guidance function and how it pushes the uncoupled data towards the keypoint trajectories. Additionally, we have clarified the importance and the role of the outputs of each optimization in the alternating optimization scheme. Lastly, we also provided a sketch of the proof of Theorem 3.2 in the Appendix to simplify our theoretical results and enhance comprehension. Lastly, the typos in the equation have been carefully detected in the paper and corrected.

2. Discussion of the limitations of FSBM and suggestions for future research

While our FSBM displayed great capabilities in distribution matching, leveraging effectively the information of partially aligned datasets to guide the transport map, we acknowledge that our framework greatly relies on the the availability and the quality of the available KP pairs. The manner in which we sample from solution of the static SB does not ensure that all modes are equally and satisfactorily represented. This implies that in cases where a minimal number of these pairs are available, the regions of attraction formed around these may not sufficiently cover the diversity of the source distribution. This could lead to wrongful guidance, suboptimal matching and poor generalizabilty. A pertinent section discussing the limitations of our framework has been added to the paper.

In future work, we have identified several directions which could improve our work even further. For instance, we aim to devise a more intricate mechanism to sample from the solution of the static SB, while guaranteeing that all modes in highly diverse datasets are covered and represented through the sampled KPs. Additionally, we aim to further explore the capabilities of the guidance function and devise more task-specific function which furhter improve perfomance generalizability by optimizing the information offered by the KP samples. Additionally, we intend to improve the way we impose guidance, by allowing the aligned pairs to preserve the relative distance to k-nearest neighboring KPs, instead of simply the nearest one.

3. More experimental results

We provided an extensive experimental comparison of our FSBM with other state-of-the-art mathcing frameworks, such as DSBM regarding the optimality of the coupling in the generated images and the preservation of structural information. Our FSBM is a novel method further enforcing such structures—much explicitly—via the semi-supervised guidance, and our method does lead to “better coupling” across the suggested metrics, such as L2 norm, LPIPS, and SSIM, as reported below. It is clear that FSBM outperform other SB baselines (DSBM [1] and LOSBM [2]) across almost all metrics. Note that the values of the $L2$-norm have been divided by the total number of pixels in the images (i.e., $1024\times 1024$)

References

1. Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching, 2023.
2. Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024

3. Number of key points on performance and comparison between guidance schemes.

• For our crowd-navigation experiments, both approaches were examined and found the distance-preserving scheme to be more suitable for our framework. Although the relation-preserving scheme was indeed the more effective choice in the static setting in [2], this was not the case in our applications. This could be attributed to the fact that the softmax averaging relation preserving scheme works more effectively in static frameworks. However, in our dynamic formulation, we observed empirically that the softmax averaging in relation-preserving schemes (similar to [2]) failed to generate a gradient field capable of providing effective guidance to the unpaired samples. In contrast, the distance-preserving scheme offered significantly better guidance. We have added a comparison of the gradient field generated by the relation-preserving and distance-preserving functions in the Appendix Section C of the updated manuscript.
• In future work, we plan to investigate this phenomenon in greater detail and examine whether a combined weighted scheme, of a distance-preserving and a relation-preserving term will further enforce structural information during the transport process. It is highlighted that our main contribution lies in the derivation of a dynamic transport map introducing a more general dynamic formulation, which effectively demonstrates that we can have a semi-supervised matching framework, leveraging optimal pairings. Furthermore, it is worth noting that our analysis can be adapted to any class of differentiable regularizing functions. Identifying the best possible guidance function for every use case was beyond the scope of this work.
• Finally, we would like to clarify the findings in Fig. 9, which is accompanied by the following table in the updated version of the manuscript. It is noted that Figure 9 in the previous version of the paper has been moved to the Appendix Section C.1.1 and is now Figure 13. First, it is noted that our algorithm is trained to transport the samples $X_0\sim\pi_0, \text{ to }X_1\sim \pi_1$. In the evaluation phase, we consider the following distinct tasks:

  1. Vanilla: we draw new samples from the same distribution $\pi_0 = \mathcal{N}(\mu_0, \sigma_0)$, and try to transport them to the same target distribution $\pi_1$.

  2. Perturbed Mean: we draw new samples from $\pi_0^\prime$ , where $\mu_0^\prime = \mu_0 + \delta\mu$ with $\delta \mu$ being the perturbation to the mean. The goal is still to transport the unpaired samples to the same target distribution $\pi_1$.

     Percentage of KPs (%)	$\mathcal{W}^2$ between $\pi_1$ and $p_1^\theta$, $x_0\sim\pi_0$ (Vanilla)	$\mathcal{W}^2$ between $\pi_1$ and $p_1^\theta$, $x_0 \sim \pi_0^\prime$ (Perturbed Mean)
     1	0.034	0.77
     2.5	0.027	0.55
     5	0.018	0.79
     8	0.015	0.82
     10	0.013	0.81
     12	0.010	1.01
     15	0.008	1.23
     20	0.008	1.48

where with $p_1^\theta$ we denote the generated distribution. It is demonstrated that the vanilla performance of our framework continues to increase as the number of aligned pairs increases. However, it is shown that the empirical generalization to unseen initial conditions (e.g. when the mean of the initial distribution is perturbed) does not necessarily improve by adding more KeyPoints, suggesting an overfitting like phenomenon. Another explanation is that the hyperparameters of the framework were tuned to maximize the performance while using a small number of KeyPoints samples. As a result, increasing the number of aligned pairs may require re-tuning the hyperparameters to maintain optimal performance. Lastly, it should be mentioned that even this decrease in generalizability is still very small compared to the benchmark methods.

We would like to thank the reviewer for their positive comments, as well as their constructive criticism and further suggestions for improving our work. In the following we address the points raised by the reviewer.

1. Experiments to quantify information and structure preservation in transported images

The reviewer is correct that, for data-to-data translation tasks, we aim to preserve structures of original samples after transportation. In most prior OT methods, this is enforced—implicitly—by using a latent space that is assumed to preserve similar structures for points close to each other. Our FSBM is a novel method further enforcing such structures—much explicitly—via the semi-supervised guidance, and our method does lead to “better coupling” across the suggested metrics, such as L2 norm, LPIPS, and SSIM, as reported below. It is clear that FSBM outperform other SB baselines (DSBM [1] and LOSBM[3]) across almost all metrics. Note that the $L2$-norm is divided by the number of pixels in the images (i.e., $1024\times 1024$).

Finally, we also used ChatGPT to assess the image couplings between 100 pairs of input and generated images in each translation task. Given that the coupling from LOSBM was considerably worse than the coupling generated by the other two methods, we only performed comparison between DSBM [1] and our FSBM. ChatGPT was instructed to give a score on the coupling quality on a scale 0-5, with 5 being the best. The coupling quality was assessed using the following 3 criteria: i) Background Consistency: consistency in terms of color, lighting, background objects, ii) Identity Preservation: maintain the identity of the original face in terms of facial features and overall resemblance iii) Structural Integrity: preserve structural integrity of the input image in terms of accessories worn, pose, orientation. The results in 4 translation tasks presented in the paper are shown in the Tables below.

2. Justification for selection of condition to $X_0$

We thank the reviewer for bringing up this question. While we do acknowledge that, the reason for selecting to condition on the source distribution was mainly empirically driven, we highlight that, from a mathematical perspective, our derivation could follow the same steps, and Eq. (9) would use $X_1$ as the reference point instead of $X_0$. That said, the main contribution and novelty of our proposed framework is not the specific instantiation of guidance functions in Eq 7, but rather a general semi-supervised matching framework for any guidance functions that can be problem-specific.

Crucially, conditioning the guidance function on samples of either function serves as an anchor (similar to [2]) to guide the interpolation trajectories of the unpaired samples and their coupling to the opposite marginal. That said, empirical results indicate that conditioning on the terminal distribution slightly increased the $\mathcal{W}^2$ distance between the generated and ground-truth terminal distributions across all tasks, as detailed in the table below.

In future work, we also aim to enhance our guidance strategy by enabling KeyPoint pairs to maintain relative distances to k-nearest neighboring keypoints, rather than solely focusing on the nearest one. Regarding conditioning on a weighted combination of both $X_0$ and $X_1$, it is not straightforward how to impose such conditioning on the guidance function in our current mathematical formulation. That said, this is definitely an intriguing idea we had not considered before, and we would like to thank the reviewer for suggesting it.

Thanks for the response. My major concerns were addressed, and I am happy to increase the score for the current version. However, I still believe the only metric of distribution difference between generated data and real data does not truly reflect the importance of paired data, though it shows improved results over baselines. I suggest the authors investigate more real-world applications (e.g., biological data) or metrics to show the significance of introducing the paired data, maybe in future work.

We would like to thank the reviewer for valuing our work.