Original: Shooting Schr¨odinger’s Cat
Revised: Shooting Schr¨odinger’s Cat

Original: Anonymous Authors

Original: Anonymous Institution

Original: Abstract
Revised: Abstract

Original: We present a variational inference scheme to learn a model that solves the Schr¨odinger
Revised: We present a variational inference scheme to learn a model that solves the Schr¨odinger

Original: Bridge Problem (SBP).
Revised: Bridge Problem (SBP).

Original: In contrast to previous work, our approach is solver-agnostic and guarantees solutions that respect the prior beyond the first fitting iteration.
Revised: In contrast to previous work, our approach is solver-agnostic and guarantees solutions that respect the prior beyond the first fitting iteration.

Original: Having this solution allows us to generate new samples from one of the distributions by first sampling from the other one and then solving the dynamical system.
Revised: Having this solution allows us to generate new samples from one of the distributions by first sampling from the other one and then solving the dynamical system.

Original: We show that our model is able to learn the transformation between the Gaussian distribution and arbitrary data, as well as learning dynamics that follow a potential function.
Revised: We show that our model is able to learn the transformation between the Gaussian distribution and arbitrary data, as well as learning dynamics that follow a potential function

Original: 1. Introduction and motivation
Revised: 1. Introduction and motivation

Original: The Schr¨odinger bridge problem (SBP) (Schr¨odinger, 1931, 1933) seeks to find a transformation between two probability distributions.
Revised: The Schr¨odinger bridge problem (SBP) (Schr¨odinger, 1931, 1933) seeks to find a transformation between two probability distributions.

Original: It has both a static version, where a direct transformation is obtained, and a dynamic version, where the flows between the two distributions are learned.
Revised: It has both a static version, where a direct transformation is obtained, and a dynamic version, where the flows between the two distributions are learned.

Original: Recent works (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2021) have tackled this dynamic scenario.
Revised: Recent works (De Bortoli et al., 2021; Vargas et al., 2021; Chen et al., 2021) have tackled this dynamic scenario.

Original: Finding these flows can be done using the Iterative Proportional Fitting Procedure (IPFP) (Fortet, 1940; Kullback, 1968; Cramer, 2000), that iteratively optimizes each process so that the Kullback-Leibler divergence ( D KL ) between the two processes is minimized at each iteration.
Revised: Finding these flows can be done using the Iterative Proportional Fitting Procedure (IPFP) (Fortet, 1940; Kullback, 1968; Cramer, 2000), that iteratively optimizes each process so that the Kullback-Leibler divergence ( D KL ) between the two processes is minimized at each iteration.

Original: One way to do this is to update the joint density of the problem using the potentials of the dual representation of the SBP.
Revised: One way to do this is to update the joint density of the problem using the potentials of the dual representation of the SBP.

Original: Approximating these coupled potentials requires estimating an integral involving them, which in higher dimensional problems can be difficult.
Revised: Approximating these coupled potentials requires estimating an integral involving them, which in higher dimensional problems can be difficult.

Original: The thorough approach of De Bortoli et al.
Revised: The thorough approach of De Bortoli et al.

Original: (2021) uses another representation that can more easily be applied where samples of both boundary distributions are available.
Revised: (2021) uses another representation that can more easily be applied where samples of both boundary distributions are available.

Original: This tackles the Schr¨odinger bridge as an iterative mean-matching procedure, where at each iteration, the mean of the reversed forwards process is matched with the mean of the backward process.
Revised: This tackles the Schr¨odinger bridge as an iterative mean-matching procedure, where at each iteration, the mean of the reversed forwards process is matched with the mean of the backward process.

Original: This requires assuming that at each time-step, the change in the state is small, so that the Euler-Maruyama discretization of the process is valid.
Revised: This requires assuming that at each time-step, the change in the state is small, so that the Euler-Maruyama discretization of the process is valid.

Original: Using a small enough time-step ensures this.
Revised: Using a small enough time-step ensures this.

Original: However, if the time-step is not small enough, it can lead to stability issues.
Revised: However, if the time-step is not small enough, it can lead to stability issues.

Original: Additionally, while this model is regularized using a prior diffusion process, this regularization is only directly applied in the first iteration.
Revised: Additionally, while this model is regularized using a prior diffusion process, this regularization is only directly applied in the first iteration.

Original: This can lead to the learned dynamics differing wildly from the prior, especially when they are represented by unbounded functions (such as neural networks).
Revised: This can lead to the learned dynamics differing wildly from the prior, especially when they are represented by unbounded functions (such as neural networks).

