Computing Pairwise Squared Euclidean Distances

def square_distance(src, dst):
    """
    Compute squared Euclidean distances between each pair of points from two sets.

    The formula used is:
    ||a - b||^2 = ||a||^2 + ||b||^2 - 2 * a^T * b

    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point squared distances, [B, N, M]
    """
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist

🍉 Explanation:

• B, N, _ = src.shape: Extracts batch size B, number of source points N, and feature dimensions C.
• _, M, _ = dst.shape: Retrieves number of destination points M.
• dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)): Computes dot producst between all pairs of points using matrix multiplication.
  • dst.permute(0, 2, 1): Transposes the destination tensor to align dimensions for matrix multiplication.
  • torch.matmul: Calculates pairwise dot products resulting in a [B, N, M] tensor.
• dist += torch.sum(src ** 2, -1).view(B, N, 1): Adds squared norms of source points.
  • torch.sum(src ** 2, -1): Computes squared magnitude for each point in src.
  • view(B, N, 1): Reshapes the result to enable broadcasting during addition.
• dist += torch.sum(dst ** 2, -1).view(B, 1, M): Adds squared norms of destination points.
• Returns the final tensor of shape [B, N, M] representing squared Euclidean distances.

This method efficiently computes squared Euclidean distances using the algebraic identity:

$$ (x_1 - x_2)^2 + (y_1 - y_2)^2 + (z_1 - z_2)^2 = x_1^2 + y_1^2 + z_1^2 + x_2^2 + y_2^2 + z_2^2 - 2x_1x_2 - 2y_1y_2 - 2z_1z_2 $$

Example Usage

Given two sets of points, src and dst:

import torch

def square_distance(src, dst):
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist

# Define source and target point sets
src = torch.tensor([[[1, 2, 3], [4, 5, 6]]])  # shape: [1, 2, 3]
dst = torch.tensor([[[7, 8, 9], [10, 11, 12], [13, 14, 15]]])  # shape: [1, 3, 3]

dist = square_distance(src, dst)
print(dist)

Expected output: For example, the squared distance between (1, 2, 3) and (7, 8, 9) is: $$ (7-1)^2 + (8-2)^2 + (9-3)^2 = 108 $$

Intermediate Computation Details

• dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
• torch.sum(src ** 2, -1) produces a tensor like tensor([[14, 77]]), indicating the squared norms for each source point.
• torch.sum(src ** 2, -1).view(B, N, 1) reshapes it to tensor([[[14], [77]]]) for broadcasting.
• Adding torch.sum(dst ** 2, -1).view(B, 1, M) completes the calculation of the full distance matrix.