Cluster-Aware Matching via Laplacian Optimal Transport

Core Methodology: Laplacian Optimal Transport (LapOT)§

The paper addresses the problem of aligning point clouds that exhibit intrinsic cluster structure. Standard optimal transport (OT) finds a point-to-point coupling, but when points within a cluster are interchangeable, a cluster-to-cluster alignment is more robust. LapOT modifies the OT objective by adding quadratic Laplacian regularization terms that promote smoothness of the transport plan with respect to the similarity graphs of each point cloud.

Given two point sets $X = \{x_i\}_{i=1}^n$ and $Y = \{y_j\}_{j=1}^m$ with similarity graphs (e.g., k-NN) having Laplacians $L_X$ and $L_Y$, the LapOT optimization problem is:

$$ \min_{P \in U(a,b)} \langle C, P \rangle + \lambda \left( \text{tr}(P^T L_X P) + \text{tr}(P L_Y P^T) \right) $$

where $C$ is the cost matrix (e.g., Euclidean distances), $a$ and $b$ are marginal distributions (typically uniform), $\lambda > 0$ is a regularization strength, and $P$ is the transport plan. The Laplacian terms penalize large differences in the coupling for nearby points in each cloud. This encourages the mass from a cluster in $X$ to be transported to a cluster in $Y$, rather than scattering individual points.

Refined Simultaneous Clustering (RSC)§

Given the LapOT coupling $P$, RSC produces consistent clusterings across the two point sets. It alternates between updating cluster assignments based on $P$ and refining $P$ using the cluster structure, leading to partitions that are aligned across domains. The method can be seen as a spectral clustering with cross-domain guidance.

Implementation Details§

Solving the regularized OT problem can be done via a generalized Sinkhorn algorithm. The Laplacian terms are quadratic in $P$, so the total objective is convex. One can use proximal methods or incorporate the Laplacian terms into the cost by a change of variables: $\tilde{C} = C + \lambda (L_X \otimes I + I \otimes L_Y)$? More directly, the gradient of the regularizer is $2\lambda (L_X P + P L_Y)$. A simple approach is to use the Sinkhorn algorithm with an additional proximal step.

PyTorch Code Snippet§

import torch
import torch.nn.functional as F

def sinkhorn_lapot(C, Lx, Ly, a, b, lambda_reg=0.1, n_iter=100):
    """Sinkhorn with Laplacian regularization.
    C: cost matrix (n, m)
    Lx: n x n Laplacian of X
    Ly: m x m Laplacian of Y
    a: marginal probabilities for X (n,)
    b: marginal probabilities for Y (m,)
    """
    n, m = C.shape
    P = torch.ones(n, m) / (n * m)  # initialize
    for _ in range(n_iter):
        # Regularization gradient
        grad_reg = 2 * lambda_reg * (Lx @ P + P @ Ly.T)
        # Log-domain update (simplified)
        K = torch.exp(-C - grad_reg)  # not exact, use proper scaling
        # Sinkhorn iterations (simplified)
        P = K * a.view(-1, 1)
        P = P / P.sum(dim=1, keepdim=True)
        P = P * b.view(1, -1)
        P = P / P.sum(dim=0, keepdim=True)
    return P

Note: A rigorous implementation would use log-domain Sinkhorn and handle the Laplacian regularizer via a double stochastic approximation or proximal operator. The snippet illustrates the idea.

Theoretical Analysis§

The paper proves that LapOT preserves cluster structures: if the point clouds have well-separated clusters and the graphs are constructed accordingly, the optimal coupling concentrates mass between corresponding clusters. This provides a theoretical guarantee for the method's effectiveness.

Experiments§

Empirical results on synthetic and real data (e.g., shape matching, document alignment) show that LapOT produces more meaningful alignments than standard OT and other alternatives, especially when clusters are present. RSC further improves clustering consistency.