PRNet: Self-Supervised Learning for Partial-to-Partial Registration

In plain words

Registration means finding the rotation and shift that line up one 3D scan with another. This is hard when each scan shows only part of the object, as when a depth camera sees one side of a chair and a second camera sees another side. Earlier methods assume the two scans contain the same points, or that the whole shape is visible.

PRNet first learns which points each scan shows that the other scan also shows. It then pairs those points up, computes the best rigid motion for the pairs, moves the scan, and repeats. The network learns all of this from synthetic pairs, where the true motion is known, so no hand-labelled keypoints are needed. Repeating the step lets a rough first guess become a finer one.

TL;DR

• Input: two point clouds, each a partial view of the same object, with an unknown rigid motion between them.
• Each pass selects keypoints in both clouds, learns a soft one-to-one matching between them, and solves for the rigid motion in closed form (SVD). The source is moved and the pass is repeated; the code runs three passes §3.2 main.py.
• Trained self-supervised on ModelNet40 with synthetic motions; the reported partial-to-partial numbers are author-reported on synthetic data only.
• Real scans (Stanford Bunny) needed inference-time fine-tuning on 100 synthetic examples per scan §4.2.

Problem

The ICP formulation assumes both clouds have the same number of points (M = N) §3.1. For partial views, usually N ≠ M and only subsets of the two clouds match §3.2. Earlier learned registration (DCP, PointNetLK) either assumes the whole shape is visible or applies one shot with no refinement §1 §2. The paper also argues that a fixed-sharpness soft matching is either too blurry for a precise final alignment or not differentiable when made sharp §3.2.

Method

Notation. The paper writes the source as X and the target as Y, with rigid motion [R_XY, t_XY] that maps X onto Y §3.1. After p passes, X^p is the source after p updates, and [R^p, t^p] is the motion predicted in pass p §3.2 eq. 7 eq. 8.

Preliminaries. For a fixed point-to-point mapping, the best rigid motion is closed-form: centre both sets, form the cross-covariance H, take its SVD, and set R = V U^T and t = -R x̄ + ȳ §3.1 eq. 2. ICP alternates this with a nearest-neighbour mapping. DCP replaces the mapping with a softmax over feature dot products, in one shot §3.1 eq. 3 eq. 4.

PRNet pass. One pass §3.2:

1. Embed each cloud with a DGCNN; the paper uses five dynamic EdgeConv layers with filter counts (64, 64, 128, 256, 512) §4.1 [supp. Actor-Critic Closest Point].
2. Contextualise the two embeddings jointly with one Transformer encoder and one decoder, 4 heads, LayerNorm, no dropout [supp. Actor-Critic Closest Point].
3. Keypoints: take the k points whose embeddings have the largest L2 norm, separately in each cloud eq. 9. The paper states no supervision for keypoints §4.3.
4. Predict a temperature λ from the average-pooled embeddings of both clouds, using a small value network §3.2 [supp. Actor-Critic Closest Point].
5. Sample a keypoint-to-keypoint matching with Gumbel-Softmax at temperature λ; the forward pass is a hard one-hot, and gradients use a straight-through estimator eq. 10 eq. 11.
6. Solve the rigid motion from the matched keypoints with the closed form of eq. 2 eq. 8.
7. Apply the motion to the source and repeat eq. 7.

The keypoint embeddings pass through the action head (Gumbel-Softmax plus SVD). The value head outputs λ, so the step is described as an actor-critic structure §3.2 §5.

# One PRNet pass  [§3.2] [fig. 1]
# Input:  source (B, 3, N), target (B, 3, M) point clouds
# Output: rotation (B, 3, 3), translation (B, 3)