Original: Our model seeks to address these issues by learning the dynamics of the two processes in a way that is independent of the solver used, and so better solvers become available, while also maintaining the regularization in every iteration.
Revised: Our model seeks to address these issues by learning the dynamics of the two processes in a way that is independent of the solver used, and so better solvers become available, while also maintaining the regularization in every iteration.

Original: Authors y 5 yy 3 y 2 y 1 ysssssxxxxxpp T

Original: 2. Solving the SBP with shooting variables
Revised: 2. Solving the SBP with shooting variables

Revised: [Figure]

Revised: sssssxxxxxpp T

Original: Writing the two boundary distributions as p 0 and p T , we define the process that transforms the first into the second as the forward process
Revised: Writing the two boundary distributions as p 0 and p T , we define the process that transforms the first into the second as the forward process

Original: [Equation]
Revised: [Equation]

Original: where k is the IPFP iteration,
Revised: where k is the IPFP iteration,

Original: f k : R d → R d is the drift function, B t is the Brownian process and γ t is the diffusion scale.
Revised: f k : R d → R d is the drift function, B t is the Brownian process and γ t is the diffusion scale.

Original: The backward process then transforms p T into p 0 , with drift function g k : R d → R d and same diffusion scales for the same values of t .
Revised: The backward process then transforms p T into p 0 , with drift function g k : R d → R d and same diffusion scales for the same values of t .

Original: We consider the time interval [0 , T ] divided into N sub-intervals ∆ t i , for i ∈ { 0 , . . .
Revised: We consider the time interval [0 , T ] divided into N sub-intervals ∆ t i , for i ∈ { 0 , . . .

Original: , N − 1 } .
Revised: , N − 1 } .

Original: For the following derivation, we assume that we have samples { y i } Ni =0 , obtained by solving the backward SDE with an arbitrary solver and then reversing them time-wise (the derivation is equivalent if we have samples from the forward SDE).
Revised: For the following derivation, we assume that we have samples { y i } Ni =0 , obtained by solving the backward SDE with an arbitrary solver and then reversing them time-wise (the derivation is equivalent if we have samples from the forward SDE).

Original: The probability density of these reversed solutions is then
Revised: The probability density of these reversed solutions is then

Original: [Equation]
Revised: [Equation]

Original: Taking inspiration from Hegde et al.
Revised: Taking inspiration from Hegde et al.

Original: (2021), we now introduce the shooting variables { s i } N − 1 i = and { x i } Ni =1 with transition probabilities p ( x i| s i − 1 ).
Revised: (2021), we now introduce the shooting variables { s i } N − 1 i =0 and { x i } Ni =1 with transition probabilities p ( x i

Original: The probability density in Equation (2)then becomes
Revised: | s i − 1 ) such that the probability density in Equation (2) becomes

Original: [Equation]
Revised: [Equation]

Original: where we have assumed that p ( y i | x i , y i − 1 )

Original: ≈ p

Original: ( y i | x i ).

Original: We now introduce the variational distribution q
Revised: We now introduce the variational distribution q

Original: ( s i − 1 |
Revised: ( s i − 1 |

Original: y i − 1 ) and the “variational” transition density q ( x i | s i − 1 ), which we define to be the forward process, to obtain the following lower bound on Equation (3)
Revised: y i − 1 ) and the “variational” transition density q ( x i | s i − 1 ), derived from the forward process, to obtain a lower bound on the joint density, L ≤ log p ( { y i } Ni =0 ), as

Original: [Equation]
Revised: [Equation]

Original: The diagram of the model is shown in Figure 1.
Revised: The diagram of the model is shown in Figure 1.

Original: The first and last terms  can be easily computed once the distributions p ( y i | s i ) and p ( s i − 1 | y i − 1 ) are defined.
Revised: The first and last terms inside the brackets can be readily computed once the distributions p ( y i | x i ) and p ( s i − 1 | y i − 1 ) are defined.

Revised: However, the middle KL term, D KL [ q ( x i | s i − 1 )

Original: The second term, though, requires us to know the transition probabilities of the forward process, which for more elaborate solvers might be very difficult (or even impossible) to obtain in closed form.
Revised: ∥ p ( x i | s i − 1 )], requires knowledge of the transition probabilities from the forward process, that might be very difficult (or even impossible) to obtain in closed form for more elaborate solvers.

