Stability of Schrödinger Potentials and Convergence of Sinkhorn's Algorithm
Marcel Nutz, Johannes Wiesel
Introduction
where denotes relative entropy with respect to the product of the marginals,
where is the set of measures on with second marginal (and arbitrary first marginal), and is defined analogously. The algorithm alternatingly “fits” the marginals and , hence is also called iterative proportional fitting procedure (IPFP). The argmin in (1.2) can be solved explicitly, and then each step of the algorithm only requires an explicit integration against , cf. Section 3. The algorithm dates back as far as and its convergence properties are well studied when the cost is bounded: the convergence of to the solution of (1.1) holds in total variation (and in relative entropy), see , among others. More precisely, relaxed the boundedness condition, but introduced several other conditions, including one (see (B1) in ) that essentially forces to be bounded from above in one variable and thus excludes quadratic cost with unbounded marginal supports.
One main result of this paper is the convergence of Sinkhorn’s algorithm (1.2) for quadratic cost and arbitrary subgaussian marginals. More generally, our result (see Corollary 3.2) holds for the continuous cost and marginals as soon as
A fairly elementary proof of the convergence , following the same idea of weak-star compactness as , was recently given in under the condition that (1.3) holds for some . We emphasize that this condition does not capture the regime of principal interest: when approximating the 2-Wasserstein distance and marginals are standard Gaussians, say, this condition forces , but the approximation often requires to be several orders of magnitude smaller (e.g., ). While the case is broadly similar to the case of bounded cost, it seems that the regime requires a fundamentally different line of attack—which brings us to the main theorem of this paper.
In all that follows, we focus on in (1.1) for notational simplicity; the general case is easily retrieved by replacing with . Assuming that
is finite, there is a unique solution , and is uniquely characterized within by having a density of the form
Our aim is to establish stability of the potentials with respect to the marginals. Consider sequences and of marginals converging weakly (i.e., in the topology induced by bounded continuous functions). Denoting by associated potentials, we want to state that in a suitable sense. In view of the above characterization, a key step in this endeavor is to establish a form of compactness. In general, it is not straightforward how to formalize the convergence of potentials. E.g., in a computational context, we may be interested in discrete measures approximating a continuous pair . Then, these measures are mutually singular and the spaces and are not immediately comparable. If (and similarly for ), this issue is milder as convergence in -probability yields a natural topology. However, compactness still turns out to be an issue in the regime of interest. If has high integrability, the weak-star topology in can be used, similarly to the arguments for Sinkhorn convergence in . In some cases one can even use the Arzelà–Ascoli theorem (see Appendix B). But in the regime of interest here, where we want to cover quadratic cost and subgaussian marginals, we have not succeeded with off-the-shelf compactness concepts.
Instead, we shall build compactness through an approximation scheme and properties specific to the problem at hand, eventually using compactness of bounded sets in Euclidean space. This construction is the main technical contribution of the paper. It will also allow us to cover the case of mutually singular measures (in fact, focusing on equivalent measures would not result in a substantial simplification). The approximation scheme has the form
The scheme in the diagram also acts as a way to formalize a strong convergence . It implies convergence in distribution; that is, weakly, where denotes the pushforward under . Convergence in distribution is a natural notion given the weak convergence setting, but it is far from strong enough to imply the desired conclusions. If , we show that our scheme implies the convergence in -probability (and similarly for ) under a fairly general condition on the Radon–Nikodym derivatives; cf. Corollary 2.4. This condition is satisfied in particular whenever the marginals convergence in total variation, allowing us to deduce via Scheffé’s lemma a result of its own interest (Corollary 2.6): the optimal couplings are stable in total variation; i.e., for marginals with and , the corresponding optimizers satisfy . Returning to the convergence of Sinkhorn’s algorithm, we interpret each iteration of the algorithm as solving an entropic optimal transport problem with changing marginals . These marginals converge in total variation to , and we infer the convergence of the algorithm to the desired limit .
Several recent works have addressed the stability of entropic optimal transport from different angles. The first result is in , for a setting with bounded cost and marginals equivalent to a common reference measure with densities uniformly bounded above and below. The authors show by a differential approach that the potentials are continuous in relative to the marginal densities. Still with bounded cost (and some other conditions), establishes uniform continuity of the potentials relative to the marginals in Wasserstein distance ; this result is based on the Hilbert–Birkhoff projective metric. Closer to the present unbounded setting, obtains stability of the optimal couplings in weak convergence for general continuous costs. Based on the geometric approach first proposed in , the main restriction of the technique is that the underlying spaces need to satisfy Lebesgue’s theorem on differentiation of measures which generally holds only in finite-dimensional spaces. As a by-product, the main result of the present paper yields a similar stability result for weak convergence; cf. Theorem 2.1 (i). The present result also applies in an infinite-dimensional context; the more important difference, however, is that we achieve a strong form of convergence, whereas is silent about any convergence of the densities or potentials. In particular, we can infer stability in the sense of total variation convergence (Corollary 2.6) and the corresponding convergence of Sinkhorn’s algorithm (Corollary 3.2). It is worth noting that these two stabilities are at opposites ends of the spectrum: in the weak topology, compactness for sets of couplings is immediate; the difficulty is to ensure that a limit is the optimal coupling for its marginals. For a limit in total variation, the latter is easy, but obtaining compactness is difficult due to the strength of the topology. In the present work, we effectively reduce the dimension by focusing on densities with a decomposition given by potentials and then obtain compactness through the potentials. A last related work is , which was conducted concurrently. Here stability of the coupling in Wasserstein distance is shown under certain growth and integrability conditions. Obtained by control-theoretic arguments through a transport inequality, the main strength of this result lies in being quantitative (which the present one is not). On the other hand, is once again silent about the densities or potentials, and does not yield a convergence in total variation. Indeed, we are not aware of previous stability results in total variation beyond bounded settings. Finally, we would like to mention the ongoing research kindly pointed out to us by Giovanni Conforti. In the setting of dynamic Schrödinger bridges satisfying a logarithmic Sobolev inequality for the underlying dynamics and marginal distributions with finite Fisher information, the authors study quantitative bounds for the relative entropy of Schrödinger bridges with different marginals and the convergence of the gradients of the potentials towards the Brenier map as .
The remainder of this paper is organized as follows. Section 2 contains the main results on stability. Section 3 details the application to Sinkhorn’s algorithm. The proof of the main result, Theorem 2.1, is split into ten steps which are reported in Section 4. For convenience, Appendix A summarizes background on entropic optimal transport. Appendix B details how stability, even uniformly on compacts, can be obtained rather directly under strong integrability conditions. Lastly, Appendix C contains some proofs that we defer in the body of the text.
Stability
Let be Polish spaces endowed with their Borel -fields and their sets of Borel probability measures. We recall that is continuous. In Theorem 2.1 below, we consider the entropic optimal transport problem (1.4) for marginals converging to marginals . The condition (2.1) of the theorem implies that and that there exist optimal couplings with associated potentials ; see Section A for these facts and further background.
Before stating the theorem, let us comment on the normalization chosen therein. As mentioned in the Introduction, and are only unique up to an additive constant. This does not affect the sum determining the density of , but in order to obtain a separate convergence for and , it is clearly necessarily to impose an additional condition to pin down this constant. There are many possible choices; in Theorem 2.1, we work with . As is strictly increasing, fixing the value of the integral is equivalent to determining the additive constant (and conversely, there is a version of the potentials such that, e.g., ). Since is bounded, it is clear that always converges after passing to subsequence. Furthermore, this type of normalization is compatible with convergence in distribution.
and that converge weakly to . Then
and the optimal couplings converge weakly: .
and analogous properties hold for .
Then the optimal values converge: . If are as in (ii), then
The next two results discuss the main condition (2.1) of the theorem.
and thus (2.1) is equivalent to boundedness in :
Remark 2.2 follows from the duality ; cf. Proposition A.2. We can provide a sufficient condition for (2.1) in terms of the given data as follows.
Let satisfy (2.9). The following condition is sufficient for (2.1) and (2.10):
The following condition is sufficient for (2.11):
The proof is deferred to Appendix C. We remark that the assumption (2.9) is mostly a matter of normalization. Indeed, suppose that and —which necessarily holds under (2.11), cf. Proposition A.1. Then by duality (cf. Proposition A.2). Therefore, (2.9) always holds after choosing a suitable normalization for , for instance the centering .
Let (2.1) hold and let converge weakly to . Suppose that , and
Then in -probability, where are arbitrary potentials for . If are as in Theorem 2.1 (ii), then in -probability and in -probability.
As is uniquely determined and convergence in probability is metrizable, it suffices to show that any subsequence of has a subsequence converging to . Thus we may further assume that are as in Theorem 2.1 (ii) and show the convergence of and . Fix and a subsequence of . By (2.6) and (2.5) in Theorem 2.1, there exist a further subsequence (not relabeled) and functions such that along this subsequence,
for some with , by (2.4). Taking , then and finally , we obtain
Condition (2.13) holds in particular for sequences converging in total variation.
Suppose that and in total variation. Then
The proof is deferred to Appendix C. Our final result is the stability of the optimal couplings in the topology of total variation, complementing the weak stability shown in Theorem 2.1 (i).
Let (2.1) hold and let converge in total variation to where and . Then in total variation.
For the sake of readability, we state here the proof under the additional assumption that and ; the general case is deferred to Appendix C. By Corollary 2.4 and Lemma 2.5, we have in -probability and in -probability after passing to a subsequence. Under the additional assumption, in and in . We see that
in -probability. But then the convergence also holds in , by Scheffé’s lemma, and we conclude that in total variation. The convergence of the original sequence follows. ∎
Convergence of Sinkhorn’s Algorithm
Fix marginals and a continuous cost . Sinkhorn’s algorithm (1.2) can be written in terms of potentials. Set and
where . One can check by direct calculation that are the same measures as in (1.2). Denoting by the marginal distributions of , the following summarizes well known properties of Sinkhorn’s algorithm (e.g., [28, Section 6]).
Let . We have and for all . Moreover, ; in particular, and in total variation. For , the marginals satisfy
In brief, are potentials for the marginals which in turn converge to in total variation. The stability result of Corollary 2.6 then yields the following convergence result. As emphasized in the Introduction, it covers quadratic costs with arbitrary subgaussian marginals and the problem (1.1) with arbitrary regularization parameter .
Then and the Sinkhorn iterates converge to in total variation.
As by Lemma 3.1, Lemma 2.3 yields that
In particular, and as defined in Lemma 3.1 are potentials for the marginals ; cf. Propositions A.1 and A.2. Next, we show that satisfy (2.1). In general, the Sinkhorn iterates satisfy as well as and ; see [28, Lemma 6.4 and its footnote]. Here, as , we have and .
Consider , then . Using (3.1), Jensen’s inequality and ,
and hence which is bounded by (3.4). Similarly, (3.1) implies and hence
The argument for is symmetric, so that (2.1) holds. As the marginals are equivalent and converge in total variation by Lemma 3.1, the claim follows by Corollary 2.6. ∎
Proof of Theorem 2.1
The proof is structured into several steps.
Further properties related to and .
Proof that .
Proof that induce a coupling and .
Identification of the limit, end of proof of Theorem 2.1 (i),(ii).
Step 1 is based on the following generalization of [29, Lemma 2.3] extending that result from a single measure to a tight set of measures.
The proof of Lemma 4.1 is an adaptation of the arguments in ; for completeness, the details are reported in Appendix C.
As a preparation for Step 4, we record the following covering lemma.
has diameter at most and boundary ,
The general relation implies that , and the other requirements are satisfied by construction. ∎
We record two more facts about the construction in Step 4 for later use. For , (4.2) and (4.5) yield
As and on , it follows that
In view of , this yields in particular
Second, we define similarly as in (4.6) the function
for . In view of , it follows that
As and uniformly (cf. Step 4), we conclude
On the other hand, using the mapping theorem and (4.9) for both and ,
Fix such that . Then by (4.9),
Above, we have introduced the functions on . Analogously, one constructs on . We can now detail the main step of the proof, showing that are indeed potentials for a coupling . To keep track of the argument more easily, we state the technical parts as lemmas and prove them at the end.
and let be measurable with ; we show . Define the auxiliary measures
Fix and consider the decomposition
We estimate separately the four terms on the right-hand side.
The lemma is proved at the end of Step 8. For the first term in (8), Lemma 4.3 shows that there exists such that
We continue with the last term of (8). Since is increasing and nonnegative, an application of the monotone convergence theorem for shows that after increasing as necessary, we have
(The second case will be eliminated by contradiction later on.) The value of is now fixed for the remainder of the proof.
Turning to the third term in (8), note that since and ,
For the remainder of the proof, is fixed.
The lemmas are proved at the end of Step 8. Together, they show
Combining this with (4.12) and (4.15) yields
This shows . Thus we have proved that and weakly [5, Theorem 2.1]. As the marginals then also converge weakly and , this implies .
where and are finite sets. Thus
and similarly for instead of . By the construction in Step 4,
As and are -continuity sets, we also have
We can now expand the difference to be estimated as
In view of (4.19) and (4.20), taking yields
defines a coupling of . Moreover, Step 7 and (2.1) imply that and . By the general verification result in Proposition A.2, the form of with implies that , that and that is the unique minimizer for the entropic optimal transport problem (1.4). It follows that also holds along the original sequence, completing the proof of Theorem 2.1 (i).
It remains to prove Theorem 2.1 (iii). Let (2.7) hold. Passing to a subsequence, we may assume that and are as in Theorem 2.1 (ii). We first show the upper semicontinuity
Indeed, the weak convergence (2.3) and the uniform integrability (2.7) imply that . Together with Portmanteau’s theorem for , the first part of (4.22) follows, and similarly for the second.
Next, we argue the lower semicontinuity of the sum,
By (2.1) and Proposition A.2, we have the duality
Together, the lower semicontinuity (4.23) of the sum and the separate upper semicontinuity (4.22) imply (2.8). This completes the proof of Theorem 2.1. ∎
Appendix A Background on Entropic Optimal Transport
Let and let be measurable. We have the following result on existence and uniqueness for the entropic optimal transport problem (1.4).
If , then and .
See [28, Theorem 4.2] for a proof. Conversely, the next result shows that the form of the density characterizes the minimizer. We also include the duality relation.
Let admit a density of the form
If , then is the minimizer and are its potentials.
If , then necessarily and
In particular, and (a) applies.
See [28, Theorem 4.2, Theorem 4.7, Remark 4.8]. When are potentials as in Proposition A.1, the fact that implies the so-called Schrödinger equations
We may choose versions of the potentials such that these relations hold without exceptional sets.
Appendix B Uniform Stability under Strong Integrability
The following result shows that stability of the potentials, even uniformly on compacts, can be obtained quite easily when the cost is sufficiently integrable. As discussed in the Introduction, the integrability condition is not satisfied in the regime of principal interest. For the statement, we choose versions of the potentials such that (A.2) holds without exceptional sets.
Then uniformly on compacts, where are arbitrary potentials for , and the corresponding optimal couplings converge weakly. If are as in Theorem 2.1 (ii), then also and , uniformly on compacts.
Similarly as in Lemma 2.3, a sufficient condition for (B.1) is that
and are normalized such that (2.9) holds, for instance .
This yields the uniform lower bound , and similarly . Fix . Using (A.2) and the lower bound,
Next, we show that is equicontinuous. On the strength of the pointwise boundedness, it suffices to show that is equicontinuous. Fix a compatible metric on , let and . Using tightness, choose a compact with for all . Let satisfy . By Hölder’s inequality,
Recalling that , the integral can be estimated by
As is continuous and is compact, is continuous. Therefore, for , for sufficiently small, and we obtain the desired equicontinuity,
We have shown that is equicontinuous and pointwise bounded, and the same arguments hold for . Passing to a subsequence, the Arzelà–Ascoli theorem shows that and uniformly on compacts. Note that is -uniformly integrable by (B.1). As uniformly on compacts and weakly, it follows for the optimal couplings that
In particular, . Proposition A.2 now shows that is the optimal coupling for and are corresponding potentials. As is unique (Proposition A.1), the claim for the original sequence follows. ∎
Following the above proof, the argument for pointwise boundedness still applies, and the argument for equicontinuity is even simpler under the additional hypothesis. The argument using uniform integrability may no longer be clear, but we can instead use Theorem 2.1 to conclude that must be potentials for .
Appendix C Omitted Proofs
and now Tonelli’s theorem yields The analogue holds for . Thus (2.11) implies (2.1), and via (2.9) also (2.10).
More generally, given a measurable set , (C.1) also implies
If , then , so that the – characterization of uniform integrability yields the claim.
(ii) The variational representation of relative entropy [28, Lemma 1.3] shows that
for any measurable function bounded from below. Choosing , we deduce
Noting that , the right-hand side is bounded under (2.12), so that (2.11) applies. To obtain the last claim, we replace with in the preceding argument and apply the la Vallée–Poussin theorem. ∎
then shows that . The claim follows. ∎
By Corollary 2.4 and Lemma 2.5, we have in -probability and in -probability after passing to a subsequence. Consider the Lebesgue decomposition into and . Then in total variation and hence in . This implies the convergence in -probability of the reciprocal, and as -a.s., that in -probability. Similarly, in -probability. Following the proof of the particular case in Section 2 but writing the reciprocals,
in -probability. Consider the Lebesgue decomposition into and . Then it follows that
in -probability, which by Scheffé’s lemma implies in total variation. As are probability measures, it follows that and finally in total variation. The convergence of the original sequence follows. ∎
which implies (C.3). Next, we observe from the definition of and (C.4) that for ,
Let and assume without loss of generality that . Then
This concludes the proof of the first estimate in the lemma. Turning to the second, note that by (C.2), (C.3) and the definition of ,
where we chose (ensuring , in particular). Define
Arguing as for (C.3) and (C), now using (C.7) instead of (C.2), we see that and that for ,
We conclude the proof by arguing as in (C.6) but with replaced by . ∎