Point-cloud Wasserstein loss

This tutorial builds a training loop that uses transport.samples.loss as a geometry-aware loss between predicted and ground-truth 3D point clouds. The pattern applies to shape autoencoders, generative models, and any task where you want the model to produce a point set close to a target — measured by the cost of moving one set of points onto the other.

Prerequisites: a CUDA device and pip install torchmatch.

The problem

Suppose a decoder network takes a latent vector z and produces a set of 3D points. A naïve MSE loss on unordered point sets requires a fixed one-to-one pairing between predicted and ground-truth points first — but no such pairing exists when both sets are unordered. The Wasserstein loss sidesteps the correspondence problem: it measures the minimum cost of moving predicted points onto ground-truth points — each predicted point contributes proportionally to nearby ground-truth points, without needing a fixed pairing.

Setup

import torch
import torch.nn as nn
import torchmatch

device = torch.device('cuda')

# --- Tiny decoder for illustration ---
class Decoder(nn.Module):
    def __init__(self, latent_dim: int = 64, n_points: int = 512):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.ReLU(),
            nn.Linear(256, 512), nn.ReLU(),
            nn.Linear(512, n_points * 3),
        )
        self.n_points = n_points

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        # Returns (batch, n_points, 3)
        return self.net(z).reshape(z.size(0), self.n_points, 3)

decoder = Decoder().to(device)
optimiser = torch.optim.Adam(decoder.parameters(), lr=1e-4)

Synthetic target distribution

For illustration, we use a fixed Gaussian mixture as the ground-truth shape. In a real pipeline this would come from a dataset loader.

def make_target(batch: int, n_pts: int, device: torch.device) -> torch.Tensor:
    """Synthetic 3-D point cloud: 4 Gaussian clusters on the unit sphere surface."""
    centres = torch.tensor([
        [1.0, 0.0, 0.0], [-1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0], [0.0, -1.0, 0.0],
    ], device=device)                               # (4, 3)
    idx  = torch.randint(4, (batch, n_pts), device=device)
    pts  = centres[idx] + 0.08 * torch.randn(batch, n_pts, 3, device=device)
    return pts / pts.norm(dim=-1, keepdim=True)     # project to sphere

Training loop

BATCH  = 16
LATENT = 64
N_PTS  = 512    # points per cloud
BLUR   = 0.05   # Sinkhorn temperature

for step in range(1000):
    z          = torch.randn(BATCH, LATENT, device=device)
    gt_clouds  = make_target(BATCH, N_PTS, device)  # (B, N, 3)
    pred_clouds = decoder(z)                         # (B, N, 3)

    # transport.samples.loss operates on pairs of (N, D) tensors, so
    # loop over the batch dimension (or use torch.vmap for batching).
    total_loss = torch.tensor(0.0, device=device)
    for pred, gt in zip(pred_clouds, gt_clouds):
        total_loss = total_loss + torchmatch.transport.samples.loss(
            pred, gt, blur=BLUR,
        )
    loss = total_loss / BATCH

    optimiser.zero_grad()
    loss.backward()
    optimiser.step()

    if step % 100 == 0:
        print(f'step {step:4d}  loss {loss.item():.4f}')

Using the Sinkhorn divergence

The Sinkhorn solver approximates the true Wasserstein loss by adding an entropy term controlled by ε (equivalently, blur²). This approximation has two side-effects: the raw Sinkhorn loss S_ε(x, y) is not symmetric and does not vanish when x == y for finite ε. The debiased Sinkhorn divergence corrects both properties:

D_ε(x, y) = S_ε(x, y) − ½ S_ε(x, x) − ½ S_ε(y, y)

Pass debias=True to use it. Each loss call makes three forward solver calls, which is roughly 3× slower but yields a proper divergence.

loss = torchmatch.transport.samples.loss(pred, gt, blur=BLUR, debias=True)

Use debias=True when:

• the loss value needs to compare meaningfully across training (it equals zero when the two clouds match, making values comparable across iterations)
• the loss must equal zero when predicted and ground-truth clouds are identical (e.g., when reporting the loss as a validation metric)

Use the default (debias=False) when training speed matters. The bias shifts loss values by a constant that depends only on each cloud individually, which usually does not change the direction of the gradient.

Handling outliers with unbalanced OT

Standard OT requires every predicted point to be fully matched to some ground-truth point, so outliers are forced to pair with the nearest GT point and inflate the loss. Unbalanced OT relaxes this requirement via a reach parameter: points that are far from any counterpart are allowed to go unmatched, at a penalty proportional to how much mass is left unaccounted for.

loss = torchmatch.transport.samples.loss(
    pred, gt,
    blur=BLUR,
    reach=0.5,  # smaller reach → more outlier tolerance
)

Or use asymmetric reach when only one side has outliers:

loss = torchmatch.transport.samples.loss(
    pred, gt,
    blur=BLUR,
    reach_x=0.3,  # predicted cloud has outliers; relax source marginal
    # reach_y left at None → target marginal is exact
)

Choosing blur

blur controls the smoothness of the transport plan. Larger values assign mass more diffusely and make the loss faster to compute but less sensitive to fine-grained geometry. It equals the square root of the regularisation strength ε used in the algorithm. Practical ranges:

To set blur automatically: normalise the point clouds to zero-mean unit variance before computing the loss, and fix blur=0.05.

Performance notes

• transport.samples.loss materialises no N × M cost matrix. The Triton streaming kernel computes pairwise costs tile by tile without ever storing the full N×M matrix, keeping memory use proportional to N+M rather than N×M.
• For the balanced case, a single forward + backward pass over 512-point clouds costs about 0.5–2 ms on an A100 depending on blur.
• Loop over the batch dimension (as in the example above) or use torch.vmap for batching. A native batched samples face is planned.

See also

• Algorithms — the mathematical derivation of Sinkhorn and debiasing.
• Reference — full signature for samples.loss.
• Choosing — when samples.loss beats matrix.solve.