We thank the reviewer for these comments.

Comment:

Can we use the probability flow ODE to avoid Brownian motion?

Response:

While probability flow ODEs are indeed powerful tools, they are not ideal for our specific goals for several reasons:

1. Inferring Individual Trajectories:

   • Probability flow ODEs focus on evolving the entire distribution and don't preserve individual particle identities or trajectories
   • Our method explicitly tracks individual particles, which is crucial for understanding the motion of individual particles (e.g., in N-body systems)

2. Statistical Framework:

   • The smooth Schrödinger bridge has an elegant quasi-Bayesian interpretation: the Schrödinger bridge matches the observed marginal distributions while keeping particle trajectories as faithful as possible to the prior $R$
   • Probability flow ODEs lack such statistical guarantees when used for smoothing

3. Other Advantages of our Method:

   • Our method naturally provides velocity and acceleration estimates. These emerge directly from the HMM model we use
   • Post-hoc smoothing via ODEs would require additional assumptions and computation

This makes our approach more suitable for applications requiring detailed trajectory analysis and closer adherence to the prior distribution.

Comment:

Can we tackle the continuous multi-marginal SB directly without tackling the discrete optimal transport?

Response:

The goal of this paper is to infer particle trajectories that are as close to the prior path distribution as possible while matching the observed empirical distributions. Empirical distributions are discrete by nature, so a discrete multi-marginal Schrödinger bridge problem is a natural starting point. However, our algorithm can be extended to continuous marginals (population distribution) by using the following approach:

1. Replace the messages $\beta$ and $\gamma$ with their continuous versions
2. Replace sums with integrals in (8) on the right-hand side of the equation

This continuous extension of our algorithm yields scaling functions $\beta^*_0, ..., \beta^*_K$ that correspond to the multi-marginal Sinkhorn algorithm for continuous distributions. Crucially, this modification preserves the computational efficiency of the discrete version, reducing the exponential complexity of multi-marginal Sinkhorn to polynomial time. Furthermore, similar to how we represent $\delta$'s using orthonormal bases in Section 5, the continuous $\beta$ functions can also be expressed through basis expansion, maintaining both mathematical rigor and computational tractability.

Comment:

I suggest the authors make a diagram of the flow of methods to facilitate the reading.

Response: We appreciate this suggestion and will add a diagram to the revised version of the paper. The diagram will illustrate the flow of methods, including the relationship between the multi-marginal Schrödinger bridge problem, the static problem, and the lifted space. This will help clarify the overall structure and flow of our approach for readers.

We thank the reviewer for these insightful comments.

Comment: Which processes are absolutely continuous with respect to the GAP? Is the minimizer of the SB problem with GAP prior smooth?

Response: When $R$ is a smooth prior, the solution to (1) is also smooth, in the same sense. To be more specific, a sample path of an order $m$ GAP is $m-1$-times differentiable with probabiliy 1. In other words, there is an event $A \subseteq \Omega$ such that each $\omega \in A$ lies in $C^{m-1}$ and $R(\Omega \setminus A) = 0$. Since the solution $P$ to (1) is necessarily absolutely continuous with respect to $R$, we must have that $P(\Omega \setminus A) = 0$ as well, so that $P$ is concentrated on $C^{m-1}$ paths. We will clarify this.

Alternatively, the explicit form of the solution to the smooth SB problem given in Theorem 4.1 shows that the optimal $P$ is a mixture of Gaussian processes, which inherit the smoothness of $R$.

Comment: Recovering the dynamic SB solution from the static one requires an additional flow-matching-style training step, which can be computationally expensive.

Response: We would like to draw a distinction between two senses of "recovering" the dynamic SB, focusing for a moment on the standard SB. The solutions of the static and dyanmic SB problems actually coincide (by convolving the solution to the static problem with a Brownian bridge, see Leonard 2014). This is the argument we use in the proof of Lemma 2.1. In this sense, the dynamic SB is "recovered" directly from the static problem, and in particular, this allows easy sampling from the dynamic SB. However, the reviewer is correct that recovering the dynamic SB as a stochastic process (e.g., recovering the drifts in the corresponding SDE) is not so direct, though it can still be accomplished from the static problem (as in Pooladian & Niles-Weed, 2024).

In our work, we recover a smooth dynamic SB in the first sense (with very reasonable computational complexity): given the solution to the static problem, we "recover" the corresponding dynamic solution in the sense that we can easily sample from it (using standard techniques in GP regression). In particular, we can sample from intermediate, unobserved times. (See our "leave-one-out" experiments in Tables 2 and 3, which shows that this leads to accurate reconstruction on synthetic data.)

It is less clear how to recover it in the second sense, since the smooth dynamic SB is not given by the solution to an SDE. We agree that it would be interesting to pursue this direction in future work.

Comment: What is the practical motivation for using smooth priors?

Response: Our motivation is twofold:

1. In many systems, the only physically relevant models involve particles whose velocities vary continuously, i.e., whose paths are at least $C^1$. This is never the case with a BM prior. As our experiments show (e.g., Figure 4), this is particularly relevant when trajectories of two particles cross and separate: the BM prior easily loses the identity of particles between timesteps when two trajectories intersect, but assuming that the velocities are continuous allows the particles to be identified.

