Registration is a Representation Problem

TL;DR

• Learn2Reg PSMAReg challenge task: Register a follow-up whole-body PSMA PET/CT scan to the baseline scan of the same patient, without introducing volume changes in the tumor.
• Instead of aligning volumes using raw intensities, we use representations of anatomy in the estimation: anatomix¹ features, MIND-SSC³ descriptors, TotalSegmentator⁴ organ labels and LION⁵ lesion masks.
• A GPU-based solver, FireANTs², is used to align these representations in two stages: dense features first, then organs and lesions.
• Volume constraints keep the volume of each lesion constant. Thus, the PET tumor biomarkers MTV and TLG do not change.
• We do not train anything on the challenge data. Despite that, this method placed first on both the validation and the test leaderboard.

Overview

The task. The PSMAReg task (Learn2Reg 2026, Task 1) is to register a follow-up whole-body PSMA PET/CT scan to the baseline scan of the same patient. The CT and the PET of one timepoint have the same grid. Thus, one deformation applies to both. The challenge ranks the methods on three criteria:

• registration accuracy: the Dice and HD95 of the CT organs and the PET lesions,
• PET biomarker preservation: the MTV and TLG errors,
• deformation regularity: %NDV.

The challenge gives biomarker preservation the same weight as accuracy.

The pipeline. The chart below lists every step from the four input volumes to the exported displacement field. The following sections show the intermediate result of each step.

Workshop talk (8-min)

This talk gives little context on the challenge, because the opening talk of the workshop covered its basics and details.

Watch on YouTube

Representations

We describe each scan with four representations, from the most detailed to the least detailed:

• anatomix¹ is a 3D biomedical foundation model, pretrained on highly variable synthetic data. It gives a 32-dimensional feature vector at each voxel. The features are approximately invariant to appearance and equivariant to spatial transformations.
• MIND-SSC³ is a hand-crafted descriptor of local self-similarity, with 12 channels.
• TotalSegmentator⁴ labels organs and bones on the CT.
• LION⁵ segments PSMA lesions on the PET.

These representations are used to drive the registration optimization instead of intensity images.

PET biomarker preservation

When registering across time in oncological settings, you want the registration to best align tissue but not change the lesion's volume itself. This is so that tumor volume change can be properly assessed. We therefore limit the local volume change in a mask M of possible lesions in the moving image. Trained PSMA lesion segmentors like LION have good specificity but can miss lesions. A lesion outside M does not have its volume preserved, so a missed lesion costs more than a false detection. We therefore build the mask M to favor recall.

We threshold the PET at SUV > 3. This threshold also selects organs with high physiological uptake. We subtract 16 of these organs, which TotalSegmentator segments, dilated by 10 mm. Then we take the union of the result and the LION prediction:

Stage 1: dense feature-based alignment

Stage 1 removes the large deformations. For each timepoint, we stack the CT, the PET, the 32 anatomix channels and the 12 MIND-SSC channels into a 46-channel feature volume F. The figure below shows F. We optimize

where \mathcal{L}_{\mathrm{LNCC}} is the squared local normalized cross-correlation. We calculate it for each channel separately, and only in the body mask. \mathcal{L}_{\mathrm{vol}} is the lesion-volume penalty, which we describe below. We use the compositive deformation model in FireANTs² on a 6×4×2×1 pyramid, with a step size of 2.0. We smooth the gradient with a Gaussian of standard deviation 2.0 voxels, and the deformation with a Gaussian of standard deviation 0.5 voxels.

Stage 2: segmentation-guided refinement

After Stage 1, the scans look well aligned. But shape-biased features can mix up adjacent abdominal organs. Stage 2 refines the alignment of each organ and lesion.

We warp the moving organ segmentations with \phi_1^*. In both volumes, we represent each structure by a distance transform. We truncate the transform at 20 mm and give it a soft boundary shell 2.0 voxels wide. We add the LION segmentation as one more structure, with weight 0.25. Then we optimize a residual deformation \phi_2:

where S are the multichannel segmentations and A are the anatomix features. We divide each hatted loss by its absolute value before Stage 2. Thus, each weight sets the size of its term at the start of Stage 2. We use \lambda_c=0.15, \lambda_r=0.25, \lambda_a=0.5 and \lambda_v=10. The Dice weight is then 1-\lambda_c-\lambda_r=0.60.

