PaperlayerView exact revision on arXiv

Source-linked AI summary

Two-Parameter Flow Map Learning for Continuous-Time Diffeomorphic Image Registration

Mohammadjavad Matinkia, Nilanjan Ray

arXiv:2609.10789v1cs.CV

TL;DR

Diffeomorphic registration must balance anatomically consistent topology with alignment accuracy, while non-autonomous flow models introduce expressiveness without avoiding integration-related training burdens. TPFM-DIR directly learns two-parameter flow maps with cocycle consistency, improving registration across nine datasets while preserving topology.

  • Problem

    Learning-based diffeomorphic registration is limited by stationary-velocity assumptions, while non-autonomous methods rely on numerical integration to enforce flow structure.

  • Method

    TPFM-DIR directly models the two-parameter solution of a non-autonomous ODE and uses cocycle regularization to learn temporally consistent flow maps without numerical integration during training.

  • Results

    TPFM-DIR consistently improves alignment across nine MRI, CT, and ultrasound datasets while preserving topology, including a 2.1% average Dice improvement on neuroimaging MRI.

  • Takeaways & Limitations

    The framework is architecture-agnostic and can be integrated into existing registration backbones while narrowing the performance gap between diffeomorphic and unconstrained deformable methods.

  • Takeaways & Limitations

    Uniform sampling underrepresents large time transitions, so inference partitions the interval into N subintervals and composes locally predicted maps.

Abstract

from arXiv · show

Diffeomorphic image registration is central to medical image analysis, enabling anatomically consistent alignment across subjects. Most learning-based diffeomorphic methods model autonomous ODEs(ordinary differential equations) by parameterizing a stationary velocity field and recovering deformations via scaling-and-squaring. While non-autonomous ODEs with time-dependent velocities increase expressiveness, existing approaches rely on numerical integration to implicitly enforce flow structure that entangles model expressiveness with discretization accuracy. We propose a framework to directly learn the continuous-time solution of a non-autonomous ODE formulated as a two-parameterflow map. By enforcing cocycle consistency, a fundamental structural property of time-varying flows, we learn the flow maps without time discretization and velocity integration during training. The framework recovers diffeomorphic mappings at inference using a small number of compositions. Our proposed framework seamlessly incorporates standard registration backbones and improves alignment accuracy consistently across nine datasets while preserving diffeomorphic structure. Notably, the proposed method achieves an average Dice improvement of 2.1% on brain MRI benchmarks, a 12% TRE reduction on lung CT, and a 2.6% Dice gain on cardiac MRI and ultrasound datasets.

1 Introduction

TPFM-DIR directly learns two-parameter flow maps for non-autonomous diffeomorphic registration, using cocycle consistency to model time-dependent deformations without relying on multiple handcrafted smoothness penalties. It consistently improves alignment while preserving topology across diverse datasets and integrates with standard registration backbones.

  • Cocycle-based regularization removes explicit numerical integration during training while promoting temporal consistency and invertibility.It replaces reliance on multiple handcrafted smoothness penalties with structural constraints from valid non-autonomous flows.
  • Global transformations are formed by composing a small number of predicted short-time maps, improving deformation regularity while maintaining computational efficiency.The approach captures complex time-dependent deformations through locally predicted flow maps.
  • TPFM-DIR matches or surpasses strong deformable baselines while preserving topology, narrowing the performance gap between diffeomorphic and deformable registration.The framework also improves alignment accuracy and topology preservation across diverse backbone architectures.
  • TPFM-DIR directly models the two-parameter solution of a non-autonomous ODE at the flow-map level rather than learning its generating velocity field.This formulation structurally models time-dependent deformations.
  • The method is validated on nine benchmark datasets spanning 2D and 3D MRI, CT, and ultrasound, with consistent improvements over strong deformable and diffeomorphic methods.The overview includes continuous registration trajectories with time-varying velocity norms and temporally smooth incremental differences.

2 Related Works

Prior work contrasts accuracy-focused deformable registration with structurally constrained diffeomorphic registration and distinguishes stationary-velocity flows from more expressive time-varying formulations. Non-autonomous flow-based models provide a broader deformation framework than autonomous models, whose stationarity assumption restricts deformation dynamics.

  • Deformable methods prioritize alignment accuracy with dense displacement fields but generally lack explicit structural constraints, which may allow folding under large deformations.Recent gains include increasingly expressive architectures such as vision transformers and correlation- or structure-aware modules.
  • Diffeomorphic registration seeks smooth, invertible transformations to ensure anatomical plausibility.This structural guarantee distinguishes it from unconstrained deformable registration.
  • Autonomous flow formulations use stationary velocity fields and scaling-and-squaring, whereas non-autonomous models allow time-varying velocities and a more general deformation framework.The related-work discussion places LDDMM, NODEO, and R2Net among non-autonomous approaches.
  • Figure 4 compares visual registration results on OASIS and LungCT using Dice or TRE and the percentage of voxels with negative Jacobian determinants.The displayed metrics pair alignment quality with folding-related topology information.

3 Proposed Method