2. Existing work on trajectory inference largely focuses on smooth trajectories (see F&S, WLR as examples), not for physical but for essentially statistical reasons. Like classic cubic splines, at a heuristic level, fitting data with smooth paths helps to ameliorate overfitting and ensure interpretability of the resulting estimates. Making this intuition rigorous in a particular statistical setting is a very interesting open problem.

Comment: The most significant limitation of the proposed framework is that it essentially reduces to a one-dimensional Schrödinger bridge. The theory does not extend to high-dimensional Gaussian processes with dependent coordinates.

Response: The fact that the theory does not extend to processes with dependent coordinates is not quite true. First, in Theorem 2.3, the covariance matrix $\sigma$ need not be diagonal; this allows the coordinates of the white-noise process defining the GAP to be correlated. More generally, if we consider a large covariance matrix $\mathbb{R}^{md \times md}$ that describes the covariance among all entries of hidden variables across dimensions, it is not necessary that this matrix be block diagonal (corresponding to independent coordinates). From an algorithmic perspective, this only changes the value of $\Gamma$ (see equations 19 and 20); the algorithm need not be changed.

However, we agree that from a computational perspective, the use of independent coordinates is beneficial, since the tensor $\Gamma$ factors. This greatly reduces the space complexity of storing $\Gamma$, and simplifies some message passing updates.

We also point the reviewer to our general comments above on dimension dependence.

General remarks

We thank the reviewers for their helpful comments and questions. We are gratified that the reviewers found our method "interesting" (XtBd), "novel," "very useful" (N7xY), and theoretically "sound" (okGt).

We would like to address an important issue that arose in several reviews: dimension dependence.

First, we agree that we should highlight more explicitly that our experiments were conducted on low-dimensional data. We plan to include this fact prominently in our revision.

Second, the reviewers are correct that the dependence on dimension is poor. The running time scales quadratically with $M$, the number of coefficients used in the approximation algorithm (see Theorem 6.1). As mentioned in Section 6, $M$ will typically be exponential in the dimension, limiting our approach to relatively low dimensional problems.

However, we offer several arguments supporting the interest of our method:

1. A low-dimensional algorithm is still valuable in applications. For example, in the fluid dynamics and astronomical object tracking examples mentioned in the introduction, the data lies in dimension 2 or 3. Our method makes an important contribution to these inherently low-dimensional problems. Based on our experiments, smooth SBs are now the leading method for N-body tracking problems in 3 dimensions.

2. Existing approaches often fail, even in low dimension. In the 5-dimensional Dyngen dataset, our algorithm (Figure 5) provides significantly more reasonable results than current leading approaches (F&S, Figure 7; MIO, Figure 8). Another algorithm, DMSB, fails to even converge (Table 2). This indicates the low-dimensional setting is far from solved.

3. A small number of coefficients go a long way. In our experiments on the 5-dimensional Dyngen dataset, we use 1024 coefficients (only 4 per dimension), yet our approach still brings substantial benefits. We have also conducted experiments on 10-dimensional data using 1024 coefficients. Our method outperforms the other methods on this data in a leave-one-out task. We will include this new, higher-dimensional experiment in our revision.

4. It is common to assume that high dimensional data has low intrinsic dimension. Many approaches to trajectory inference begin with non-linear dimension reduction. For example, the pipline developed by Schiebinger et al. (2019) begins by reducing dimension to 30. While this is still larger than we can handle, "working up" from low-dimensional examples can yield biologically relevant insights as computational power improves.

Response to XtBd

Comment: In the end of Sec 2.2 and Footnote 1, it's a bit unclear why "smooth" paths corresponds to stationary processes.

Response The footnote about stationary processes has been removed. To clarify, we've added a new lemma in the appendix proving that any Markov Gaussian process with differentiable paths is essentially trivial:

Let $\omega: t \mapsto \omega(t)$ be a real-valued Gaussian process that is Markovian and a.s. differentiable. For arbitrary $t \ge 0$ such that $\mathrm{Var}(\omega(t)) > 0$, we have $\omega(s) = \mathbb{E}[\omega(s)|\omega(t)]$ for all $s \ge 0$ a.s.

In other words, any such Gaussian process is essentially deterministic, indicating that obtaining smooth paths requires moving beyond Markov processes.

Comment: Relation to [2006.14113]?

