Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory

arXiv:2609.33814 · math.OC, cs.LG · Submitted 2026-09-27 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Annealed Sinkhorn with Momentum".

Tom: The gist Annealed Sinkhorn with Momentum provides an exact algebraic characterization of Bregman Douglas–Rachford splitting for unregularized discrete optimal transport and develops an anytime primal–dual certificate in linear memory.

Jane: First, who's behind it and why it matters.

Paper summary: Tom: So we’re looking at this paper right now, "Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory". Basically, they take something called Bregman Douglas–Rachford splitting and show how it works for unregularized discrete optimal transport.

Jane: That sounds really technical, Tom. What’s the main claim here?

Tom: The big claim is that they develop an anytime primal–dual certificate that fits into linear memory, which means you don't need a huge amount of space to run these kinds of calculations. It also shows that this splitting method is equivalent to warm-started Inexact Proximal point method for exact Optimal Transport using just one inner Sinkhorn iteration.

Lu: From an AI perspective, the idea of eliminating the primal transport plan from the updates and deriving a dual formulation is interesting because it changes how we think about these optimization problems.

Meng: So, if this is equivalent to IPOT with a single Sinkhorn iteration, what does that mean for practical speed?

Tom: It means they can get results faster. They show that by eliminating the plan from the updates, you get a dual formulation that acts like annealed Sinkhorn with an implicit inverse-linear temperature schedule and some momentum plus a cooler kernel.

Jane: That sounds like a lot of moving parts in the math, but what about this overrelaxed version they introduce?

Tom: They introduce Overrelaxed BDRS, or OBDRS. This combines the annealing and overrelaxed scaling into one single recursion governed by a parameter lambda between one and two that steepens that implicit cooling schedule <ref:2609.33814#pg1>.

Lalam: I see the benefit here from an LLM perspective; having a certificate computable in linear memory means we can deploy these solvers on devices with limited resources without needing massive amounts of storage for the transport plan itself.

Meng: So, what’s the actual practical performance they show? Are we talking about hours or minutes?

Tom: They have some scale results. For pixel-level color transfer between one thousand twenty-four by one thousand twenty-four images, Dual BDRS hits a relative duality gap of about one point six percent in twenty-four minutes, which is a nine times speedup compared to MDOT–TNT.

Jane: And for larger problems, they tested bidirectional color transfer between two images that are four thousand two hundred thirty-eight by two thousand three hundred sixty-five pixels. How does that compare?

Tom: For those large images, Dual BDRS completed both directions in about thirty-five hours each, achieving best relative duality gaps of two point four one percent and two point eight zero percent <ref:2609.33814#pg1>. That’s solid performance for transport problems of that size.

Lu: The connection they make between the dual formulation revealing it as annealed Sinkhorn under that specific schedule is a nice piece of theoretical grounding for why the method converges in this way.

Jane: It connects the theory of annealing directly to how the solver operates, which makes it easier to understand why it performs well.

Paper summary: Meng: But I wonder about what this actually means for real-world engineering applications? Does this linear memory requirement really solve a problem for large-scale data processing?

Tom: It does address a major bottleneck. The paper notes that the memory requirement goes from quadratic down to linear, which is huge when dealing with many source and target atoms in dense problems.

Lalam: And if we think about how this advance in computation could impact the way we train large models, having an anytime certificate means we have a reliable stopping rule without needing to know the exact final plan beforehand.

Jane: It’s about making optimization processes more robust and predictable when you don't have perfect information upfront.

Tom: The authors also laid out some limitations they have to be clear about. They say their algebraic equivalences and the certificate they derive don't establish global convergence or an a-priori convergence rate for BDRS or O-BDRS, which is important context.

Meng: And what about the memory issue again? The paper mentions that linear memory doesn't eliminate the work per iteration; it still has that O(NM) pairwise work for dense problems with N source and M target atoms.

Tom: True. So, while the certificate is in linear memory, they aren't magically eliminating all the computational work required inside each step for very large dense setups.

Jane: And they also point out that since their certificate is proved in exact arithmetic, we still need to validate how it behaves when we switch to floating-point evaluation and what tolerances are needed.

Lu: The paper’s work on characterizing the temperature schedules under which annealed Sinkhorn converges to unregularized OT, quantifying the relaxation error induced by annealing, and proposing a debiased variant is significant context for understanding these trade-offs.