def prnet_pass(source, target, k):
    # Per-point embeddings from a DGCNN, one per cloud  [§3.2] [supp. ACP]
    source_feat = dgcnn(source)                  # (B, d, N)
    target_feat = dgcnn(target)                  # (B, d, M)

    # Jointly contextualise both clouds with one Transformer block  [supp. ACP]
    source_feat, target_feat = transformer(source_feat, target_feat)

    # Keypoints: the k points with the largest feature L2 norm  [eq. 9]
    source_keys, source_keys_feat = top_k_by_norm(source, source_feat, k)
    target_keys, target_keys_feat = top_k_by_norm(target, target_feat, k)

    # Temperature from average-pooled features of both clouds  [eq. 11]
    temperature = value_head(source_keys_feat, target_keys_feat)  # (B, 1)

    # Scores between keypoint features, one row per source keypoint  [eq. 10]
    logits = transpose(source_keys_feat) @ target_keys_feat  # (B, k, k)

    # Sharp, approximately differentiable matching  [eq. 10] [eq. 11]
    matching = gumbel_softmax(logits, temperature, hard=True)

    # Closed-form rigid fit of the matched keypoints  [eq. 2] [eq. 8]
    rotation, translation = svd_fit(source_keys, target_keys, matching)
    return rotation, translation

The closed-form fit used inside each pass §3.1 eq. 2 eq. 8:

# Rigid fit from matched keypoints  [§3.1] [eq. 2]
# Input:  keys_src (B, 3, k), keys_tgt (B, 3, k), matching (B, k, k)
# Output: rotation (B, 3, 3), translation (B, 3)

def svd_fit(keys_src, keys_tgt, matching):
    # Target point assigned to each source keypoint  [eq. 10]
    matched_tgt = keys_tgt @ transpose(matching)              # (B, 3, k)
    # Centre each set on its own centroid  [eq. 2]
    src_mean = mean(keys_src, dim=2, keepdim=True)            # (B, 3, 1)
    tgt_mean = mean(matched_tgt, dim=2, keepdim=True)         # (B, 3, 1)
    src_centered = keys_src - src_mean
    tgt_centered = matched_tgt - tgt_mean
    # Cross-covariance H = U S V^T  [eq. 2]
    H = src_centered @ transpose(tgt_centered)                # (B, 3, 3)
    U, S, V = svd(H)
    # Rotation R = V U^T  [eq. 2] [eq. 8]
    rotation = V @ transpose(U)
    # Translation t = -R x_mean + y_mean  [eq. 2] [eq. 8]
    translation = -rotation @ src_mean + tgt_mean
    return rotation, translation

Training runs P = 3 passes on each pair, moving the source after each pass, and sums a discounted loss §3.2 eq. 12 [supp. Training Protocol]. The local ground truth for pass p is the true motion, composed with the inverse of the motion predicted so far eq. 5 eq. 6:

# Discounted training loss over P passes  [§3.2] [eq. 12]
# Input:  source, target, R_gt, t_gt (true motion source -> target)
# Output: scalar loss

def training_loss(source, target, R_gt, t_gt):
    loss = 0
    # Motion accumulated over earlier passes  [eq. 6]
    R_before = identity()
    t_before = zeros()
    for p in range(P):
        # Forward and reverse motion plus the pooled features  [§3.2]
        R_p, t_p, R_back, t_back, psi_src, psi_tgt = both_ways(source, target)
        # Local ground truth for this pass  [eq. 5]
        R_local = R_gt @ transpose(R_before)
        t_local = t_gt - R_local @ t_before
        # Rigid motion loss  [eq. 13]
        loss_motion = sq_norm(transpose(R_p) @ R_local - identity())
        loss_motion = loss_motion + sq_norm(t_p - t_local)
        # Cycle consistency between the two directions  [eq. 14]
        loss_cycle = sq_norm(R_p @ R_back - identity())
        loss_cycle = loss_cycle + sq_norm(t_p - t_back)
        # Global feature alignment  [eq. 15]
        loss_feat = norm(psi_src - psi_tgt)
        # Discounted sum, weights alpha and beta  [eq. 12]
        weight = gamma ** p
        loss = loss + weight * (loss_motion
                                + alpha * loss_cycle
                                + beta * loss_feat)
        # Move the source for the next pass  [eq. 7]
        source = R_p @ source + t_p
        # Update the accumulated motion  [eq. 6]
        R_before = R_p @ R_before
        t_before = R_p @ t_before + t_p
    return loss