Response: We appreciate the reviewer mentioning this important work, which we failed to cite though we cited its companion works (''Learning hidden markov models from aggregate observations'' and ''Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem"), which have the same authors.

In [2006.14113], the authors develop a belief propagation algorithm for multi-marginal optimal transport. The use of belief propagation for multi-marginal OT has a long history (Teh & Welling, 2001), and like [2006.14113], we use this connection for an efficient algorithm. However, we believe we're the first to exploit this connection to develop an algorithm for the Schrodinger Bridge problem with a smooth prior.

Our Algorithm 2 is indeed analogous to Algorithm 2 in [2006.14113], tailored to our setting:

1. The potential functions are:
   • Between $\eta_{k-1}$ and $\eta_k$: Given by $\Phi_k$ (defined in Section 4)
   • Between $\eta_k$ and $\omega_k$: Given by the indicator function $\mathbf{1}_{\{z_k^{(0)} = x_k\}}$
2. Our message passing variables ($\delta^-$, $\delta^+$, $\beta$, $\gamma$) correspond to their factor-to-variable messages $m_{\alpha \rightarrow j}
3. Our operators $\mathcal{L}$, $\mathcal{I}$, $\mathcal{R}$ implement the operations in equations (38a,b)

An important technical difference is that since our $\eta_k$ variables are continuous, the operators $\mathcal{L}$, $\mathcal{I}$, $\mathcal{R}$ cannot be directly implemented. This fact requires us to develop an approximation scheme, which gives rise to our practical implementation detailed in Section 5.

Response to N7xY

We thank the reviewer for these insightful comments.

Comment: What is the assumptions on the marginals in Lemma 2.1 (and 3.1)? In fact I would like the authors to double check if there is any missing assumptions for the marginals across the paper.

Response: Thank you for allowing us to clarify this important point, which we'll make clearer in the revision.

A complete proof of Lemma 2.1 is now included in the appendix. Lemma 2.1 consists of two parts:

1. The first part demonstrates that the multi-marginal Schrödinger bridge problem with finite marginal constraints (1) can be reduced to a finite-dimensional KL divergence minimization problem (2). This reduction requires no assumptions about the marginals.

2. The second part addresses a technical challenge: we want to define a multi-marginal Schrödinger bridge problem with the prior $R$ being a GAP and the marginals being discrete. However, the KL divergence in the optimization problem is ill-defined when $R$ has absolutely continuous marginals and the marginals are discrete. We address this by proving that problem (2) is equivalent to formulation (3) if the marginals have densities. Since (3) is well defined even for discrete marginals, we take (3) as the definition of the multi-marginal Schrodinger Bridge for general marginal distributions. This is in fact a standard approach in statistical analysis of the Schrodinger Bridge problem, see, e.g., Pooladian & Niles-Weed (2024).

The equivalence between (2) and (3) uses two assumptions:

• The marginals $\mu_k$ are absolutely continuous with respect to the Lebesgue measure, with finite entropy $\int \mu_k(x) \log \mu_k(x) dx$
• The joint marginals of the prior $R$ are absolutely continuous

In the remainder of the paper, we adopt the setting described in the beginning of Section 2.2 where the marginals are sampled data (i.e., $\mu_k$ is discrete). In particular, this is true of Lemma 3.1. Note that, before the lemma, we defined $p(x, y)$ to refer to a "mixed" discrete-continuous density, where $x$ coordinates are discrete but $y$ coordinates are continuous. We will emphasize this point in the revision.

Comment: Is the Markovian assumption in the GAP prior essential? How about the Gaussianity?

Response: The GAP itself is not Markov; however, the Markovian structure of the lifted process is crucial for our algorithm because:

1. It enables factorization of the joint distribution into pairwise terms, allowing efficient message passing.

2. It ensures the solution $p(z)$ of (6) is Markovian (by Theorem 4.1), enabling efficient trajectory sampling.

Gaussianity in the GAP prior provides computational efficiency through closed-form message passing updates but isn't essential. The method can be generalized to work with any process that can be lifted to a Markov process on the phase space by replacing the potential functions between $\eta_k$ and $\eta_{k+1}$ with the conditional density $r(\eta_{k+1} | \eta_k)$ of the desired prior process.

Comment: Does the algorithm translate to general missing data problem?

Response: Yes. As the reviewer mentioned, not observing higher-order derivatives is a form of missing data handled by our algorithm. Other forms of missing data can be incorporated by modifying the factor nodes $\alpha_k$ in the graphical model (Figure 3). Currently, these nodes are indicators enforcing that only positions $\omega_k$ are observed. By modifying these factor nodes, one can incorporate other types of missingness in the observations. We plan to pursue this in future work.

Comment: Is it true that an inferred trajectory necessarily pass one particle at each time point? And does is makes sense in the context of single cell example? Is there an assumption that each time snapshots has same number of particles?

Response: Yes, in the current formulation, inferred trajectories must pass through exactly one particle at each observed time point, due to the marginal constraints.

For point cloud matching applications, while tracking individual "particles" may not be physically meaningful, our algorithm remains valuable for inferring collective motion patterns:

1. We can infer velocity and acceleration fields that characterize point cloud dynamics
2. We can interpolate the point cloud distribution at unobserved times by sampling from the solution

Our results in Table 2 show that this approach yields good results.

Relaxing the requirement that trajectories pass through observed points can be achieved by modifying the factor nodes $\alpha_k$ to model observational noise. This allows inferred trajectories to deviate from exact observed positions. We've added this extension as a remark in the revision, and plan to pursue it in follow-up work.

Our method does not require the same number of particles across different time snapshots. Both the theoretical framework and algorithm support arbitrary discrete marginal distributions.