Tom: So, looking at it simply, this paper provides a way to describe a complex transport algorithm using simpler mathematical concepts like annealing and momentum while giving us a memory-efficient way to check if the solution is good enough.

Jane: And the introduction of OBDRS with that overrelaxation parameter lambda gives us another lever to control how aggressive the cooling schedule gets during the iterative process.

Meng: It’s about adding control knobs so engineers can tune these solvers more effectively than just tweaking one single temperature setting.

Lalam: For culture, this shows a path where complex mathematical proofs can directly translate into tools that are usable by people who are focused on building applications, not just pure mathematics.

Tom: So to wrap up on "Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory", it’s about making these powerful transport methods more accessible through better memory usage and a clearer way to monitor convergence anytime.

Conclusion: Tom: So we’re wrapping up our look at "Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory." This paper shows how to describe a complex transport problem using simpler math, and it gives us a memory-efficient way to check if our solution is good enough.

Jane: It really boils down to taking a complicated optimization method, Bregman Douglas–Rachford splitting, and showing that it’s actually just an advanced version of something we already know called Sinkhorn.

Lu: Exactly. The big deal is they connect it directly to Inexact Optimal Transport using just one step of the Sinkhorn iteration. It’s a very clean algebraic characterization you get there.

Meng: So, if this holds up, it means we can use these transport methods on systems that have very limited memory without crashing our computers. That’s a huge practical win for large-scale data processing.

Lalam: From my side, the ability to compute these certificates in linear memory is pretty powerful. It suggests we can deploy solvers on devices that don't have massive storage constraints for the plan itself.

Tom: And they show this method scales surprisingly well in practice. For color transfer between two one thousand twenty-four by one thousand twenty-four images, it gets a result in twenty-four minutes with only about a one-and-a-half percent gap compared to the best known methods.

Jane: That speed is impressive, especially when you compare it to those other transport algorithms we’ve been using. It shows how much efficiency this new characterization brings to the table.

Lu: The authors also introduce this overrelaxed version, OBDRS, which lets you control the cooling schedule with a parameter called lambda between one and two. It gives more knobs to turn in the optimization process.

Meng: That extra knob is useful because it helps steer the algorithm toward a better solution faster during those iterative steps. It’s about having more control over the convergence path, not just one fixed setting.

Lalam: And that anytime certificate they developed means we have a reliable stopping rule without needing to guess how many iterations we need to run beforehand. That predictability is valuable for any complex process.

Tom: So, looking at the whole picture, this paper takes some heavy machinery—splitting methods and annealing—and gives us a clear map of how it works while making the computational requirements much more manageable.

Jane: It’s about taking these deep theoretical ideas and showing exactly how they translate into a tool that can actually run on real hardware with limited memory.

Lu: The implication is that we can build more complex transport systems without being bottlenecked by massive data storage needs for the intermediate plans.

Meng: For engineers, it means we can tackle bigger problems in color transfer or data movement than we thought possible because the memory hurdle is lowered.

Lalam: And for the broader culture of AI research, having these robust, verifiable solvers that are memory-efficient opens up new avenues for applying transport models to areas where they weren't before.

Tom: This paper really lays out a solid foundation for how we can understand and optimize these kinds of complex matching problems in a way that’s both mathematically sound and computationally feasible.

MIT · HEC Paris

math.OC, cs.LG

Submitted: 2026-09-27

Updated: 2026-09-27

Comments: 38 pages, 11 figures

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 85/100

The gist: The gist Annealed Sinkhorn with Momentum provides an exact algebraic characterization of Bregman Douglas–Rachford splitting for unregularized discrete optimal transport and develops an anytime

Key concepts

Bregman Douglas–Rachford Splitting (BDRS): This is a method used to solve optimal transport problems by splitting the problem into two simpler parts and iteratively updating them. The paper shows it has a specific structure that relates it directly to other known optimization methods, specifically Inexact Proximal point methods for exact Optimal Transport.
Annealed Sinkhorn with Momentum: This describes the underlying mathematical structure of BDRS. It means the algorithm behaves like a standard Sinkhorn iteration where the temperature (or regularization) is gradually reduced over time, but it also incorporates a momentum term to help guide the updates more effectively during this cooling process.
Overrelaxed BDRS (OBDRS): This is an extension of BDRS that incorporates an overrelaxation parameter. This parameter steepens the implicit cooling schedule, meaning it speeds up how quickly the algorithm converges towards a solution by adjusting the rate at which the temperature is reduced during iterations.
Primal–Dual Certificate in Linear Memory: This is a mathematical proof that allows researchers to determine if an algorithm has found an optimal solution without needing to store all intermediate steps of the transport plan. The key finding is that this certificate can be computed using only linear memory, making it practical for large-scale problems.