TPFM-DIR directly learns two-parameter flow maps for non-autonomous ODEs, using cocycle consistency to avoid numerical integration during training while preserving valid diffeomorphic structure.

  • Two-Parameter Flow Maps: The non-autonomous formulation uses time-dependent velocities, offering greater expressiveness but coupling deformation structure to time discretization and solver accuracy.The proposed approach addresses this training-time dependence by learning the flow map itself.
  • Two-Parameter Flow Maps: TPFM-DIR models the two-parameter flow map directly, replacing velocity-field parameterization and numerical integration with structural flow modeling.The map represents transformations from time s to t and is constrained by identity and cocycle properties.
  • Flow-Map Parameterization: The flow map is parameterized as ϕ_s,t = Id + (t - s)f_θ(s,t; I_m,I_f), with a time-conditioned neural network and shared spatial coordinates.The parameterization enforces identity at equal times, while cocycle regularization supplies the additional structural constraint needed for valid flow maps.
  • Cocycle Regularization: Cocycle regularization enforces temporal consistency and, under standard regularity assumptions, yields a two-parameter diffeomorphism solving the non-autonomous ODE.The method uses anchored cocycle constraints as a computationally efficient alternative to enforcing every arbitrary time triplet.
  • Inference: Four composed local maps are sufficient in experiments to construct an accurate diffeomorphic deformation at inference.Composing short-time maps improves stability and reduces the risk of foldings compared with directly predicting the full interval.

4 Experimental Results

TPFM-DIR is evaluated across nine datasets and multiple modalities using overlap, surface, landmark, intensity, and topology metrics, achieving strong alignment while preserving diffeomorphic structure.

  • Datasets and Metrics: The evaluation spans nine datasets across brain MRI, thoracic and abdominal CT, cardiac MRI, and cardiac ultrasound, covering inter-subject and intra-subject registration.Metrics include Dice, HD95, ASSD, TRE, SSIM, negative-Jacobian percentage, and SDLogJ.
  • Baselines: Comparisons include diffeomorphic, deformable, classical, learning-based non-autonomous, and instance-optimization baselines across the datasets.The baseline set includes R2Net, NODEO, LDDMM, TransMorph-diff, and HViT among others.
  • Brain MRI: 2.1% average Dice improvement over the next-best method is reported on five neuroimaging 3D MRI datasets while maintaining near-zero folding ratios.Tables 1 and 2 cover OASIS, IXI, LPBA40, Mindboggle101, and CANDI.
  • Thoracic and Abdominal CT: 17% and 12% TRE reductions relative to GradICON and CorrMLP, respectively, are reported on LungCT, with zero folding on that dataset.LungCT and AbdomenCT results are summarized in Table 3.
  • Cardiac MRI and Ultrasound: TPFM-DIR improves CAMUS Dice by 3% while preserving topology with zero folding across the cardiac ultrasound evaluation.ACDC and CAMUS results are summarized in Table 4, with qualitative comparisons shown in Fig. 5.

5 Ablation Studies

The ablations show that TPFM-DIR improves alignment and topology preservation across backbones, while its regularization and inference compositions provide practical accuracy–regularity trade-offs. The framework also retains favorable computational characteristics by avoiding numerical integration during optimization.

  • 5.1 Effect of Regularization Factor: λ = 10 provides a favorable accuracy–topology trade-off, while λ = 0 substantially degrades both performance and topology.Increasing λ progressively removes negative-Jacobian voxels but moderately decreases the performance metric; reducing λ increases folding without improving alignment.
  • 5.2 Effect of Inference Compositions: Four inference compositions are sufficient to approach zero |J|<0%, with further compositions offering little additional topology benefit.Progressive composition improves smoothness and reduces folding with only minor alignment degradation, while benefits saturate beyond four steps.
  • 5.3 The Choice of Architecture: TPFM-DIR is a framework rather than a fixed architecture and adapts existing backbones through temporal conditioning in their decoders.The study modifies CorrMLP and TransMorph using context addition to incorporate time information.
  • 5.3 The Choice of Architecture: TPFM-DIR improves Dice while reducing folding across CorrMLP and TransMorph, outperforming scaling-and-squaring variants in alignment with comparable or better topology preservation.The comparison includes original deformable models, conventional diffeomorphic variants, and TPFM-DIR across OASIS, IXI, Mindboggle, and AbdomenCT.
  • 5.4 Computational Analysis: TPFM-DIR achieves the fastest inference while maintaining low inference memory and faster training and inference than methods using non-autonomous ODE integration.The computational comparison evaluates parameters, runtime, and GPU memory on OASIS using an NVIDIA RTX3090 GPU.

6 Conclusion

