Thank you for the time invested in reviewing our paper.

The paper is mathematically heavy and I find its clarity poor. A lot of important mathematical quantities are assumed to be known by the reader as well as some of their properties. [...] While I find the theory to be strong and interesting, it might be possible that the paper's scope is too theoretical for a machine-learning venue like ICLR.

We are aware that our paper includes many mathematical details that can confuse – rather than inform – the reader and ultimately conceal the relevance of the content.

We have nonetheless chosen to submit our manuscript to ICLR in order to highlight the applicability of our proposed method in practical applications, such as those in cancer biology. Our submission reflects the bigger trend involving SBs – long studied theoretically but only recently coupled with efficient solvers – which are now becoming a tool to apply in real-world scenarios.

I find the experiments not extensive enough. The authors only compared their method with one competitor. They could have also considered the following Schrödinger Bridge problem approaches (that deal with the balanced case) [4, 5, 6] [...].

Even though we explicitly mention one baseline in our manuscript, we effectively also compare our results against [8], since the latter coincides with UDSB, after disabling its birth/death mechanism.

We report the performance of this method (e.g., in Table 1) with the label Ours, no births/deaths. We explicitly reference Liu et al. [8] only in the caption since we use our own implementation (which coincides with our UDSB code, just without killing function), rather than the one published by the authors. We do this to highlight that the performance improvement cannot be simply explained by differences in implementation (e.g., training, optimizers, …) but must be linked to the killing mechanism.

Moreover, we have compared our algorithms also against De Bortoli et al. [4], but did not report the results because we encountered convergence issues and eventually got worse-quality predictions, compared to [10]. We nonetheless provide our measurements below:

What happens when the measures have the same mass? Do you recover the original IPF and Schrödinger Bridge?

Yes, that is the case, since when the mass at both marginals is equal, $\Psi = 0$ (this can be seen from eq. 33). As a consequence, the jump-related terms in infinitesimal generators disappear and all the processes revert to standard diffusions in $\mathbb{R}^D$. We will more explicitly state that our method is a (strict) generalization of the standard SB problem in the manuscript.

Thank you for your review.

[...] I believe most of the ICLR audiences are not familiar with that.

We agree that we did not sufficiently stress the fundamental differences between flavors of SBs and, instead, directly delved into the details of our proposed method. Below, we include a short survey of SBs for the benefit of the reviewer. We will also take care of clearly defining these terms in the introduction of our manuscript.

I'm curious about the uniqueness of SB/DSB/UDSB methods? Are there any real applications where SB/DSB/UDSB works well and other methods do not work?

Different terms related to Schrödinger Bridges:

• SBs are Stochastic Differential Equations (SDEs) which describe the most likely evolution of a system with known initial and final states. When assuming that the members of a population (in the original example, the particles of a gas) evolve continuously in time, it turns out to be possible to recover the dynamics of the system by solving (coupled) SDEs for which no analytical solution is known in the general case. This means that, although conceptually very powerful, the SB framework by itself does not provide an easy way to compute the evolution of a real-world system for which (i) the state space is not low-dimensional or (ii) the initial/final distributions are only known empirically.

• DSBs are one (very successful) attempt at solving SBs numerically. They heavily rely on the (somehow unintuitive) observation that, as originally proved by the physicist Schrödinger, diffusion processes are fundamentally invertible, i.e., they are still diffusion processes even if the time is played backward. Mathematically, this realization justifies the use of an iterative algorithm, similar to the Sinkhorn algorithm, to compute SB solutions.

• Although computationally sound, DSBs require the (sometimes limiting) hypothesis of mass conservation. In other words, solving DSBs means ‘matching’ two populations (the initial population, time $t=0$, to the final one, time $t=1$) under the assumption that particles change but none of them disappears. This requirement is often violated in practice (since living beings die and reproduce, and natural systems are rarely isolated from their environments) and this is the fundamental motivation of our work. We show that UDSBs (properly) extend DSBs to the case of loss/gain of mass, while still being efficiently solvable via a Sinkhorn-style algorithm, grounded in the theory presented in Sec. 4.

The difference in quality between UDSB and standard SB solutions rests on the nature of the open system under consideration. In particular, we find that if

• the amount of mass exchanged between the system and its environment is non-negligible,
• additional information about constraints or physically unattainable states is available

then UDSB-F learns dynamics that are more meaningful than the vanilla SB ones. In fact, when the interaction system/environment can be encoded using death/birth zones, which do not find satisfying equivalents in standard SBs, our experiments point to improved modeling of the phenomenon itself (as opposed, e.g., to the mere refinement of training procedures).

Thank you for the valuable suggestions provided. Below, we address your comments.

Experimental evidence: Despite presenting a very interesting theory. Results are presented on one small single-cell dataset. It is difficult to tell how and when to apply UDSB given it is only on a single example with what seems like strong prior knowledge on birth and killing rates. Furthermore, there is only a single baseline for the cell dynamics interpolation task when many specialized methods exist for this task. For me this is the single weakest part of the paper. [...]

In our experiments, we use the single-cell dataset in [9] because it provides reliable cell counts, which are usually not available in single-cell datasets. In it, the number of cells recorded at every step is, in fact, chosen as a fixed proportion of the total population at that time.

This information is vital for us to prove that our method adjusts births and deaths so as to match the observed evolution in population size, which is the distinctive trait of our contribution. As correctly observed, it is relatively easy to define birth/death mechanisms in this scenario but, as noted below (see answer on virus spread), we can also apply our method on weaker priors.

