Learned Suppression for 3D Keypoint Detection

Abstract

Detecting 3D keypoints is a long-standing challenge in computer vision. Most detectors end with a heuristic post-processing step that is not learned. We propose a 3D keypoint detector that improves on this step with a learned suppression module, paired with a Point Transformer backbone that we extend with a directional graph neural network. The module is a graph network over candidates that learns which to keep, which to suppress, and how to relocate the remaining ones. Paired with three backbones, it improves over DBSCAN and greedy non-maximum suppression, and because it operates on candidate features rather than raw geometry, the same formulation applies to both structural and semantic keypoints. Our model surpasses the per-category trained KeypointDETR on 12 of 16 KeypointNet categories, attains the best Corner F1 on the Building3D Entry-Level benchmark, and remains competitive with BWFormer on the larger Tallinn split.

Method

Backbone

The backbone predicts, for every point, a keypoint probability and an offset to the nearest keypoint. It builds on a hierarchical Point Transformer encoder with two extensions: a spectral positional encoding that adapts its frequency bands to the geometry of each local neighborhood, and a directional graph neural network after the decoder that passes signed, direction-aware messages between neighboring points to reinforce geometrically distinctive locations such as corners.

Point Transformer block and spectral positional encoding

Learned suppression

High-confidence points vote at their offset-displaced positions, and the votes are reduced to a small candidate set. Instead of a fixed rule deciding which candidates survive, a graph network over the candidates makes that decision. In each of two message-passing rounds, every pair of nearby candidates predicts how strongly one suppresses the other, each candidate updates its state from the suppression it receives, and then moves toward a better position before the next round. A final keep score, discounted by the strongest suppression a candidate received, selects the output keypoints.

Earlier learned NMS only rescored detections in place. In 3D a keypoint rarely coincides with any candidate, so relocation is what lets the module improve localization, not just selection. The module adds 0.15M parameters and is trained on top of the frozen backbone.

def learned_suppression(z, c):                                       # z: [K, D], c: [K, 3]
    h = mlp_emb(z)                                                   # 1
    for t in range(T):                                               # 2
        d = torch.cdist(c, c)                                        # 3
        A = (d < rho) & ~torch.eye(len(c), dtype=torch.bool)
        s = torch.sigmoid(mlp_rel[t](pairs(h, d)).squeeze(-1)) * A   # 4  s[i, j]: i suppresses j
        u = s.T @ h                                                  # 5  incoming suppression
        h = h + mlp_upd[t](torch.cat([h, u], dim=-1))                # 6
        c = c + mlp_mov[t](h)                                        # 7  relocate
    keep = torch.sigmoid(mlp_keep(h)).squeeze(-1)
    pi = keep * (1 - s.max(dim=0).values)                            # 9
    return c[pi > tau]                                               # 10

BibTeX

Copy@misc{bekar2026learnedsuppression3dkeypoint,
  title={Learned Suppression for 3D Keypoint Detection with a Graph-Transformer Backbone}, 
      author={Batuhan Arda Bekar and Can Sarı and Hüseyin Can Gülkan and Barış Özcan},
      year={2026},
      eprint={2605.15088},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2605.15088},
}