ALGORITHMS

Point Clouds and Shapes

Computing OT directly on raw point sets with samples.loss, training a neural network shape generator, handling outliers with unbalanced OT, and computing Wasserstein barycenters.

The cost matrix bottleneck

transport.matrix.solve requires you to pass an explicit (B, N, M) cost tensor. For N = M = 10000 in float32, that is 400 MB per problem in the batch. Even for moderate batch sizes, materialising the cost matrix exhausts GPU memory long before the solver itself becomes the bottleneck.

torchmatch.transport.samples.loss solves this by computing the OT loss directly on raw point clouds, never allocating the full N×M matrix. Internally, a Triton kernel streams through the cost block by block, fusing squared-Euclidean cost computation with log-sum-exp accumulation so that only the dual potentials — two vectors of length N and M — need to live in memory.

The API is simpler too: pass the two point clouds and get a scalar loss.

Computing the loss

import torch
import torchmatch

device = "cuda"   # samples.loss requires CUDA

N = 512
# Source: a ring in 2-D.
angles = torch.linspace(0, 2 * torch.pi, N)
x = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)  # (N, 2)

# Target: uniform random points in a square.
y = torch.rand(N, 2) * 2 - 1                                    # (N, 2)

x_gpu = x.to(device).requires_grad_(True)
y_gpu = y.to(device)

loss = torchmatch.transport.samples.loss(x_gpu, y_gpu, blur=0.1)
print(f"Loss: {loss.item():.4f}")

loss.backward()
print(f"Gradient shape: {x_gpu.grad.shape}")   # (N, 2)

blur plays the role of sqrt(reg) in the matrix face: it sets the effective length scale of the entropic regularisation. Larger blur → smoother, denser matching; smaller blur → sharper, sparser matching.

Gradients flow through both x and y, so both point clouds can be learned parameters.

Training a point-cloud generator

Any differentiable loss can be a training objective. Here, train a small MLP decoder to map random latent vectors to a target shape:

import torch
import torch.nn as nn
import torchmatch
import numpy as np

device = "cuda"
rng = np.random.default_rng(42)

class PointDecoder(nn.Module):
    def __init__(self, latent_dim=8, n_points=256):
        super().__init__()
        self.n_points = n_points
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 64),  nn.Tanh(),
            nn.Linear(64, 128),         nn.Tanh(),
            nn.Linear(128, n_points * 2),
        )

    def forward(self, z):
        return self.net(z).reshape(z.size(0), self.n_points, 2)

def make_s_shape(n_points=256):
    t = np.linspace(0, 2 * np.pi, n_points // 2)
    top = np.stack([0.5 * np.cos(t), 0.35 + 0.35 * np.sin(t)], axis=1)
    bot = np.stack([0.5 * np.cos(t + np.pi), -0.35 + 0.35 * np.sin(t + np.pi)], axis=1)
    pts = np.concatenate([top, bot], axis=0)
    pts += rng.normal(0, 0.03, pts.shape).astype(np.float32)
    return pts

LATENT_DIM = 8
N_POINTS = 256
model = PointDecoder(latent_dim=LATENT_DIM, n_points=N_POINTS).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-3)

target = torch.from_numpy(make_s_shape(N_POINTS)).to(device)

for step in range(300):
    z = torch.randn(1, LATENT_DIM, device=device)
    pred = model(z).squeeze(0)   # (N_POINTS, 2)

    loss = torchmatch.transport.samples.loss(pred, target, blur=0.05)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 50 == 0:
        print(f"step {step:3d}  loss={loss.item():.4f}")

The Wasserstein loss is well-suited to this task: it measures how far the generated point cloud is from the target in the geometric sense, independent of which generated point corresponds to which target point.

CPU fallback

samples.loss raises RuntimeError on CPU because the Triton kernels require a CUDA device. For testing or environments without a GPU, fall back to the matrix face with torch.cdist:

if pred.device.type == "cuda":
    loss = torchmatch.transport.samples.loss(pred, target, blur=0.05)
else:
    C = torch.cdist(pred.unsqueeze(0), target.unsqueeze(0)).pow(2)
    log_plan = torchmatch.transport.matrix.solve(C, reg=0.05 ** 2, n_iter=100)
    loss = (log_plan.exp() * C).sum()