Furthermore, we would like to stress that we explicitly mention only one baseline in our manuscript, but we effectively compare our results also against [8], since the latter coincides with UDSB, after disabling its birth/death mechanism. We report the performance of this method (e.g., in Table 1) with the label Ours, no births/deaths. We explicitly reference Liu et al. [8] only in the caption since we use our own implementation (which coincides with our UDSB code, just without killing function), rather than the one published by the authors. We do this to highlight that improvement in performance cannot be simply explained by differences in implementation (e.g., training, optimizers, …) but must be linked to the killing mechanism.

Lastly, we have compared our algorithms also against De Bortoli et al. [4], but did not report the results because we encountered convergence issues and eventually got worse-quality predictions, compared to [10]. We nonetheless provide our measurements below:

Placement with regard to existing work: I believe the authors missed work on Trajectory Inference [1,2,3].

Thank you for pointing this out. We will duly reference these papers in the manuscript. In particular, we will mention that, whenever multiple marginals are available, it is possible to solve a global regularized OT optimization problem which admits (heuristical) extensions to the case of deaths if the growth mechanism is known [2, 1] or a characterizable bias [3] when the latter is unknown. The main difference with respect to our method is that we do not necessarily rely on individual particle states at intermediate times (but only at the initial and final endpoints) and encode instead prior knowledge about the system dynamics as death/birth functions.

UDSB requires known killing and / or birth rates which may limit in practical applications. It seems like these may be difficult to tune in practice, but the freedom to set them is nonetheless quite interesting.

True. UDSBs require some knowledge of the killing mechanism: the more precise its description, the more relevant the use of Unbalanced SBs. However, as we try to show with the aid of the toy dataset, birth (and death) functions do not need to mirror the physical (dis-)appearance of mass. They may be used more broadly as a way to include additional information about the nature of the system at hand. For example, in Fig. 3 we use killing functions to compensate for the missing observations in the marginal (empirical) distributions: doing this effectively helps the SB discart the outliers (e.g., the group of particles that is not observed in one of the marginals) and prevent the deterioration of the quality of SB solutions.

Thanks for your review and for giving us the chance to elaborate on the relevance and limitations of our work.

The experimental results are not very convincing. Even taking std into account still differences between results (for Ours and Ours - no death/births) are not that big. Is this difference really important for a considered downstream task? Ok, MMD is slightly smaller, and so what? Can we really better predict, e.g., responses to cancer drugs in practice?

The difference in the quality of solutions computed by UDSB-F and standard DSB solvers rests on the nature of the open system under consideration.

In particular, we find that, if

• the amount of mass exchanged between the system and its environment is non-negligible and
• additional information about constraints or physically unattainable states is available,

then UDSB-F learns dynamics that are more meaningful than the vanilla SB ones. In fact, when the interaction system/environment can be encoded using death/birth zones, a scenario that has no counterpart in standard SBs, our experiments point to better modeling of the phenomenon itself (as opposed, e.g., to the mere refinement of training procedures).

More importantly, besides being (modestly) visible in better distributional metrics (e.g., in MMD), the improvement consists in building SB dynamics that look more similar to the natural phenomenon itself. This helps reconcile practical applications of SBs with the spirit of Schrödinger’s original idea: that of faithfully modeling the probable evolution of a physical entity (which evolves according to well-defined admissible behaviors).

Moreover, the dimensionality of considered data is rather limited. Although the authors consider an important applied problem, it is not clear whether the proposed method is efficient for more high-dimensional tasks

Even though we did not include examples of Unbalanced SB solutions on high-dimensional data, we would like to stress that UDSB-F scales roughly like DSB [4]. That is because the two algorithms share very similar structures, with the only difference being that our method also relies on the Ferryman loss (which is conceived to scale well).

A similar observation, however, does not apply to UDSB-TD. This theory-compliant version of our algorithm becomes, in fact, unstable in high dimensions because it directly parametrizes the Schrödinger Potential ($\hat\varphi$), which is a (scaled) probability density and therefore decreases exponentially in magnitude upon increasing the dimensionality of the state space.

The authors do not discuss computational complexity and robustness of the method. How does the capacity of NNs influence the results?

Thanks for pointing this out. We should emphasize that, while our theory-compliant algorithm UDSB-TD is significantly slower than balanced SB solvers —because of the TD loss —UDSB-F is not. At its core, UDSB-F consists of gradient updates determined by the MM loss (like in [4]) and by the Ferryman loss, which has a simple form.

Unlike UDSB-TD, it directly parametrizes the scores ($\nabla \varphi$ and $\nabla \hat\varphi$), rather than the Schrödinger Potentials ($\varphi$ and $\hat\varphi$), hence avoiding to differentiate through the network at every score evaluation.

In fact, the improvement of the performances of UDSB-TD was, along with its poor behavior in high dimensions, one of the motivations driving the inclusion of our second algorithm UDSB-F.

I have some expertise in diffusion processes, but the text is difficult to follow for non-experts in diffusion processes. Some important propositions, used by the authors while proving their main results, can be provided for completeness in the appendix

We acknowledge that we rely on a few results from the theory of stochastic processes which are not properly stated in the manuscript. We will properly include them in a new ‘Background’ section of the appendix.

To further improve its readability, we also kindly ask the reviewer to contribute a list of definitions and results that are currently not properly introduced/recalled.

The theoretical derivations look OK, but the practical implementation is very difficult to understand, follow and reproduce. The description of the algorithm and of its each step are difficult to follow in the appendix

We acknowledge that it is non-trivial to turn the pseudo-algorithms that we provide into actual code, also due to the notational burden needed to treat forward/backward SDEs.

To make it easier for readers, we plan to release the entire code repository upon publication. We can already offer an anonymized version to the reviewer, which can be downloaded from here: https://sendanywhe.re/2967NAGS (the link expires on 25th Nov). Our SB solver has been written from scratch, with the aim of following our theory as closely as possible.