• The generalized Dice loss⁶ measures the overlap of the labels. It corrects for the difference in size between large and small organs. The one-sided chamfer loss measures the distance from the warped moving boundary to the fixed boundary.
• The region loss measures if the tissue in each fixed organ maps into the corresponding moving organ. For each organ, we sample a maximum of 4,000 points in the fixed organ. We map these points through the current deformation. A point in the moving organ has zero error. For a point outside the moving organ, the error is its distance to the organ in millimetres. When two organs overlap only a little, the Dice gradient is weak. But this loss still gives a direction for the correction.
• The anatomix term keeps the alignment of the tissue that the organ labels do not cover. The volume term prevents more change in lesion volume.

We optimize \phi_2 at one resolution for a maximum of 35 iterations, with a step size of 2.0. We smooth the gradient with a Gaussian of standard deviation 1.0 voxel, and the deformation with a Gaussian of standard deviation 0.5 voxels. Then we compose \phi^*=\phi_1^*\circ\phi_2^*.

Lesion volume preservation

MTV and TLG reach their ideal values when the deformation preserves the volume of the ground-truth lesions. We encourage this in both stages with a size-normalized penalty^7,8 on local volume change inside the lesion mask,

Volume projection

To further reduce any volume change introduced by the registration, we post-process the deformation after each stage. We use the volume-preserving projection of Hameeteman et al.⁹. The projection works on the fixed grid. Thus, it uses the lesion mask after the warp, M\circ\phi, dilated by one voxel. On M\circ\phi, let d=\log\det\nabla\phi be the local volume change.

We compose \phi with a correction v, which is a small displacement field. The field v is the gradient of a scalar potential p. The potential p is the solution of a Poisson equation, with -d on M\circ\phi as its source:

Because \log\det(I+\nabla v)\approx\nabla\!\cdot v=-d, each pass cancels the residual volume change to first order. We solve the Poisson equation with FFTs. On one NVIDIA RTX PRO 6000 GPU, the full projection after Stage 1 (a maximum of four passes) takes 0.4 to 0.8 s, and 0.44 s on average over 118 runs. We keep v only in a window around the lesions, limit it to 0.5 voxels and smooth it.

Applying the correction only when it's needed

If we apply this correction to all of M, it can decrease the anatomical alignment where the deformation already preserves volume. In these regions, it does not improve the biomarkers. Thus, we correct a lesion only if its volume changes. For each connected component C_i of M, let B_i=|C_i| be its voxel count. Let A_i=|C_i\circ\phi| be its voxel count after the warp. We do not change a component if

The other components first get four passes of the correction above. Each of these passes corrects all components together and is exact only to first order. Thus, a small volume error remains in each lesion. Then we correct the lesions one at a time, in a maximum of eight more passes. In each pass, we measure the voxel count A_i of each lesion after the warp. We correct the volume of the lesion until A_i is equal to B_i.

Last-mile details

Background fold correction

We calculate the image losses only in the patient. Thus, the losses do not constrain the deformation much in the air around the patient, and the deformation can fold there. On the training pairs that are above the %NDV threshold, 99.5 to 100% of the folded voxels are outside the body mask. Thus, we repair folds¹⁰ only outside the body mask of the fixed image. We dilate this mask by a safety margin of 5 mm.

In each of six passes, we expand the detected region by two voxels. In this region, we replace the deformation with a feathered blend of the original field and a smoothed field. The smoothing is Gaussian, with a standard deviation of 2.0 voxels. The blend weight is zero in the body.

(This entire procedure could likely have been avoided if we just padded the volumes with empty space. Oh well.)

Meeting resource requirements

The test phase of the challenge required each pair to be registered in about 90 s on one 80 GB H100 GPU. A direct serial implementation of our pipeline takes 812 s and 83 GB. The final pipeline takes 79 s and 72 GB. The subsections below give the four most important changes.

Fewer sliding windows

anatomix runs on overlapping windows of 128 voxels. The number of windows increases with the cube of the inverse stride. We decrease the overlap from 0.6 to 0.25. This decreases the number of windows on each pyramid level of 288³ voxels from 125 to 27. It also decreases the feature extraction time from 78 s to approximately 20 s for each pair.

We skip the windows that contain no body. We also accumulate the windows in FP16. This decreases the accumulator from 43.8 to 28.0 GB. The cosine similarity to the FP32 result is 1.000000.

A body mask without a network

The image losses need a body mask, and so does the step that skips empty windows. The body model of TotalSegmentator takes 29 to 31 s. We use these steps instead:

1. Threshold the CT at −500 HU.
2. Erode by one voxel, to separate the patient from the scanner table.
3. Keep the largest connected component.
4. Fill the holes in each slice.
5. Close by two voxels.

These steps take 0.7 s. The result agrees with TotalSegmentator at a Dice of 0.993 to 0.995.

Faster organ labels