TPFM-DIR directly parameterizes non-autonomous flow maps and enforces anchored cocycle consistency to obtain diffeomorphic trajectories without numerical integration during optimization. The method uses composed short-time maps at inference to improve stability and preserve diffeomorphic structure.

  • TPFM-DIR directly parameterizes the flow map of a non-autonomous differential equation and learns temporally consistent diffeomorphic trajectories without numerical integration during optimization.Its anchored cocycle constraint provides the structural consistency required for valid two-parameter flows.
  • Under regularity assumptions, identity and anchored cocycle conditions guarantee that the learned map is a two-parameter diffeomorphism solving a non-autonomous ODE.The proof constructs local maps for sufficiently small increments, composes them into a global diffeomorphism, and derives the full cocycle property.
  • The construction uses partitions with increments smaller than 1/L_f, making each short-time map locally invertible before composing them into a global transformation.Invertibility follows from the Jacobian bound and the Inverse Function Theorem.
  • The instantaneous velocity can be recovered analytically from the flow-map parameterization as v(x,t) = f_θ(t,t,x).Evaluating the time derivative at s = t removes the additional derivative term.
  • Because uniformly sampled time pairs favor smaller increments, inference composes multiple short-time maps; N = 4 provides a favorable efficiency–accuracy balance.This strategy relies on better-trained local transitions and promotes stability and diffeomorphic preservation.

C Datasets and Preprocessing

The evaluation spans nine datasets covering brain MRI, thoracic and abdominal CT, cardiac MRI, and cardiac ultrasound. Dataset-specific splits, pairings, preprocessing, and evaluation annotations support inter-subject and intra-subject registration experiments.

  • Nine datasets cover multiple modalities, anatomical regions, and spatial dimensions, including brain MRI, LungCT, AbdomenCT, ACDC cardiac MRI, and CAMUS ultrasound.The collection includes five brain MRI datasets, two CT datasets, and two 2D cardiac datasets.
  • Dataset statistics and train, validation, and test splits are summarized in Table 7, with split totals matching each dataset’s available scans.
  • Registration pairs are sampled independently across splits for inter-subject datasets, while LungCT and cardiac datasets use defined respiratory or cardiac phase pairs.LungCT uses end-inhalation and end-exhalation scans; ACDC and CAMUS use end-systolic and end-diastolic phases.
  • Preprocessing standardizes spatial resolution and intensity ranges through cropping, resampling, Hounsfield-unit clipping, and min-max normalization.Brain scans are placed in MNI152 space, while cardiac images are resized or cropped to 128 × 128.

D Implementation Details

TPFM-DIR uses a time-embedded U-Net whose decoder receives the two time samples through FiLM modulation. The implementation combines sinusoidal time embeddings, gated-attention skip connections, and dimension-specific deformation outputs.

  • TPFM-DIR uses a time-embedded U-Net with encoder dimensions [2, 16, 32, 64, 128, 512] and decoder dimensions [512, 128, 64, 32, 16, 3].The final decoder dimension is 2 for 2D deformations.
  • The architecture uses skip connections with gated attention modules, while the input concatenates moving and fixed images.
  • Each time sample is encoded with 256-dimensional sinusoidal positional encoding, transformed separately, summed, and injected into every decoder layer.
  • FiLM modulation transforms decoder features as ĥ_i = (1 + γ(c))h_i + ζ(c), with γ and ζ implemented as learnable linear layers.

E TPFM-DIR Backbone Architectures

TPFM-DIR is architecture-agnostic: CorrMLP and TransMorph are adapted by modifying their decoders to accept time context. The adaptations use FiLM modulation while leaving other architectural components unchanged.

  • CorrMLP and TransMorph are evaluated as TPFM-DIR backbones after adapting their decoders to incorporate time dependence.
  • Time context is injected into the adapted decoder layers through FiLM modulation to make the backbones suitable for TPFM-DIR.
  • Only the decoders are modified, while the original models’ other architectural components, including encoder and decoder channel dimensions, remain unchanged.

F Complete Quantitative Results

The supplementary section provides complete-metric quantitative results across all evaluated dataset groups, including brain, abdominal, lung, and cardiac imaging.

  • Full-metric results are reported for OASIS, IXI, LPBA40, CANDI, Mindboggle101, AbdomenCT, LungCT, and 2D cardiac datasets.Tables 8–11 organize the results by dataset group.

G Visual Results

The supplementary visual results examine regional Dice distributions, forward and backward warping, and comparisons across additional datasets using deformation-quality metrics.

  • Regional Dice-score distributions are shown for OASIS, IXI, CANDI, and ACDC.Figure 11 covers different anatomical regions across these four datasets.
  • LPBA40 visualizes TPFM-DIR forward and backward warping, with Dice, SDLogJ, and negative-Jacobian percentages reported alongside deformation grids.Forward warping maps moving images toward fixed images, while backward warping performs the reverse.
  • IXI and Mindboggle101 provide visual comparisons with other diffeomorphic and deformable methods, reporting Dice and deformation-grid metrics.Figures 13 and 14 place Dice, SDLogJ, and |J|<0% values around the warped images or grids.
  • Quantitative comparison tables identify best and second-best results across OASIS, IXI, LPBA40, Mindboggle101, CANDI, AbdomenCT, LungCT, ACDC, and CAMUS.Tables distinguish second-best diffeomorphic and deformable results where applicable.
Loading 2609.10789v1…