Terminology

Annealed Sinkhorn — 54× in this paper
Optimal Transport — 37× in this paper · explained in 5 episodes
Ablation — 18× in this paper · explained in 1 episode
Accuracy — 8× in this paper · explained in 1 episode
Annealed Sinkhorn with Momentum — 8× in this paper · explained in 1 episode
Benchmark — 6× in this paper · explained in 2 episodes
Core — 6× in this paper · explained in 2 episodes
Debiased Annealed Sinkhorn — 4× in this paper
Regularization — 4× in this paper · explained in 3 episodes
Generative AI — 3× in this paper · explained in 2 episodes
Log-Averaged Mirror Prox — 2× in this paper
Python Optimal Transport — 2× in this paper
Computational Biology — 1× in this paper · explained in 2 episodes
Grounding — 1× in this paper · explained in 3 episodes
Objective Function — 1× in this paper · explained in 2 episodes
Optimal Transport (OT) — 1× in this paper · explained in 4 episodes
Trace — 1× in this paper · explained in 2 episodes
Verification — 1× in this paper · explained in 2 episodes
Wasserstein Distance — 1× in this paper · explained in 4 episodes

Summary

The gist Annealed Sinkhorn with Momentum provides an exact algebraic characterization of Bregman Douglas–Rachford splitting for unregularized discrete optimal transport and develops an anytime primal–dual certificate in linear memory.

How it works

The paper establishes that BDRS coincides with warm-started Inexact Proximal point method for exact Optimal Transport (IPOT) using a single inner Sinkhorn iteration
By eliminating the primal transport plan from the updates, an equivalent dual formulation is derived, revealing BDRS as annealed Sinkhorn under an implicit inverse-linear temperature schedule, augmented by a momentum-type extrapolation of the dual potentials
The paper introduces Overrelaxed BDRS (OBDRS), which combines annealing and overrelaxed scaling within a single recursion, governed by a single parameter λ ∈ (1, 2) that steepens the implicit cooling schedule
A primal–dual certificate for both Dual BDRS and O-BDRS can be computed in linear memory without the need to compute the primal transport plan, providing a computable stopping rule

Key Results and Characterizations