The paper's local ground truth in eq. 5 is written with the product R^{1..p} defined in eq. 6; the code's loss differs, see Paper vs code.

Hyperparameters: λ-prediction, temperature handling and the value head are described in the paper but their range is not stated §3.2; the code clamps them (see below).

Hyperparameters

Results

Author-reported. Setup: ModelNet40 partial-to-partial, 1024 input points, 768 kept per cloud, evaluated on held-out objects §4.1; inference timing on an Intel 16-core CPU and a GTX 1080 Ti [supp. Efficiency].

Rotation errors are on Euler angles in degrees §4.1; the units of the translation errors are not stated. Bold values are the paper's.

Table 1: unseen point clouds, ModelNet40 test split tab. 1

Table 2: unseen categories tab. 2 (PRNet and DCP trained on the first 20 ModelNet40 categories; PRNet (Ours*) trained on ShapeNetCore)

Table 3: unseen point clouds with Gaussian noise tab. 3

Table 4: transfer, linear SVM on DGCNN features, ModelNet40 accuracy tab. 4 §4.4

The paper reports classification accuracy, in percent. The z/z column is the setting trained and tested with azimuthal rotations; SO(3)/SO(3) with random rotations in both.

Supervised baselines in this table use labels; PRNet's features are learned without them. The paper calls PRNet "comparable" and does not claim to beat them §4.4.

Real data §4.2 fig. 2 fig. 3. On the Stanford Bunny (10 real scans), the ModelNet40-trained model was fine-tuned for 10 epochs at lr 0.0001 on 100 synthetic examples per scan. The paper gives only qualitative figures for real scans, with no numeric table.

Choice of λ [supp. Choices of λ].

Ablations [supp. ablation tables]. Only MAE and R² are reported. Values as printed. Keypoint selection [supp. (a)]:

Number of keypoints k [supp. (d)]:

Full point clouds, 768 points per cloud by farthest-point sampling [supp. (b)]:

Discount factor, printed as "λ" in the caption [supp. (c)]:

Missing data ratio at test [supp. (e)]:

Data noise [supp. (f)]:

Inference time, seconds [supp. Efficiency]. Per pair, as printed.

The paper states that PRNet is about as fast as PointNetLK and slower than DCP [supp. Efficiency].

Headline (author-reported). On the ModelNet40 partial-to-partial split PRNet has the best value in all eight Table 1 columns, with rotation MSE 10.235 against 45.005 for DCP-v2 tab. 1. It also has the best value in all eight columns of the noise table tab. 3. On unseen categories it does not: FGR has lower MAE on rotation (1.952 vs 2.329) and on translation (0.007 vs 0.015) tab. 2.

Paper vs code

Official code: https://github.com/WangYueFt/prnet, linked from the paper's contribution list §1; its owner handle matches author Yue Wang. Cloned at commit 82d466a91c45f45acb57c839acdd7369d3ba10b8 (2022-09-21). The repository contains main.py, model.py, data.py, util.py and a README; there is no LICENSE file. The README states "MIT License" README.md; that is recorded in code_license and is unverified. I read the code and did not run it.

Differences between paper and code, each with both citations:

