Designing the Loss for a Multi-View Keypoint Detector

6 minute read

Consider a static, calibrated multi-camera rig looking at a scene, and a model that must find the same set of $K$ semantic keypoints in every view and lift them to 3D. Because the rig is static and calibrated, the intrinsics $\mathbf{K}_v$ and extrinsics $[\mathbf{R}_v \mid \mathbf{t}_v]$ are known and fixed — which means geometry is available to the loss function, not just to a post-processing step.

Most of the design work is deciding how much of the supervision should come from per-image annotation and how much from geometric consistency between views.

Start with per-view heatmaps

Direct coordinate regression from an image is a hard, badly conditioned mapping. Predicting a spatial heatmap per keypoint is easier to learn: it keeps the output aligned with the convolutional feature grid, and it represents ambiguity honestly — an occluded keypoint can come back diffuse instead of confidently wrong.

Targets are Gaussians centred on the ground-truth pixel,

\[H_k(p) = \exp\left(-\frac{\|p - x_k\|^2}{2\sigma^2}\right)\]

and the loss is mean squared error over pixels and keypoints:

\[\mathcal{L}_{\text{hm}} = \frac{1}{VK}\sum_{v=1}^{V}\sum_{k=1}^{K} \nu_{vk} \left\| \hat{H}_{vk} - H_{vk} \right\|_2^2\]

The visibility flag $\nu_{vk}$ matters more than it looks. With $V$ cameras, any given keypoint is invisible in most of them. If you supervise an all-zero target wherever the keypoint is occluded, the background term dominates and the model learns to predict nothing with great confidence. Either mask those terms out, or — if you have a reliable occlusion label — supervise them explicitly with a separate visibility head.

$\sigma$ is worth a sweep. Too tight and the gradient signal is a handful of pixels; too loose and you cap your localization accuracy.

Make the argmax differentiable

A hard argmax over the heatmap is not differentiable and quantizes to the stride of the output grid. Soft-argmax fixes both — take the expectation over a spatially normalized heatmap:

\[\hat{x}_{vk} = \sum_{p \in \Omega} p \cdot \frac{\exp(\hat{H}_{vk}(p)/\tau)}{\sum_{q}\exp(\hat{H}_{vk}(q)/\tau)}\]
def soft_argmax_2d(heatmap, tau=1.0):
    """heatmap: (B, K, H, W) logits -> (B, K, 2) pixel coords in (x, y)."""
    b, k, h, w = heatmap.shape
    prob = torch.softmax(heatmap.reshape(b, k, -1) / tau, dim=-1)
    prob = prob.reshape(b, k, h, w)

    ys = torch.arange(h, device=heatmap.device, dtype=prob.dtype)
    xs = torch.arange(w, device=heatmap.device, dtype=prob.dtype)
    exp_x = (prob.sum(dim=2) * xs).sum(dim=-1)   # marginalize over y
    exp_y = (prob.sum(dim=3) * ys).sum(dim=-1)   # marginalize over x
    return torch.stack([exp_x, exp_y], dim=-1)

Two things to know. The expectation is only meaningful for a unimodal heatmap — with two competing peaks, soft-argmax lands in the empty space between them. And $\tau$ trades sub-pixel precision against gradient smoothness; low $\tau$ approaches a hard argmax and the gradient largely vanishes.

Because these coordinates are differentiable, everything downstream of them can now be part of the loss.

Add the geometry

Triangulation and reprojection errorThree calibrated cameras observe the same semantic keypoint. Their rays are triangulated into a single 3D estimate, which is then projected back into each image. The gap between the projected point and the detected point is the reprojection residual the loss penalizes. cam 1 cam 2 cam 3 X̂ₖ (triangulated) image 1 π(X̂) residual image 3 agrees — near-zero residual
The residual the geometric loss penalizes. A view that localizes the keypoint poorly is pulled toward the position the confident views agree on — supervision you cannot get from per-image annotation.

Triangulate the per-view detections into a 3D point $\hat{X}_k$ — DLT is differentiable, so this stays inside the graph — then project it back into every view and penalize the disagreement:

