Optimal Transport and Sinkhorn Iterations in AI

This article explores how optimal transport and Sinkhorn iterations are used to align probability distributions in artificial intelligence representation learning. It breaks down the mathematical concept of moving probability mass with minimal effort, explains why standard optimal transport is too slow for modern deep learning, and demonstrates how Sinkhorn iterations solve this computational bottleneck using entropic regularization. Finally, it covers practical applications of these techniques in domain adaptation, self-supervised learning, and generative modeling.

Understanding Representation Learning and Distributions

In artificial intelligence, representation learning teaches models to convert raw data—like images, text, or audio—into compact numerical vectors called embeddings. When an AI processes a large dataset, the resulting embeddings form a high-dimensional probability distribution.

Often, an AI system needs to compare or align two different distributions. For example, a model might need to align the embeddings of English sentences with French sentences, or match real images with AI-generated images. Traditional metrics like Kullback-Leibler (KL) divergence fail when these distributions do not overlap in the embedding space. Optimal transport solves this problem by taking the underlying geometry of the space into account.

The Role of Optimal Transport

Optimal transport (OT) is a mathematical framework originally developed to find the most cost-effective way to move a pile of dirt from one location to another. In AI, the "dirt" is probability mass, and the "locations" are data representations in an embedding space.

To align two distributions, optimal transport calculates the Wasserstein distance (also known as Earth Mover's Distance). This metric measures the minimum effort required to transform one probability distribution into another.

  1. The Cost Matrix: First, the algorithm measures the distance (usually Euclidean or cosine distance) between every data point in the first distribution and every point in the second. This creates a cost matrix.
  2. The Transport Plan: Next, the algorithm finds a "transport plan"—a matrix that dictates exactly how much probability mass should move from each point in the source distribution to each point in the target distribution to minimize the total distance traveled.

Because OT relies on the distance between points rather than distribution overlap, it provides smooth, meaningful gradients for training neural networks, even when the two distributions are far apart.

The Computational Bottleneck

While exact optimal transport is mathematically ideal, it is computationally expensive. Finding the exact transport plan requires solving a linear programming problem. The computational complexity of this process scales cubically with the number of data points.

If a neural network processes a batch of 1,000 samples, calculating the exact optimal transport becomes too slow for standard training loops. Furthermore, the exact solution is not always differentiable, making it incompatible with backpropagation, the core mechanism by which neural networks learn.

How Sinkhorn Iterations Solve the Problem

Sinkhorn iterations bypass the computational bottleneck of exact optimal transport by introducing a concept called entropic regularization.

Instead of forcing the algorithm to find the single sharpest, most exact path between points, entropic regularization adds a penalty term that encourages the transport plan to be "blurry" or spread out. It maximizes the entropy of the transport plan alongside minimizing the transport cost.

This small mathematical adjustment changes the problem entirely, allowing it to be solved using the Sinkhorn-Knopp algorithm.

  1. Initialization: The cost matrix is converted into a similarity kernel using an exponential function controlled by a regularization parameter.
  2. Iterative Scaling: The Sinkhorn algorithm then alternates between normalizing the rows and normalizing the columns of this matrix so that they sum to the respective marginal distributions of the source and target data.
  3. Convergence: After a fixed number of these simple row-and-column division steps (Sinkhorn iterations), the matrix converges to an approximate optimal transport plan.

This approach reduces the computational complexity from cubic to quadratic. More importantly, Sinkhorn iterations rely entirely on matrix-vector multiplications. These operations run rapidly on GPUs and are fully differentiable, allowing gradients to flow seamlessly backward through the Sinkhorn steps to update the neural network's weights.

Applications in AI Representation Learning

By combining the geometric awareness of optimal transport with the speed of Sinkhorn iterations, AI researchers can align distributions across several domains:

  • Self-Supervised Learning: Algorithms like SwAV (Swapping Assignments between Views) use Sinkhorn iterations to cluster image representations on the fly. The model aligns the distribution of incoming image features with a uniform distribution of "prototype" clusters, preventing the model from collapsing all representations into a single point.
  • Domain Adaptation: When a model trained on synthetic data is deployed on real-world data, performance drops due to distribution shift. Sinkhorn-based optimal transport aligns the feature distribution of the real-world data with the synthetic data, allowing the model to generalize without needing new labels.
  • Generative Modeling: In generative adversarial networks (GANs) and autoencoders, Sinkhorn distances are used as loss functions to force the distribution of generated data to match the distribution of real data, stabilizing the training process and preventing mode collapse.