• Loss target (rigid motion). The paper supervises each pass with the local ground truth of eq. 5, which is built from the previous passes' motion eq. 6. The code compares the accumulated pose after each pass with the global ground truth: R_acc^T R_gt and the accumulated translation against t_gt model.py model.py.
• Mean instead of squared norm. The paper writes squared Frobenius and squared Euclidean norms eq. 13 eq. 14. The code uses F.mse_loss, which averages over the 9 rotation entries and the 3 translation entries model.py model.py. The relative weights of the motion, cycle and feature terms therefore differ from the written loss by constant factors.
• Cycle translation term. The paper writes the translation part of eq. 14 as a difference, t_XY - t_YX. The code computes mse(t_ab, -t_ba), which is ||t_ab + t_ba||² model.py. The code's form matches the physical inverse-motion relation; the paper's written form does not.
• Feature alignment. The paper uses the L2 norm of Ψ_X - Ψ_Y eq. 15. The code takes the element-wise absolute difference of the pooled features, then the mean over channels model.py model.py.
• Pooling for the value head. The paper averages Φ over all points §3.2. The code pools the keypoint embeddings, because keypoint selection runs before the temperature net model.py.
• Temperature input. The paper says λ is predicted from Ψ_X and Ψ_Y §3.2. The code's value head gets only |Ψ_X - Ψ_Y| model.py. The code clamps λ to [1/100, 100] model.py with temp_factor = 100 main.py. Neither range is in the paper.
• Matching logits. The paper's logits are Φ_Y Φ_X^T with no scaling eq. 10 eq. 11. The code divides by sqrt(d_k), with d_k = 512 by default model.py main.py.
• Gradient of the matching. The paper says the backward pass uses the DCP softmax of eq. 4, ignoring the one-hot and argmax §3.2. The code uses PyTorch's F.gumbel_softmax(..., hard=True) model.py, whose straight-through gradient flows through the Gumbel soft sample at temperature λ.
• Embedding width. The paper says the embedding dimension is 1024 §4.1. The code default is 512 main.py; the feed-forward width is 1024 main.py.
• DGCNN depth. The paper says five dynamic EdgeConv layers §4.1. The code has four EdgeConv layers, with max-pooled neighbourhood features concatenated and a 1×1 convolution to 512 channels model.py.
• DGCNN neighbours. The paper does not state k for the kNN graph. The code default is 20 model.py.
• Feed-forward activation. The code uses LeakyReLU(0.2) in the Transformer feed-forward block model.py. The paper does not state the activation.
• SVD reflection guard. The paper's closed form eq. 2 has no reflection check. The code adds one: if det(R) < 0, it flips the smallest singular direction model.py.
• Partial crop. The paper says the point at which the crop is centred is placed randomly in space, and the 768 nearest neighbours are kept §4.1. The code places the probe at a random point 500 units away along a diagonal (±500 on each axis plus U[0,1) noise) and keeps the 768 points nearest to it data.py. The function name farthest_subsample_points does not describe farthest-point sampling. The same probe coordinate is used in both clouds' own frames data.py.
• Sampling of clouds. The paper's stated sampling step is farthest-point, taken over the CAD surfaces §4.1. The code loads the pre-sampled ModelNet40 HDF5 files and takes the first n_points rows data.py data.py. No farthest-point sampling appears in the repository.
• Noise. The paper writes N(0, 0.01) clipped to [-0.05, 0.05] §4.1. The code uses standard deviation 0.01 and clips the noise to ±0.05 data.py. The ablation labels its noise as N(0, σ²), which suggests σ is the standard deviation in both places [supp. (f)].
• Noise in training. The paper does not say whether training used noise. The code applies the same gaussian_noise flag to both the training and test loaders main.py, and the README's noise command sets it.
• Unseen-category comment. The comment in data.py says training is on the last 20 categories data.py. The code trains on labels below 20 and tests on 20 and above data.py, which matches the paper §4.1.
• Learning-rate step. The scheduler is stepped at the start of each epoch main.py. The 30/60/80 milestones therefore take effect from the 30th, 60th and 80th epoch in 1-based counting.
• Metric convention. The paper gives the rotation error in degrees without an Euler convention §4.1. The code converts the predicted matrix to Euler angles in zyx order, in degrees, and compares with the ground-truth [z, y, x] angles util.py model.py.
• Head choice. The code's default head is the SVD head main.py. It also has a quaternion MLP head model.py that the paper does not describe.
• Keypoints at test. The paper's keypoint count for test is not stated. The code uses the training default of 512 main.py.
• Real-data code. The repository has no fine-tuning or Bunny script. main.py accepts --model_path only main.py.