Unbalanced OT: handling outliers

Standard OT enforces the marginal constraints exactly: every unit of source mass must be transported to the target, and vice versa. When the source contains outliers — points far from any plausible target — this forces nearby inliers to stretch toward the outliers to satisfy the row-sum constraint, distorting the matching.

Unbalanced OT relaxes the marginal constraints with a KL penalty. Points that cannot find a reasonable match are soft-discarded rather than forced into the plan. The reach parameter controls this: smaller values discard outliers more aggressively.

import torch
import torchmatch

device = "cuda"
N_CLEAN = 200
N_OUTLIERS = 30

rng_t = torch.Generator().manual_seed(0)
angles = torch.linspace(0, 2 * torch.pi, N_CLEAN)
x_clean = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)

# Outliers far from the ring.
outliers = torch.rand(N_OUTLIERS, 2, generator=rng_t) * 1.5 + 2.5
x_noisy = torch.cat([x_clean, outliers], dim=0).to(device)

y = x_clean.to(device)   # clean target

loss_balanced   = torchmatch.transport.samples.loss(x_noisy, y, blur=0.1)
loss_unbalanced = torchmatch.transport.samples.loss(x_noisy, y, blur=0.1, reach=0.3)

print(f"Balanced   loss: {loss_balanced.item():.4f}")
print(f"Unbalanced loss: {loss_unbalanced.item():.4f}")

The unbalanced loss is lower because the outlier points are not forced to participate in the matching. Gradient-wise, the gradient through the inlier points is less contaminated by the outliers' distorting pull.

The reach parameter is in the same units as blur. A practical rule: set reach to roughly the largest inlier-to-target distance you are willing to tolerate; anything further will be soft-discarded.

Wasserstein barycenters

The Euclidean mean of two distributions is a pixel-wise average: blurring. The Wasserstein barycenter is the distribution that minimises the average OT distance to a set of input distributions. It preserves structure in a geometrically meaningful way.

Consider three concentric rings at radii 0.5, 1.0, and 1.5. Their Euclidean mean is a smeared annulus. Their Wasserstein barycenter is a clean ring at radius 1.0.

import torch
import numpy as np

rng = np.random.default_rng(7)

def make_ring(n, radius=1.0, noise=0.05):
    angles = np.linspace(0, 2 * np.pi, n, endpoint=False).astype(np.float32)
    pts = np.stack([radius * np.cos(angles), radius * np.sin(angles)], axis=1)
    pts += rng.normal(0, noise, pts.shape).astype(np.float32)
    return torch.from_numpy(pts)

N = 64
shapes = [make_ring(N, r) for r in [0.5, 1.0, 1.5]]

# Euclidean mean: blurs the three rings into a smeared annulus.
mean_euclidean = torch.stack(shapes).mean(dim=0)

# A free-support Wasserstein barycenter could be computed by optimising
# a point cloud Z to minimise the average OT loss to each input shape:
# loss = sum_i samples.loss(Z, shapes[i], blur=0.05) / len(shapes)
# Gradient descent on Z converges to the Wasserstein barycenter.
Z = make_ring(N, radius=1.0).to("cuda").requires_grad_(True)
opt = torch.optim.Adam([Z], lr=1e-2)
shapes_gpu = [s.to("cuda") for s in shapes]

for _ in range(200):
    loss = sum(
        torchmatch.transport.samples.loss(Z, s, blur=0.05)
        for s in shapes_gpu
    ) / len(shapes_gpu)
    opt.zero_grad()
    loss.backward()
    opt.step()

print(f"Barycenter radius ≈ {Z.detach().norm(dim=1).mean().item():.2f}")
# ≈ 1.0 — the geometric mean of 0.5, 1.0, 1.5

The barycenter is a ring at the average radius, not a blurred combination of all three. This generalises: the Wasserstein barycenter of a set of shapes is itself a shape with coherent structure.

Summary of `samples.loss` parameters

Parameter	Role	Default
`blur`	Regularisation length scale (≈ sqrt(ε))	required
`reach`	KL relaxation for unbalanced OT; `None` = balanced	`None`
`debias`	Apply Sinkhorn divergence debiasing	`False`
`half_cost`	Use `‖x−y‖` instead of `‖x−y‖²`	`False`