The Sinkhorn-Knopp algorithm is an iterative technique used to convert a non-negative matrix into a doubly stochastic matrix, which means all rows and columns sum to 1. This property is useful in applications such as optimal transport, machine learning, and statistical normalization.

In this article, we will explore how to implement it in Python using NumPy.

What is the Sinkhorn-Knopp Algorithm?

The Sinkhorn-Knopp algorithm adjusts a non-negative matrix by alternately normalizing its rows and columns until it becomes doubly stochastic.

Doubly stochastic matrix: A square matrix \( A \in \mathbb{R}^{n \times n} \) such that:

• Each row sums to 1
• Each column sums to 1

This makes it useful in fields where matrix marginals must be preserved, such as matching problems, probabilistic models, and deep learning attention mechanisms.

Step-by-Step Algorithm

The algorithm follows a straightforward process:

1. Initialize: Start with a non-negative matrix.
2. Row Normalization: Divide each row of the matrix by its sum so that each row sums to 1.
3. Column Normalization: Divide each column of the matrix by its sum so that each column sums to 1.
4. Repeat: Alternate between row and column normalization until the matrix converges. This means the row and column sums get sufficiently close to 1.

NumPy Implementation

Here’s a Python implementation using NumPy:

import numpy as np

def sinkhorn_knopp(matrix, max_iters=1000, tol=1e-9):
    """
    Perform Sinkhorn-Knopp normalization to make the input matrix doubly stochastic.

    Parameters:
        matrix (np.ndarray): A non-negative square matrix.
        max_iters (int): Maximum number of iterations.
        tol (float): Tolerance for convergence.

    Returns:
        np.ndarray: A doubly stochastic matrix.
    """
    assert matrix.ndim == 2 and matrix.shape[0] == matrix.shape[1], "Matrix must be square"
    K = matrix.copy()
    
    for i in range(max_iters):
        # Row normalization
        K /= K.sum(axis=1, keepdims=True)
        # Column normalization
        K /= K.sum(axis=0, keepdims=True)
        
        # Check convergence (both row and column sums close to 1)
        row_sums = K.sum(axis=1)
        col_sums = K.sum(axis=0)
        if np.allclose(row_sums, 1, atol=tol) and np.allclose(col_sums, 1, atol=tol):
            break

    return K

Example Usage

Let’s apply the algorithm to a sample 3×3 matrix:

# Example usage
A = np.array([
    [0.105, 0.121, 0.107],
    [0.059, 0.193, 0.081],
    [0.169, 0.020, 0.145]
])

A_ds = sinkhorn_knopp(A)
print("Doubly stochastic matrix:")
print(np.round(A_ds, 3))
print("Row sums:", np.round(A_ds.sum(axis=1), 3))
print("Col sums:", np.round(A_ds.sum(axis=0), 3))

Output:

Doubly stochastic matrix:
[[0.316 0.362 0.322]
 [0.178 0.578 0.244]
 [0.506 0.06  0.434]]
Row sums: [1. 1. 1.]
Col sums: [1. 1. 1.]

Applications

The Sinkhorn-Knopp algorithm is applicable across multiple disciplines:

• Optimal Transport: Used to compute entropic regularized transport plans between probability distributions.
• Machine Learning: Helps in attention mechanisms, domain adaptation, and structured prediction tasks.
• Data Normalization: Useful in balancing co-occurrence or contingency tables while preserving marginal sums.
• Economics and Social Sciences: Applied in trade matrices and migration flows where supply and demand must balance.

Limitations and Considerations

• Convergence Issues: The algorithm may not converge if the matrix contains zero rows or columns, or if target marginals are incompatible with the matrix’s support.
• Numerical Stability: Small values can lead to instability due to division operations. Log-domain versions of the algorithm can address this.
• Approximation Quality: The matrix is only approximately doubly stochastic. Precision depends on the tolerance and number of iterations.

Conclusion

The Sinkhorn-Knopp algorithm transforms a non-negative square matrix into a doubly stochastic one, where all rows and columns sum to 1. It works by repeatedly normalizing the rows and then the columns. This process is useful in several fields, including machine learning, data analysis, and economics. The algorithm is easy to implement, but zero entries or small values in the matrix can cause problems. In general, it is a useful technique for solving problems that involve probabilities and optimization.