No discrepancy found for: the rotation sampling range (0° to 45°) data.py; the closed-form R = V U^T with centred means model.py; training P = 3 with γ^{p-1} weighting model.py; the Adam optimiser with the 30/60/80 schedule main.py; the 768-of-1024 crop size main.py.

Pitfalls

• Test-time matching is random. F.gumbel_softmax is called inside SVDHead.forward with no check of self.training model.py. Test predictions therefore depend on the torch RNG state. main.py seeds torch once at start-up main.py; the test loader seeds NumPy per item only data.py. The paper does not say whether it averaged over samples or fixed a seed.
• The temperature range and the clamp are not in the paper. λ is clamped to [0.01, 100] in code model.py. The paper states no range for λ, so as written the value head is unbounded.
• Partial overlap is not controlled. Each cloud keeps 768 of 1024 points around a probe 500 units away, so the shared fraction depends on the probe's direction and the object data.py. The ablation tests 50% and 75% missing-data settings [supp. (e)].
• The keypoint loss is absent. The paper says keypoints are learned without supervision §4.3. The only signal is the registration and cycle losses §3.2. Keypoint quality is shown in figures only fig. 4 fig. 5.
• Matchings are not one-to-one. The conclusion says the computed matchings are not bijective, with examples on cars and chairs §5 fig. 6. A proposed fix (Gumbel-Sinkhorn) is future work §5.
• Gradient estimator. The straight-through estimator is biased §3.2. The paper gives no variance or convergence comparison for its gradient choice.
• Notation collisions. The discount factor is γ in the text and "λ" in the ablation caption [supp. (c)]. λ is also the Gumbel temperature eq. 11. The ablation text also has "small large discount factors", which is an error [supp. (c)].
• Training randomness. Gumbel samples make each training step random, and each pass uses the previous pass's sampled matching. The paper reports no seed-variance numbers.
• Inference cost. The closed-form SVD runs per sample, in a Python loop on the CPU, after moving the matrices off the GPU model.py. PRNet is slower than DCP in the paper's timing [supp. Efficiency].
• Real scans need fine-tuning. The paper says PRNet needs inference-time fine-tuning on real scans, and calls this a cost of the method §4.2 §5. Fine-tuning code is not in the repository main.py.
• Limited real-data evidence. The real-data experiment is qualitative; there is no numeric real-scan metric §4.2.
• Synthetic-only metrics. All quantitative results use synthetic motions on ModelNet40 and ShapeNetCore §4.1. Generalisation is shown only by the Ours* row and the qualitative bunny and Stanford examples §4.1 §4.2.
• Baselines are the authors' runs. DCP-v2 was trained from the authors' released code "as they suggest" [supp. Actor-Critic Closest Point]. The other baselines' settings are not described beyond the paper's text §4.1.
• Sign of the probe. The code's probe alternates between +500 and -500 along each axis with equal probability data.py. This is a per-sample random choice, not a fixed side.
• Batch size and seed are not in the paper; the code uses 36 / 12 and seed 1234 main.py main.py.

Lineage

• Builds on: ICP (classic iterative registration) §3.1; Deep Closest Point (Wang and Solomon, ICCV 2019), whose notation and softmax matching this paper uses §3.1; PointNetLK (Goforth et al., CVPR 2019) as a registration baseline §2. No arXiv ids for these are in the source; they are named, not linked.
• Uses: DGCNN (dynamic graph CNN) §3.2, Transformer (Vaswani et al. 2017) §3.1, Gumbel-Softmax (Jang et al. 2017) with a straight-through estimator §3.2.
• Compared against: ICP, Go-ICP, FGR, PointNetLK, DCP-v2 tab. 1 tab. 2 tab. 3; PointNet, PointNet++, DGCNN, VoxNet, SubVolSup, MVCNN, RotationNet, Spherical CNNs for classification tab. 4.
• Superseded by: not stated in the source; this card does not record later work.