Original: To circumvent this, we replace the KullbackLeibler divergence in this second term by the kernelized Stein discrepancy (KSD, Liu et al.
Revised: To circumvent this, we replace the KL divergence in this second term by the kernelized Stein discrepancy (KSD, Liu et al.

Original: Chwialkowski et al.
Revised: Chwialkowski et al.

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: where δ q,p ( x ) = ∇ log q ( x ) −∇ log p ( x ), (with q and p arbitrary densities) and k ( x, x ′ ) is the (twice differentiable) strictly positive kernel.
Revised: where δ q,p ( x ) = ∇ log q ( x ) −∇ log p ( x ), (with q and p arbitrary densities) and k ( x, x ′ ) is the (twice differentiable) strictly positive kernel.

Original: It can be shown that there is an estimator for the KSD that only requires the score function of one of the densities, which we can calculate for the prior transition density of the shooting variables, since we define it as the transition density of the reference process used in IPFP.
Revised: It can be shown that there is an estimator for the KSD that only requires the score function of one of the densities, which we can calculate for the prior transition density of the shooting variables, since we define it as the transition density of the reference process used in IPFP.

Original: 3. Experiments
Revised: 3. Experiments

Original: In all of these experiments, the SDEs in our model were solved using the weak solver proposed in R¨oßler (2009).
Revised: In all of these experiments, the SDEs in our model were solved using the weak solver proposed in R¨oßler (2009).

Original: Brownian prior
Revised: Brownian prior

Original: Firstly, we test whether our model is able to transform a Gaussian distribution into several synthetic data sets and where the prior process is simple Brownian motion.
Revised: Firstly, we test whether our model is able to transform a Gaussian distribution into several synthetic data sets and where the prior process is simple Brownian motion.

Original: The results are presented in Figure 2.
Revised: The results are presented in Figure 2.

Original: Whirlpool prior
Revised: Whirlpool prior

Original: We then verify that the model is able to learned dynamics that are different from a Brownian motion.
Revised: We then verify that the model is able to learned dynamics that are different from a Brownian motion.

Original: The first test consists of a “whirlpool” prior on the
Revised: The first test consists of a “whirlpool” prior on the

Original: [Figure]
Revised: [Figure]

Original: dynamics: the drift at any particular point is perpendicular to it.
Revised: dynamics: the drift at any particular point is perpendicular to it.

Original: So, for a 2D vector { x, y } the drift would be f ( x, y )
Revised: So, for a 2D vector { x, y } the drift would be f ( x, y )

Original: = {− 72 x, 72 y } .
Revised: = {− 72 x, 72 y } .

Original: The results are shown in Figure 3 (a).
Revised: The results are shown in Figure 3 (a).

Original: Double well prior
Revised: Double well prior

Original: We also test our model with a double well potential prior, where the drift follows the negative gradient of the potential
Revised: We also test our model with a double well potential prior, where the drift follows the negative gradient of the potential

Original: [Equation]
Revised: [Equation]

Original: We compare our model with De Bortoli et al.
Revised: We compare our model with De Bortoli et al.

Original: (2021) (DSB) in Figure 3(b).
Revised: (2021) (DSB) in Figure 3(b).

Original: While our model is able to avoid the hill, DSB does not.
Revised: While our model is able to avoid the potential hill between the wells, DSB does not.

Original: We believe this is because in the first iteration DSB prefers to match the prior paths (see top-left plot of Figure 3(b)) at any cost, and the quickest way to do so is to go over the hill.
Revised: We believe this is because in the first iteration DSB prefers to match the prior paths (see top-left plot of Figure 3(b)) at any cost, and the quickest way to do so is to go over the hill.

Original: Since in the following IPFP iterations, DSB only tries to match the learned paths of the other direction, it will never avoid the hill.
Revised: Since in the following IPFP iterations, DSB only tries to match the learned paths of the other direction, it will never avoid it.

Original: 4. Conclusion
Revised: 4. Conclusion

Original: We have presented a model that is able to obtain solutions to the dynamic Schr¨odinger
Revised: We have presented a model that is able to obtain solutions to the dynamic Schr¨odinger

Original: Bridge Problem that is both solver-agnostic and maintains the regularization throughout the training iterations, not just the first one.
Revised: Bridge Problem that is both solver-agnostic and maintains the regularization throughout the training iterations, not just the first one.

Original: We have shown that it is able to learn transformations from a Gaussian distribution to arbitrary data distributions, both using a simple Brownian prior and a elaborate priors.
Revised: We have shown that it is able to learn transformations from a Gaussian distribution to arbitrary data distributions, both using a simple Brownian prior and a elaborate priors.

Original: Finally, we have shown that our model can also be applied to potential priors.
Revised: Finally, we have shown that our model can also be applied to potential priors.