- Exact characterization and linear memory: BDRS is equivalent to IPOT with a single Sinkhorn iteration and derives a dual formulation of BDRS, revealing it to be AS (Chizat, 2024) under an implicit inverse-linear cooling schedule with a momentum term and cooler kernel (

- Overrelaxation: Guided by the annealing interpretation, OBDRS is introduced by applying an overrelaxation parameter λ ∈ [1, 2) that steepens the implicit cooling schedule (

- Anytime certificate in linear memory: A computable primal–dual certificate for both Dual BDRS and O-BDRS is derived without the need to compute the primal transport plan, providing a stopping rule that preserves linear memory (

Performance and Practical Consequences

- Scale: For pixel-level color transfer between 1024×1024 images, Dual BDRS reaches a relative duality gap of 1.6% in 24 minutes, attaining a lower feasible transport cost than MDOT-TNT with a 9× speed up (

- Large-Scale Pixel-Level Color Transfer: For bidirectional color transfer between two 4238 × 2365 images, Dual BDRS completes both directions in approximately 35 hours each, achieving best relative duality gaps of 2.41% and 2.80% (

Limitations

- Limitations: Our algebraic equivalences and certificate do not establish global convergence or an a-priori convergence rate for BDRS or O-BDRS (

- Limitations: Linear memory also does not remove the O(NM) pairwise work per iteration for dense problems with N source and M target atoms (

- Limitations: The certificate is proved in exact arithmetic and thus its floating-point evaluation and attainable tolerances require numerical validation (

The paper demonstrates that BDRS is equivalent to IPOT with a single inner Sinkhorn iteration and reveals it to be Annealed Sinkhorn with log scaling momentum on a cooler kernel Furthermore, the introduction of O-BDRS allows for a principled overrelaxation and an anytime optimality certificate that can be evaluated in linear memory using the available scalings Finally, the experiments show that combining momentum with a cooler kernel produces substantially smaller rounded primal optimality gaps than annealed Sinkhorn under the same schedule

The paper establishes that BDRS is identical to IPOT with a single Sinkhorn iteration and derives an equivalent dual formulation Furthermore, the dual formulation reveals BDRS to be Annealed Sinkhorn with log scaling momentum on a cooler kernel Moreover, the introduction of O-BDRS allows for a principled overrelaxation and an anytime optimality certificate that can be evaluated in linear memory using the available scalings Finally, the experiments show that combining momentum with a cooler kernel produces substantially smaller rounded primal optimality gaps than annealed Sinkhorn under the same schedule The paper establishes that BDRS is identical to IPOT with a single Sinkhorn iteration and derives an equivalent dual formulation Furthermore, the dual formulation reveals BDRS to be Annealed Sinkhorn with log scaling momentum on a cooler kernel Moreover, the introduction of O-BDRS allows for a principled overrelaxation and an anytime certificate that can be evaluated in linear memory using the available scalings Finally, the experiments show that combining momentum with a cooler kernel produces substantially smaller rounded primal optimality gaps than annealed Sinkhorn under the same schedule The paper establishes that BDRS is identical to IPOT with a single Sinkhorn iteration and derives an equivalent dual formulation Furthermore, the dual formulation reveals BDRS to be

Improvements for AI systems

Improved solvers can achieve near-linear complexity for unregularized discrete optimal transport by utilizing Dual BDRS, which eliminates this bottleneck by removing Xk from the intermediate updates and requires only the one-dimensional vectors. This allows for solving large-scale OT problems on dense matrices at scale, as demonstrated by achieving a 9× speed up compared to MDOT-TNT in pixel-level color transfer.
Improved stopping rules can be implemented with linear memory by deriving a computable primal–dual certificate for both Dual BDRS and O-BDRS without the need to compute the primal transport plan. This provides a computable stopping rule that preserves linear memory, allowing practitioners to monitor solution quality using the duality gap without materializing the transport plan or its feasible repair.
Improved algorithms can incorporate principled overrelaxation via Overrelaxed BDRS (O-BDRS) governed by a parameter λ ∈ (1, 2) to steepen the implicit cooling schedule, which allows for a principled manner to obtain a faster annealing rate and offers an advantage in achieving lower rounded primal optimality gaps at fixed η.
Improved solvers can achieve faster convergence at fixed temperatures by using O-BDRS, as the results indicate that much of its fixed-η advantage is explained by reaching lower effective temperatures sooner, which allows O-BDRS to attain a comparable objective value using approximately 1/λ as many iterations when compared to BDRS.
Improved solvers can handle very large-scale problems, such as pixel-level color transfer between 4238×2365 images, by employing Dual BDRS with η = 0.01 and λ = 1, which achieves relative primal–dual gaps of 2.41% and 2.80% within approximately 35 hours on a single NVIDIA L40S GPU.

Abstract

We characterize Bregman Douglas-Rachford splitting (BDRS) for unregularized discrete optimal transport and develop an anytime primal-dual certificate in linear memory. We first establish that BDRS coincides with warm-started Inexact Proximal point method for exact Optimal Transport (IPOT) using a single inner Sinkhorn iteration. By eliminating the primal transport plan from the updates, we derive an equivalent dual formulation that reveals BDRS as annealed Sinkhorn under an implicit inverse-linear temperature schedule, with an additional log-scaling momentum term and a cooler kernel. While this explains the role of the temperature parameter in BDRS as an initial temperature, it also reduces the solver's memory requirement from quadratic to linear. Utilizing this annealing perspective, we introduce overrelaxed BDRS, which combines annealing and overrelaxed scaling within a single recursion. We derive a primal-dual certificate for both methods that can be evaluated in linear memory without transport plan construction, thus providing a computable stopping rule. On the DOTmark benchmark, combining momentum with the cooler kernel produces substantially smaller optimality gaps than annealed Sinkhorn under the same schedule. For pixel-level color transfer between 1024 times1024 images, BDRS attains a lower repaired transport cost than MDOT-TNT with a 9 times speed up, reaching a relative duality gap of 1.59% in 24 minutes. We further demonstrate a color transfer with 4238 times2365 images, yielding 10 million pixels per image and approximately one hundred trillion implicit transport entries, reaching a best relative duality gap of 2.41% and 2.80% within 35 hours in each direction on a single NVIDIA L40S GPU.