The full-resolution model of TotalSegmentator takes 158 s for each timepoint. This is more than the time limit for a full pair. Its fast 3 mm model takes 29.5 s, but it is less accurate, mostly on the ribs and vertebrae. On both timepoints, we run three models at the same time:

• the fast model,
• a full-resolution model on a tight crop, for the 52 rib and vertebra classes,
• a full-resolution model on a tight crop, for the other 65 classes.

We merge their outputs into one map. The two crop models take 37 and 46 s for each timepoint. On every pair, the merged map agrees with the challenge labels better than the unrestricted full-resolution model. The lesion mask needs only 16 dilated organs. Thus, it uses the faster 6 mm model.

Wait points

All segmentation jobs start with the pipeline and use the same GPU. The GPU steps wait for a job only at the first step that needs its output. The volume penalty of Stage 1 needs the lesion mask approximately 32 s after the start. For this reason, the mask subtracts organs from the 6 mm model.

Stage 2 needs the merged organ maps approximately 60 s after the start. It also needs the LION mask of the fixed image for its lesion channel. LION segments both PET volumes in one process. Thus, the second volume adds approximately 3 s, and not the time of a second model start-up.

Results

The challenge organizers scored the methods on the PSMAReg test set. Our method is first. We do not train or fine-tune a model on the challenge data. We use the 202 training pairs only to select the hyperparameters.

The top three methods on the test set, as the organizers ranked them. The organizers did not give standard deviations.

Ablations, pair by pair

We ran every configuration of the ablation on 65 pairs with ground truth, from 50 training subjects. We did not use these pairs to tune the method. Thus, each step is a paired comparison. Some steps improve all 65 pairs:

• The addition of anatomix and MIND-SSC improves CT Dice and HD95 on every pair.
• Stage 2 also improves CT Dice and HD95 on every pair (compare S1 · + volume projection with S2 · + anatomix).
• The anatomix term decreases the CT intensity error on every pair, compared with labels alone.
• The final projection and smoothing decrease %NDV on every pair.

Limitations

There are plenty! For example, we use a single diffeomorphic model to represent several distinct motion models across the body (your arms and head move in a different way than your chest and abdomen move). This leads to the estimation being underconstrained for articulated limbs and sliding organs, and overconstrained where topological correspondence does not hold. E.g., when the arms are crossed in the baseline and above the head in the follow-up, finding a fluid/elastic warp that aligns them isn't happening. Bone-specific polyrigid models are one direction for future work.

Citation

@inproceedings{dey2026registration,
  title     = {Registration is a Representation Problem: Solving Whole-Body {PET/CT} Registration},
  author    = {Dey, Neel and Dalca, Adrian V.},
  booktitle = {Learn2Reg Challenge, MICCAI},
  year      = {2026}
}

Acknowledgments. This work was supported by NIH grants R01 EB033773 and S10 OD038222 and an MGB Neuroscience Institute Transformative Scholar Award.

References

[1] N. Dey, B. Billot, H. E. Wong, C. J. Wang, M. Ren, P. E. Grant, A. V. Dalca, P. Golland. Learning general-purpose biomedical volume representations using randomized synthesis. ICLR 2025. code

[2] R. Jena, P. Chaudhari, J. C. Gee. FireANTs: Adaptive Riemannian optimization for multi-scale diffeomorphic registration. Nature Communications, 2024. code

[3] M. P. Heinrich, M. Jenkinson, B. W. Papież, M. Brady, J. A. Schnabel. Towards realtime multimodal fusion for image-guided interventions using self-similarities. MICCAI 2013.

[4] J. Wasserthal et al. TotalSegmentator: robust segmentation of 104 anatomic structures in CT images. Radiology: Artificial Intelligence, 2023.

[5] L. K. Shiyam Sundar, M. Pires, S. Gutschmayer. LION: lionz-v.0.9.1, 2024. code

[6] C. H. Sudre, W. Li, T. Vercauteren, S. Ourselin, M. J. Cardoso. Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations. DLMIA 2017.

[7] T. Rohlfing, C. R. Maurer, D. A. Bluemke, M. A. Jacobs. Volume-preserving nonrigid registration of MR breast images using free-form deformation with an incompressibility constraint. IEEE TMI, 2003.

[8] M. Staring, S. Klein, J. P. W. Pluim. A rigidity penalty term for nonrigid registration. Medical Physics, 2007.

[9] R. Hameeteman, J. F. Veenland, W. J. Niessen. Volume preserving image registration via a post-processing stage. SPIE Medical Imaging 2008.

[10] B. Karaçali, C. Davatzikos. Estimating topology preserving and smooth displacement fields. IEEE TMI, 2004.