Paper: Unbalanced Diffusion Schrödinger Bridge (ICLR-2024)

Reviewer comment (Writing):
The paper is mathematically heavy and its clarity is poor, with many important mathematical quantities and their properties assumed to be known by the reader.

Author response:
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.

Reviewer comment (Experiments):
The experiments are not extensive enough; the authors only compared their method with one competitor.

Author response:
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:

|                                                                      |              **MMD**               | $\textbf{W}_\varepsilon$  |
|:--------------------------------------|:---------------------------:|:--------------------------------:|
| De Bortoli et al. [6]                           |       12.8e-2 (0.07e-2)     |                10.53 (0.02)                |
| Ours                                                          |  **1.13e-2** (0.07e-2)  |              **5.52** (0.09)            |

Reviewer comment (Experiments):
Additional Schrödinger Bridge problem approaches for the balanced case ([1, 2, 3]) could have been considered, and if they cannot be applied to this dataset, other single-cell datasets ([4]) could have been explored.

Author response:
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:

|                                                                      |              **MMD**               | $\textbf{W}_\varepsilon$  |
|:--------------------------------------|:---------------------------:|:--------------------------------:|
| De Bortoli et al. [6]                           |       12.8e-2 (0.07e-2)     |                10.53 (0.02)                |
| Ours                                                          |  **1.13e-2** (0.07e-2)  |              **5.52** (0.09)            |

Reviewer comment (Theory):
It is unclear what happens when the measures have the same mass—whether the method recovers the original IPF and Schrödinger Bridge.

Author response:
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.

Reviewer comment (Experiments):
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. I would like to see additional comparisons to cell dynamics methods, particularly those based on branching SDEs.

Author response:
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:

|                                                                      |              **MMD**               | $\textbf{W}_\varepsilon$  |
|:--------------------------------------|:---------------------------:|:--------------------------------:|
| De Bortoli et al. [6]                           |       12.8e-2 (0.07e-2)     |                10.53 (0.02)                |
| Ours                                                          |  **1.13e-2** (0.07e-2)  |              **5.52** (0.09)            |

Reviewer comment (Novelty):
Placement with regard to existing work: I believe the authors missed work on Trajectory Inference [1,2,3]. Comparison to these methods which also account for cell growth and death in an SDE formulation would greatly strengthen this work and provide context on other ways this problem has been addressed.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
I don’t understand why the reconstruction quality at the end timepoint should be any better. This is an IID generative modeling task. Is DSB underfit here?

Author response:
The quality of the final population reconstruction is not _necessarily_ better on every dataset, but should ideally improve if the introduction of the killing mechanism allows for improvement in the trajectories at intermediate times, e.g., by accurately modeling births and deaths or by helping remove outliers. For instance, in our cell dataset, some cells are likely to be killed by the drug (i.e., are underrepresented in the final empirical distribution) while others proliferate (i.e., are overrepresented at $t=1$): the presence of this phenomenon is supported by the fact that, for some drugs, the total population of cells grows over time, signaling that the net balance of births vs. deaths is positive.

We argue that the use of standard DSBs in those cases would not be ideal. That is because DSBs operate under the assumption that the evolution of cell statuses is **continuous in the state space**. As such, to accommodate for the presence of ‘outliers’ they need to create artifacts in the reconstructed dynamics, which are more subtle examples of the diagonal trajectories appearing in the DSB solution to our toy example (Fig. 3c and 3d).

Reviewer comment (Experiments):
I don’t understand how the data is partitioned into training and test for this experiment. If the middle timepoint is left out, how is it known that the observed mass increases 35% but then decreases to -25% from the original total mass? The appendix just states a 80/20 test split.

Author response:
Thanks for the remark, we apologize for not mentioning this clearly and will update the text accordingly. We use the training set to (i) obtain the individual statuses of cells at initial ($t=0$) and final ($t=1$) times and (ii) compute the target mass variation at all time points (including the intermediate one): as mentioned above, we can deduce it from the dataset, by comparing the number of cells analyzed at every time instant.

