A Novel Approach to Sinkhorn ε-Annealing: A 3.62x Speedup for the Optimal Transport Problem

Optimal Transport

The optimal transport problem is found across many different fields. The problem is widespread because it answers the question: what is the cheapest way to transform distribution A into distribution B, given a cost function? Originally, this was named the earth mover’s problem, and today it is used frequently in 2D and 3D.

The fields in which the optimal transport problem is found include Generative AI, Machine Learning, Computer Vision, Robotics, and Biology. (see figures #1-5)

Sinkhorn Algorithm

The Sinkhorn algorithm is the most widely used algorithm for solving the optimal transport problem due to its speed. It achieves this speedup through the parameter, epsilon.

To understand epsilon, first think of a piece of wood that you want to sand down. Sand paper has grits ranging from 40 to 200. When sanding wood, you start with 40 grit, then gradually switch your sand paper until you are at 200 grit: you do not start at 200 grit because it is too slow, and you do not end at 40 grit because it is too rough.

ε is very similar to sandpaper grit. The difference is that ε starts large and then becomes smaller. In the research literature, the terms epsilon scaling and annealing are used interchangeably, and they simply mean to change epsilon as the algorithm progresses. Current methods of epsilon annealing are predetermined, meaning that before starting on the problem, the program has already decided on which values of ε to use throughout the program.

Algorithm Convergence Behavior

To improve on prior methods, the characteristics of the Sinkhorn algorithm must first be understood. The Sinkhorn algorithm converges as you repeat the algorithm. The objective gap is how much the transport cost at this iteration differs from the minimum transport cost possible. As shown in the graph (see figure #6), the objective gap is not as well behaved as its marginal gap counterpart. The marginal gap is mathematically closer to the problem’s core processes, and is thus a better proxy for convergence.

Technique: Projection

The first novel contribution of my project is derived from the marginal gap’s convergence behavior. The marginal gap’s convergence is known to be approximately exponential (figure #6). I repurposed a path-finding ODE to calculate the subtle curvature in the marginal gap convergence graph. Using this technique, an accurate estimate for the number of iterations needed for the algorithm to converge is cheaply calculated.

Technique: Contextual Inference

The second novel contribution of my project builds on convergence guarantees. Saving sample problem snapshots when the initial epsilon path has not yet converged is algorithmically equivalent to having converged problem snapshots for a larger epsilon. Given the current epsilon, there exists a smaller epsilon that has the current epsilon as its most efficient warm start. Put simply, there is a smaller epsilon that has our current epsilon as its “favorite” (figure #8).

The issue is that our decisions are made under limited information due to computational cost restrictions. It is only optimal to select 3 sample points to use in our decision (figure #7). To resolve this issue, the best choice under limited information can be deduced to be at the candidate epsilon where the last two saved points change from increasing to decreasing cost.

Comparison Baselines

The two baselines that I compared my algorithm against are Schmitzer and Chizat’s epsilon annealing technique. They are chosen due to their prominence in optimal transport code libraries and for being two of the most frequently used annealers in practice. The Schmitzer scaler/annealer is the most widely used epsilon scaler in practice, and it is geometric. For example, a schmitzer sequence of epsilons could be 1, 0.5, 0.25, etc. Chizat’s technique was also compared against due to its prominence in the literature’s theory. Chizat’s technique decreases epsilon after every 2 matvecs in the algorithm.

The reason why Matvecs (Matrix-Vector multiplications) are used to measure complexity instead of runtime is because Matvecs are not affected by uncontrolled sources of error such as GPU capacity, CPU speed, etc.

Datasets Used

The four datasets shown on the slides are Kather Texture, Dotmark White Noise, Dotmark Cauchy Density, and ImageNet. Model-Net-40 and Xenium datasets were also used, however, they are not shown on the slides. Across the results, the average improvement over the best previous technique is shown to be a 3.62x reduction in matvecs spent.

General Optimal Transport Benchmark

The Dotmark dataset is used by researchers for benchmarking their optimal transport methods, and serves as a general test for program proficiency.

Computer Vision and Generative AI Benchmark

The data used on the results slide was from ImageNet. 2D optimal transport on natural images is used for training AI to classify and generate images. 3D optimal transport (Model-Net-40 dataset, not on slides) was also used for benchmarking algorithm performance on point clouds. An approximate 3.81x speedup was achieved, significant in computer vision applications where a 3D item model must be reconstructed from 2 camera angles.

Biology Benchmark

The Kather Texture dataset shown on the slides has data of stained human cancer cells. Such datasets are used in AI representation learning for identification of cancer cells. On Xenium data (not on slides) for spatial transcriptomics, the coupling of cell tissue architecture to gene expression, a similar speedup of 3.54x was achieved.

Controlled Problem Parameters

Target Epsilon Value (ending epsilon value)

Convergence Tolerance (maximum marginal gap to count as complete)

Cost Metric (L2)

Scaling Vector Initializations

Significance

Optimal transport is useful in many fields, but it is still expensive to compute. My algorithm reduced the number of matvec operations by an average of 3.62x compared with the best previous technique. In practice, this reduces the runtime and economic cost of the Sinkhorn algorithm.

This improvement has direct applications in AI, Computer Vision, and Biology. In AI training, speedups translate to reduced GPU rental costs and training durations. In Computer Vision, especially for robotics, speedups on optimal transport for 3D object tracking (point cloud matching) will enable significantly lower reaction times.  In Biology, the computational speedup allows for larger sample sizes to be analyzed while staying within the same resource constraints.

Contribution to Python Code Libraries

The open source Python Optimal Transport library is used widely by researchers and engineers to implement optimal transport algorithms in real world applications. Improvements to the code library are impactful across multiple fields. Due to the speedup, I am in the process of contributing to the Python Optimal Transport library, providing my technique as an improvement upon the current Schmitzer epsilon annealing technique.

Conclusion

This innovation shows that adaptive epsilon annealing can outperform predetermined epsilon schedules by using the Sinkhorn algorithm’s own convergence guarantees to inform decisions. The project’s technique uses ODE and exponential convergence information to project convergence costs, then contextual inference selects the best choice under limited information. Together, these methods reduce computational cost while preserving the same target accuracy.

Future Work

To further reduce the runtime of the algorithm, the probing and projection costs of the algorithm could be decreased. This can be done by making mathematical progress that allows for the ODE to generalize or project across an even larger section of iterations without incurring excessive computational expenses.

Once the condition for swapping to the next epsilon is met, the algorithm always makes a choice to switch to another epsilon. To further increase efficiency, it would be viable to introduce the option to continue on the current epsilon for several more iterations before switching.

I would like to thank my parents, teachers and the Edmonton Regional Science Fair Team for their support. I would also like to thank Professor Underhill at the University of British Columbia for providing access to Xenium cancer data.

Journal articles

Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F.-X. (2018). Scaling algorithms for unbalanced optimal transport. Mathematics of Computation, 87(314), 2563–2609.

Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T. H., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V., Sutherland, D. J., Tavenard, R., Tong, A., & Vayer, T. (2021). POT: Python optimal transport. Journal of Machine Learning Research, 22(78), 1–8.

Kather, J. N., Weis, C.-A., Bianconi, F., Melchers, S. M., Schad, L. R., Gaiser, T., Marx, A., & Zöllner, F. G. (2016). Multi-class texture analysis in colorectal cancer histology. Scientific Reports, 6, Article 27988.

Peyré, G., & Cuturi, M. (2019). Computational optimal transport. Foundations and Trends in Machine Learning, 11(5–6), 355–607.

Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., & Fei-Fei, L. (2015). ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision, 115, 211–252.

Schmitzer, B. (2019). Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing, 41(3), A1443–A1481.

Schrieber, J., Schuhmacher, D., & Gottschlich, C. (2017). DOTmark—A benchmark for discrete optimal transport. IEEE Access, 5, 271–282.

Sinkhorn, R., & Knopp, P. (1967). Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2), 343–348.

Images

SuperAnnotate. (2025). Introduction to diffusion models for machine learning [Image]. SuperAnnotate. https://www.superannotate.com/blog/diffusion-models

Van Otten, N. (2023, December 11). Representation learning made simple & top 10 machine learning and deep learning models [Image]. Spot Intelligence. https://spotintelligence.com/2023/12/11/representation-learning/

Lin, Y., & Yang, J. Y. H. (2022). 3D reconstruction of spatial expression [Image]. Nature Methods, 19, 526–527. https://doi.org/10.1038/s41592-022-01476-5

Carnegie Mellon University Robotics Institute. Point Cloud Registration (PCR) [Image]. Carnegie Mellon University. https://www.ri.cmu.edu/project/point-cloud-registration-pcr/

Khokhar, S. (2020, October 21). How to choose the right LiDAR sensor for your project [Image]. Clearpath Robotics. https://store.clearpathrobotics.com/blogs/blog/how-to-choose-a-lidar

Proceedings, conference abstracts:

Cuturi, M. (2013). Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems 26.

Su, H., Maji, S., Kalogerakis, E., & Learned-Miller, E. (2015). Multi-view convolutional neural networks for 3D shape recognition. In Proceedings of the IEEE International Conference on Computer Vision.

Wu, Z., Song, S., Khosla, A., Yu, F., Zhang, L., Tang, X., & Xiao, J. (2015). 3D ShapeNets: A deep representation for volumetric shapes. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.

Webpages

Flamary, R., & Courty, N. POT: Python Optimal Transport. PythonOT. https://pythonot.github.io/

Princeton University. Princeton ModelNet. https://modelnet.cs.princeton.edu/

TensorFlow. (2024). Colorectal histology. TensorFlow Datasets. https://www.tensorflow.org/datasets/catalog/colorectal_histology

ImageNet. Large Scale Visual Recognition Challenge. https://www.image-net.org/challenges/LSVRC/