\[\mathcal{L}_{\text{rep}} = \sum_{v,k} \nu_{vk} \, \rho\!\left( \left\| \pi_v(\hat{X}_k) - \hat{x}_{vk} \right\| \right)\]

Use a robust $\rho$ (Huber, or Geman-McClure) rather than squared error. A single bad view otherwise drags the triangulated point off, which then poisons the reprojection term in every other view — one outlier becomes $V$ bad gradients.

def reprojection_loss(X, x_det, K, T_cam_world, visible, delta=5.0):
    """X: (B, K, 3) triangulated. x_det: (B, V, K, 2) detections in pixels.
    T_cam_world: (V, 4, 4). visible: (B, V, K) in {0, 1}."""
    Xh = torch.cat([X, torch.ones_like(X[..., :1])], dim=-1)        # (B, K, 4)
    Xc = torch.einsum('vij,bkj->bvki', T_cam_world, Xh)[..., :3]    # into each camera

    z = Xc[..., 2:3].clamp(min=1e-3)          # points behind the camera would flip sign
    proj = torch.einsum('ij,bvkj->bvki', K, Xc / z)[..., :2]

    residual = (proj - x_det).norm(dim=-1)
    loss = F.huber_loss(residual, torch.zeros_like(residual),
                        reduction='none', delta=delta)
    return (loss * visible).sum() / visible.sum().clamp(min=1.0)

The clamp on $z$ is not defensive padding — without it, a point triangulated behind a camera projects with a flipped sign and produces a large, confidently wrong gradient that pushes the solution further behind the camera.

This term does real work: it couples the views. A keypoint that a view sees poorly gets pulled toward the position the confident views agree on, which is exactly the supervision signal you want and cannot get from per-image annotation.

If you have 3D ground truth, supervise it directly as well:

\[\mathcal{L}_{3\text{D}} = \sum_k \left\| \hat{X}_k - X_k \right\|_1\]

If you don’t — the common case, since 3D keypoint labels are expensive — you can still constrain the geometry with an epipolar term, penalizing how far a detection in view $v$ lies from the epipolar line induced by the detection in view $v’$:

\[\mathcal{L}_{\text{epi}} = \sum_{k} \sum_{v \neq v'} \nu_{vk}\nu_{v'k} \, d\!\left(\hat{x}_{vk},\; \mathbf{F}_{vv'} \, \hat{x}_{v'k}\right)\]

with $\mathbf{F}_{vv’}$ the fundamental matrix from the known calibration. This is weaker than triangulation — it constrains each point to a line, not to a point — but it needs no 3D labels at all, and it is a good way to exploit unlabelled multi-view captures.

Putting it together

\[\mathcal{L} = \lambda_{\text{hm}}\mathcal{L}_{\text{hm}} + \lambda_{\text{rep}}\mathcal{L}_{\text{rep}} + \lambda_{3\text{D}}\mathcal{L}_{3\text{D}} + \lambda_{\text{vis}}\mathcal{L}_{\text{vis}}\]

Ordering matters as much as the weights. Train on $\mathcal{L}{\text{hm}}$ alone first. Reprojection loss on a randomly initialized detector triangulates noise, and pulling every view toward a meaningless 3D point is an efficient way to reach a bad local minimum. Warm up the heatmaps, then ramp $\lambda{\text{rep}}$ in.

Keep the heatmap term for the whole run. It is the only term anchored to the annotation — drop it and the geometric losses will happily agree with each other on the wrong answer, since a consistent-but-wrong configuration satisfies them perfectly.

For the relative weights, the same multi-task balancing question applies here, and uncertainty weighting is a reasonable default. But treat $\lambda_{\text{hm}}$ as fixed rather than learned — you want one term the optimizer cannot turn down.

Evaluating it

Report per-view 2D PCK and 3D MPJPE separately. They fail in different ways, and an aggregate number will hide the interesting case: a model whose 2D detections got slightly worse but whose 3D points got much better, because the geometric terms finally started resolving depth ambiguity. That trade is usually the one you want.