We then let the test cells evolve and benchmark their final distribution to assess the quality of the learned dynamics.
Note, however, that we do not feed the individual statuses of cells at the intermediate time to our UDSB algorithm, meaning that our method cannot simply be traced back to some kind of multi-marginal SB.

Reviewer comment (Theory):
The requirement that $k$ is non-negative is not a standard requirement in unbalanced dynamic OT. This seems to be necessary in the theory, but is not a requirement in practice. I did not understand why. Could the authors explain this / clarify in the text?

Author response:
The non-negativity of $k$ is linked to the fact that, when presenting the unbalanced SB theory, we restrict ourselves to death-only forward processes, i.e., processes in which the total mass is non-increasing in time. The killing function $k \delta t$ is the (spatial) rate of mass variation and it is therefore natural to have $k \geq 0$ in this case. Analogously, we represent births with _negative_ values.

Reviewer comment (Presentation):
More of a comment than a question, “simulating virus spread” is mentioned once in the intro, in the conclusion with an extended experiment in the supplement. I think either this should appear in the main text or not be included. I don’t really understand why this is not in the experiments.

Author response:
We believe that this additional experiment on virus spread is insightful because it shows that birth/death functions do not always need to represent the physical (dis-)appearance of mass but can also be used to reproduce the (soft) transition of population members across categories (in our case, across virus variants). We consider this kind of UDSB application very promising as it allows us to selectively forgo the continuity of particle trajectories in the state space and work with discrete features (if they can be embedded in Euclidean space).

We, however, eventually chose to not include this experiment in the main text because we were unable to fairly compare it against existing methods. Our reconstruction task, namely the inference of virus propagation given initial and final known states, is non-standard in epidemiology (which mostly deals with predicting the future course of the disease given its past behavior) and our attempts at surveying the (very) big amount of recent research did not reveal a suitable baseline candidate.

We will remove the references to this experiment from the introduction, as they may be confusing for readers, and only mention it in the outlook section.

Reviewer comment (Evaluation):
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?

Author response:
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).

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
The authors do not discuss computational complexity and robustness of the method. How does capacity of NNs influence the results?

Author response:
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.

Reviewer comment (Reproducibility):
The theoretical derivations look OK, but the practical implementation is very difficult to understand, follow and reproduce.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Writing):
The description of the algorithm and of its each step are difficult to follow in the appendix.

Author response:
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.

Reviewer comment (Experiments):
In Sec. 5, page 6, the authors introduce a TD loss. Why is it needed? Can the algorithm work without it? Any ablation study on how this loss influences the final quality of the results?

Author response:
As also discussed in [8], the combination of the MM and TD losses is important to learn the Schrödinger densities. Intuitively, the MM loss drives the learning of the scores ($\nabla\log \varphi$ and $\nabla \log \hat{\varphi}$), while the TD removes unwanted terms which may appear as additive factors to the quantity $\log \varphi$ (and $\log \hat{\varphi}$). This is made more precise in [8], particularly in the proof of Proposition 4 (Section A.3.5).

Without the TD loss, the MM loss would just constrain the **gradient** of  $\varphi$. This would however not be sufficient in our case, since we use $\varphi$ to compute the normalization quantity $\Psi$.

This loss is therefore theoretically _important_ in the sense that it ensures that we solve the true generalization of SBs to unbalanced marginals. However, as noted above, the TD loss has many disadvantages that render it of limited interest in practice. For that reason, we try to emulate its contribution to training by means of the (more tractable) Ferryman loss.

Finally, to test whether the addition of the Ferryman loss truly improves the quality of the reconstructed dynamics, we convert our algorithm into the Deep GSB by [8] (a standard SB solver), which is similar in structure to UDSB, upon disabling the death/birth mechanism of the latter. The results – contained, e.g., in Table 1 – show that the good performance of our algorithm (3rd row) with respect to the baseline (1st row) cannot be simply explained by differences in implementation (e.g., training, optimizers, …), since by executing our same code and only disabling the birth/death mechanism (2nd row) we get worse results.