Topological Schrödinger Bridge Matching

Dear reviewer, thank you for taking the time to review our work. We appreciate your feedback. Please find our response to your questions: **on code implementation:** We confirm that the code will be released upon publication. **on broadening the experimental scope:** We included additional experiments on human brain networks and a single-cell dataset for matching tasks. They are conducted on real-world data and make no assumptions of Gaussian boundaries. Please see our [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) for the details and results. We hope they, together with those on real-world seismic, traffic and ocean flow datasets, show the practical generalizability of the method. **on the Gaussian boundary case**: If we may, we want to clarify the main focus of the paper is in terms of _general_ boundaries. In Section 4, we derived the closed-form for the Gaussian case. This study is however of relevance for both theory and practice, offering insights beyond the Gaussian case. Specifically: 1. _understanding topological SB for general cases_. This is particularly manifested by the two-step procedure of in optimizing Gaussian TSBP (i.e., proving Thm 6): first, reducing the original problem to a static problem, then building the optimal process based on that the optimal solution shares the bridge with the reference process. Importantly, this procedure applies to general non-Gaussian boundary cases, offering a blueprint for solving general cases by first optimizing the joint distributions between boundaries, and then constructing a bridge based on the reference dynamics. This principle underlines IPF based training methods for SB models, corresponding to the alternating forward and backward training. 2. _explaining and inspire methods_. The Gaussian TSB time-marginal in eq. 5 aligns with the _stochastic interpolant (SI)_ framework by Albergo et al. (2023), which serves as an alternative approach to generative modeling. So we can really view this framework as a bridge model generalized from the closed-form Gaussian SB, potentially enabling the development of SI models for topological signals.. 3. _advancing theoretical results for SBPs_. The latest work on the Gaussian SB solution is Bunne et al. 2023 which considered the limited case where the reference process has _scalar_ drift coefficients. Our results in fact showed the closed-form SB solution now can be extended for a broader class of reference processes. For even more general reference processes, this is still an open question. **on insights to adopting TSB for non-Gaussian boundaries**: Thank you for this question. - Our first point above in fact showed how solving TSBP for the non-Gaussian case can be approached. - On the other hand, to further see the role of topology, we want to refer reviewer to [this response](https://openreview.net/forum?id=WzCEiBILHu&noteId=qxx5z0bRXh) where we explain why topological approach is needed. We hope these two points provide further clarity and insights. **on computational complexity:** Two types of computational costs are discussed in this work. 1. _Learning TSB vs. SB Models_: The computational demands for learning the two models are comparable. In general, we do not need to handle matrix exponentials in learning TSB models. Our analysis (please refer to our [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) for the details) compared the computational costs of two models when the graph size is up to 40k. _We included plots for this comparison in [appendix F.2.5](https://openreview.net/pdf?id=WzCEiBILHu#page=42) on p42. Moreover, across all datasets, the _additional memory_ and the _increase in training time_ introduced by TSB remain ```below 4%```, shown in table F.8._ 2. _closed-form Gaussian TSB solution_: For this, we need to handle operators involved with matrix exponentials. In our practice for the large dimension of 20k, we applied fast solver (LOBPCG, available in pytorch, scipy.sparse, and large-scale learning libraries like megaman) to pre-compute the eigenbasis of the Laplacian. Note that this is a __one-time-only__ computation, taking 868.07 secs on an NVIDIA rtx 3080 GPU using ```torch.linalg.eig``` (without using sparse computations). In moderate-sized cases, we can apply scalable Padé approximations (e.g., scipy.linalg.expm ) for this, requiring only 5.7 secs for dimension 5k, or Chebyshev approximation (please see this [anonymized tutorial](https://anonymous.4open.science/r/topological_SB_matching-387D/GTSB_ClosedForm/readme.md)), as suggested by reviewer 1. To improve the presentation, we plan to add a notation table to mitigate the dense symbols. We made updates accordingly in the revision. Please refer to the [revision summary](https://openreview.net/forum?id=WzCEiBILHu&noteId=o7PEfrdNne). Thank you again for your thoughtful feedback. We hope our responses addressed your concerns.

Dear reviewer, Thank you for appreciating our rebuttal and the updated rating. Your suggestions and questions have helped our work. We appreciate them very much, as well as your time and effort in reviewing our work. In case you have any further questions or feedback, please feel free to leave a comment.

Dear reviewer, we appreciate your effort in reviewing our work and thank you for the feedback. Please find our response to your concerns: **On experiments**: We added new applications on human brain networks and single-cell dataset for matching tasks. We hope, together with the previous ones, they are sufficient to show the generalizability of the work. Please see our [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) for the details and results, and [_revision summary_](https://openreview.net/forum?id=WzCEiBILHu&noteId=o7PEfrdNne) for the according updates in the revision. Or, more conveniently, please refer to this [anonymized repo](https://anonymous.4open.science/r/topological_SB_matching-387D/TSBLearning/readme.md) where we put the new experiments together, and tabular results related to the following two points. **On scalability of the methods:** Please allow us to clarify that the ocean current is performed over the space of around ```20k``` edges in a simplicial 2-complex. Similarly, the newly added single-cell data is performed over a graph with around ```18k``` nodes. In both large-scale networks, our method scales well. Our experimental analyses showed, compared to the SB-models, the _additional memory_ and the _increase in training time_ introduced by TSB remain ```below 4%``` for these two large-scale experiments. We also reported the training time and memory with respect to the size of the graph up to 40k in the [general response](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b). In [appendix F.2.5](https://openreview.net/pdf?id=WzCEiBILHu#page=42) on p42, we now included an illustrative comparison plot, as well as the comparisons of two models across all datasets. - Analytically, as the TSB-based learning is built upon the SB-based framework, the additional computational cost is only due to the extra drift term in the SDE and it exhibits the same complexity as the SB-based methods for the rest. Let us consider a graph with $n$ nodes and $m$ (often $m\ll n^2$) edges: if an SB-BM model has a computational cost $O(F(nd))$ in one iteration for some complexity factor $F(nd)$ with feature dimension $d$, then a corresponding TSB-BM model, due to the matrix multiplication $LX$, has a cost of $O(md + F(nd))$ if using the sparse computation; and, a cost of $O(n^2d+F(nd))$, if not. Modern GPUs really allow us to perform such large-scale matrix multiplications very much like allowing us to train deep neural networks. - In future work on model training, we aimed to be more inclusive on related work, where we cited recent advances (Deng 2024) that designed (partially) simulation-free training for SB models with a Gaussian final distribution. We wanted to point out this (among others) potential directions on training TSB models. While the other two references (Peluchetti 2023; Shi et al. 2023) used an alternative training method, they necessarily have comparable computational costs. **On normalizations of the graph Laplacian:** - We apologize for the unclarity here and have updated the normalization details in the appendix F.2.1. Specifically, for the synthetic and seismic experiments, we used the symmetric normalization, as well as for the newly added brain signal matching task. For the traffic flow and ocean current matching, we considered the normalized 1-Hodge Laplacian given by _Schaub et al. 2020_ for the former and a weighted Hodge Laplacian following _Chen et al. 2021b (Algorithm 1)_ for the latter. - We additionally tested the performance with different graph normalizations on the newly added brain fMRI data. Please see the results below. | Method | $W_1, L_{\rm{sym}}^{\rm{normalized}}$ | $W_2, L_{\rm{sym}}^{\rm{normalized}}$ | $W_1, L_{\rm{RW}}$ | $W_2, L_{\rm{RW}}$ | $W_1, L_{\rm{comb}}$ | $W_2, L_{\rm{comb}}$ | |:----:|:----:|:----:|:----:|:----:| :----:|:----:| | TSB-BM | $7.51 \pm 0.08$ | $5.51 \pm 0.06$ | $7.51 \pm 0.08$ | $5.52 \pm 0.06$ |$8.06 \pm 0.05$|$5.80 \pm 0.04$| | TSB-VE | $7.59 \pm 0.05$ | $5.55 \pm 0.04$ | $7.62 \pm 0.09$ | $5.58 \pm 0.07$ | $9.21 \pm 0.06$|$6.62 \pm 0.05$| | TSB-VP | $7.67 \pm 0.11$ | $5.64 \pm 0.09$ | $7.65 \pm 0.09$ | $5.62 \pm 0.06$ | $9.29 \pm 0.05$|$6.73 \pm 0.03$| ** Note that for $L_{\rm{comb}}$, we normalized it by the maximal eigenvalue for stability. We notice that using the symmetric normalized Laplacian has comparable performance with using the random walk one, and using the combinatorial one is worse than the other two. This is now included as an [ablation study on graph normalizations](https://openreview.net/pdf?id=WzCEiBILHu#page=41.69) on p41. We appreciate your feedback again and hope our responses have addressed your questions on experiments, the scalability of the methods and the graph normalizations.

Dear reviewer, thank you for your positive feedback. Please see our response to your questions. **On why a topological approach is needed, and vs. Euclidean one**: We want to explain this from the following perspectives: - When a time stochastic process runs over a graph (for example), in addition to the knowledge (e.g., Markov property by Itô processes) over time dimension, a topological approach also utilizes the spatial statistics of the process, e.g., the covariance dependency between pairs of nodes, based on the topological structure. - Let us consider the processes driven by $\rm{TSHeat-BM}$ [cf. Example 1 in line 197] and $\rm{BM}$ (Wiener process). Conceptually, the latter is a diffusion over time, while the former is a diffusion over both time and graph (space). To be more concise: - In their transition densities from $t_1$ to $t_2$, the former has a cross-covariance $K_{t_1t_2}$ [cf. eq. 2] depending on the coupling of time and graph, while the latter has an only-time-dependent cross-covariance as $t_2-t_1$. The former really gives rise to a nonseparable graph-time process. So if the underlying topology is ignored, the process will not be able to capture the spatial dependency between nodes. - Furthermore, we can see that the topological SB gives rise to an E-OT with a transport cost taking into account the topological information stored in $\Psi_t$ and $K_{t_1t_2}$, while classical SB returns an E-OT with a quadratic transport cost without any topological dependency. - In the Gaussian case [eq. 5], the topological bridge really incorporates the topological information via $K_{t_1t_2}$ and $\Psi_t$. - Experimentally, on the ocean current (Figure 4) data, as well as the newly added brain signals (```Figure F.7``` on [p41](https://openreview.net/pdf?id=WzCEiBILHu#page=41)) and the single-cell one (```Figure. 3``` on [p9](https://openreview.net/pdf?id=WzCEiBILHu#page=9)), we observed that when the topology is ignored, the matching trajectories either get very noisy or remain largely unchanged. This is because the temporal statistics alone are difficult to transport one type of high-dim signals to another. Please refer to our [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) for details and results on the newly added tasks. **On point clouds tasks:** - We apologize for the confusion here. We meant to cite _Deng et al. 2024_ for time series and images. For point clouds, we meant they are often embedded/defined in terms of 2D or 3D coordinates, yet we agree with the reviewer on the challenging and general nature of point clouds tasks. - From a broad sense, our experiments are carried out on point clouds: the seismic magnitude supported on mesh points and the newly added fMRI signals supported on brain regions are scalar fields on point clouds, and the current velocity fields supported on ocean buoys are vector fields on point clouds. However, note that our setting considers the underlying network structure (built upon the point clouds) remains fixed, so the points are not moving and we focus on the signals associated to it. - For the typical point clouds tasks, 1. like object classification, the point clouds belonging to the same class of objects are possibly different, yet recently TDA-based methods have been applied based on the consistent topological features of the point clouds in the same class, which are rather different from our setting. 2. like object tracking, where the point clouds are moving over time, we initially thought our current work does not directly support that scenario because we did not want to have a topology tangled with time dimension. However, as suggested by the first reviewer [here](https://openreview.net/forum?id=WzCEiBILHu&noteId=GIKXfhIesn), if we have observations over the time dimension, we can build a big graph which connects all data points over the time span, and transport observations from one timepoint to another based on their indicator functions. We did add this more point clouds-like application using the single-cell data. Please refer to this [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) for the details. - So overall, our work can be applied to the point clouds tasks, though sometimes not trivial. **On computational aspects:** - Please see our [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) on this. We now included experimental complexity analyses in [appendix F.2.5 on p42](https://openreview.net/pdf?id=WzCEiBILHu#page=42). The main message is that TSB models have comparable complexity cost compared to SB ones. We hope our responses addressed your questions by adding additional experiments on point clouds and experimental analysis on complexity. We once again appreciate these suggestions and made updates in the paper accordingly.

Dear reviewer, thank you for taking the time to review our work. Please find our response to your questions: **On experiments:** - We added an additional dataset on human brain networks. Please refer to our [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) for the details and results. - We added Swiss Roll point clouds to corroborate the theoretical results. Please refer to this [anonymized tutorial](https://anonymous.4open.science/r/topological_SB_matching-387D/GTSB_ClosedForm/readme.md) for the implementation. We initially did not use this because it was commonly used to transport discrete probability measures from one point clouds to another. This is different from our problem setting where we have a fixed topology and can be viewed as a limitation of our framework, as discussed in lines 518-526. For the same reason, our method cannot be directly applied to the tasks in [3,4] where the data points, together with the possible underlying topology, are moving over time. A possible (perhaps debatable) topology-based method might need a time-varying topology structure. While it is a promising future direction, the extension of topology-based SB models to time-varying topology can be nontrivial. - For figures 2 & 3, we forgot to cross-refer to the Wasserstein metrics in appendix Table F.4, which we have fixed. For this experiment, we only have initial and final samples. - We added Wasserstein distances for the seismic and traffic experiments in the appendix Table F.5. For convenience, we show them here: | Method | Seismic $W_1$ | Seismic $W_2$ | Traffic $W_1$ | Traffic $W_2$ | |---------|---------------|---------------|---------------|---------------| | SB-BM | 11.73 $\pm$ 0.05 | 8.29 $\pm$ 0.04 | 18.69 $\pm$ 0.02 | 13.36 $\pm$ 0.01 | | SB-VE | 11.49 $\pm$ 0.04 | 8.13 $\pm$ 0.03 | 19.04 $\pm$ 0.02 | 13.61 $\pm$ 0.02 | | SB-VP | 12.61 $\pm$ 0.06 | 8.92 $\pm$ 0.04 | 18.22 $\pm$ 0.03 | 13.02 $\pm$ 0.02 | | TSB-BM | 9.01 $\pm$ 0.03 | 6.37 $\pm$ 0.03 | 10.57 $\pm$ 0.02 | 7.62 $\pm$ 0.01 | | TSB-VE | 7.69 $\pm$ 0.04 | 5.44 $\pm$ 0.03 | 10.51 $\pm$ 0.02 | 7.58 $\pm$ 0.01 | | TSB-VP | 8.40 $\pm$ 0.04 | 5.95 $\pm$ 0.03 | 9.92 $\pm$ 0.02 | 7.16 $\pm$ 0.01 | **On solution complexity:** - We did use approximation for this. In our practice, we have applied a fast eigensolver _LOBPCG_ for the 20k dimension computation, which is implemented in scipy.sparse, pytorch and other large-scale manifold learning libraries. We did this because this computation is required _only once_, and all other operators involved in Thms 6 & 7 can be computed efficiently when using certain amount of principal eigenbasis. In the synthetic experiments, we used scipy.expm, which is based on a Padé approximation. - We also included now the _Chebyshev polynomial approximation_ based on _Huguet et al. IEEE MLSP 2023_ and _Marcotte et al. ICASSP 2022_. We appreciate your suggestion and reckon this is feasible thanks to the closed-form Chebyshev coefficients for matrix exponentials (eq. 5 in Marcottee et al. 2022), despite the need for all $t\in[0,1]$. Please see the [anonymized tutorial](https://anonymous.4open.science/r/topological_SB_matching-387D/GTSB_ClosedForm/readme.md) for this implementation, where we also provided a running time comparison between the two ways of approximations. - We added experiments on complexity and training time on different sized swiss roll point clouds. Please refer to our [_general response_](https://openreview.net/forum?id=WzCEiBILHu&noteId=7fX8aKiH6b) for the details. **On minor questions:** 1. It is at the numerator. To avoid confusion, we updated it more clearly as $\exp(\frac{-\kappa^2}{2}L)$. 2. Yes, it is a matrix exponential. 3. It is $\xi_{t=1}$ by taking time $t=1$, and $\xi_t$ is defined in eq. (cond. mean) at line 119. 4. Yes, we need the backward one. We updated the text to make it more clear: 'eq. (3a) runs forward and eq.(3b) runs backward in time with the backward Wiener process'. We appreciate your feedback on our work and hope that our responses addressed your concerns. ref: _Marcotte, et al. "Fast multiscale diffusion on graphs." ICASSP 2022._

Thank you, I appreciate your willingness to try these new experiments and provide results so fast. For the interpolation evaluation, when you write “despite that we did not incorporate the intermediate observations in the training”, does the graph still includes all the data ? For example, if you are holding out timepoint `2`, is your graph constructed from timepoints 1,3,4,5 ? I think both can be interesting, but it should be clear in the details.

Yes, we made this clear in the new revision, in both main text (lines 427-428) and appendix on data details. Meanwhile, we added the accuracy when training on the entire dataset (Table F.7). A quick observation is that the final prediction accuracy can even reach 100%. However, the intermediate predications (interpolations) remains poor (still using indicators). Thank you also for the other suggestion. That is indeed a way to go when not using the entire dataset in this case. Yet, we agree that can be nontrivial and goes beyond this work. Overall, we hope we fully addressed your concerns on this work. It was quite pleasant to get your suggestions and work on applications outside our comfort zone. This greatly improved our work.