PaperlayerView exact revision on arXiv

Source-linked AI summary

Factorized Gradients for Scalable Highly-expressive Parametric Diffeomorphisms

Amit Aflalo, Eran Treister, Chaim Baskin, Oren Freifeld

arXiv:2610.34766v1cs.CV

TL;DR

CPAB offers compact, tessellation-based diffeomorphisms, but standard gradient computation couples parameter count to trajectory integration and limits practical tessellation fineness. FG-CPAB factorizes gradients into cell-wise moments and a final parameter projection, making fine-tessellation optimization practical in 2D and 3D with improved expressiveness and competitive registration accuracy.

  • Problem

    Standard CPAB gradients couple all d parameters to T integration steps for each of N points, making fine tessellations and large spatial inputs costly to optimize.

    Source: §1 · p. 2

  • Method

    FG-CPAB accumulates parameter-independent cell-wise geometric moments during a discrete-adjoint sweep and recovers gradients through a final sparse CPA-basis projection.

    Source: §1 · pp. 2–3, §5 · p. 10

  • Results

    FG-CPAB enables finer-tessellation optimization with better practical expressiveness, competitive registration accuracy, and strong topology preservation.

    Source: §5 · p. 10

  • Takeaways & Limitations

    FG-CPAB makes CPAB a more practical diffeomorphic primitive for large-scale spatial or temporal modeling in classical optimization and deep-learning pipelines.

    Source: §5 · p. 10

    Empirical claims are restricted to spatial dimensions K ≤3 because tessellation cell counts grow exponentially with dimension.

    Source: §5 · p. 10

Abstract

from arXiv · show

Diffeomorphisms provide a flexible, topology-preserving mathematical tool for modeling complex spatial or temporal deformations and are used in various scientific fields. However, making these deformations simultaneously highly expressive and computationally efficient remains challenging. Continuous Piecewise-Affine Based (CPAB) transformations (Freifeld et al., ICCV 2015; IEEE TPAMI 2017) offer an attractive solution by parameterizing a family of diffeomorphisms via continuous velocity fields that are piecewise affine with respect to a chosen tessellation of the domain. Importantly, $d=\dim(\theta)$, the dimension of the parameter vector $\theta$, depends on the fineness of the tessellation rather than on the data resolution. Despite this compact parameterization, CPAB optimization has long been hindered by the tight coupling between trajectory integration and gradient computation with respect to $\theta$. We introduce FG-CPAB, a factorized-gradient formulation for scalable CPAB transformations. By separating integration from parameter-space projection, our method reduces gradient computation complexity from $\mathcal{O}(dTN)$ to $\mathcal{O}(TN+dC)$, where $T$ is the number of integration steps, $N$ is the number of transformed points, and $C$ is the number of tessellation cells. This reformulation yields speedups of up to four orders of magnitude (e.g., $\approx 10^4\times$ in 2D), substantial memory savings (e.g., $\approx 500\times$ reduction in peak GPU memory usage in 2D), and dramatically higher expressiveness. This makes fine-tessellation CPAB practical in 2D and 3D for the first time, enabling highly expressive parametric diffeomorphisms.

1 Introduction

CPAB parameterizes topology-preserving deformations with complexity set by tessellation rather than data resolution, but its gradients have made fine tessellations expensive. FG-CPAB factorizes gradient computation, enabling more practical optimization and stronger expressiveness.

1 Introduction

  • Dense-grid diffeomorphism methods scale parameter count with spatial resolution, contributing to high memory use, slow optimization, and difficult optimization landscapes.Diffeomorphisms are used where smooth, invertible, topology-preserving deformations are required.Source: §1 · p. 2
  • CPAB integrates a continuous piecewise-affine velocity field over tessellation cells, so parameter count depends on tessellation rather than data resolution.Under a mild regularity condition, CPAB transformations are diffeomorphisms.Source: §1 · p. 2
  • Standard CPAB gradients propagate sensitivities for all d parameters through T steps for each of N points, giving O(dTN) complexity and limiting practical tessellation fineness.The bottleneck is particularly restrictive for large spatial inputs and 3D optimization.Source: §1 · p. 2
  • FG-CPAB separates trajectory integration from parameter projection, using cell-wise geometric statistics and a sparse CPA basis to reduce backward complexity to O(TN + dC).The method recovers parameter gradients in a final projection step.Source: §1 · pp. 2–3
  • 75× faster forward passes, 104× faster backward passes, and over 500× lower peak memory are reported in 2D versus the standard implementation.The gains support efficient optimization of fine-tessellation CPAB in both 2D and 3D.Source: §1 · p. 2

2 Background and Related Work

Diffeomorphisms can be constructed by integrating velocity fields, but dense parameterizations tie model size to spatial resolution. CPAB instead uses continuous piecewise-affine fields over tessellations, with continuity constraints yielding a compact parameter space.

2 Background and Related Work

  • Diffeomorphisms are obtained by integrating velocity fields; SVF assumes a time-constant field, while LDDMM can use time-varying fields.The mapping is the endpoint of each point's trajectory.Source: §2 · p. 3
  • Dense velocity-field parameterizations can require hundreds of thousands of parameters in 2D and millions in 3D, tying model size to domain resolution.The resulting parameter space can require substantial memory and be difficult to optimize without explicit regularization.Source: §2 · p. 3
  • CPAB integrates a continuous piecewise-affine velocity field defined over convex polyhedral cells; simplicial cells are intervals, triangles, or tetrahedra in dimensions one, two, or three.With zero boundary conditions the implementation uses simplicial tessellations; otherwise it augments them with unbounded cells.Source: §2 · p. 3
  • Continuity constraints require neighboring cells' affine fields to agree on shared boundaries, restricting the concatenated coefficients to the constraint system's null space.The coefficient vector has length D = CK(K + 1), with d no greater than D.Source: §2 · pp. 3–4
  • The parameter dimension d shifts from spatial resolution to tessellation complexity and increases as the tessellation is refined.Optional linear constraints can impose zero boundary conditions or volume preservation.Source: §2 · p. 4
  • Prior CPAB gradients propagate parameter sensitivities through trajectories at O(dTN), restricting fine tessellations; generic adjoints do not remove this parameter dependence.MCMC-based inference scales poorly in the fine-tessellation regime and does not provide gradients for end-to-end learning.Source: §2 · p. 4

3 Method

FG-CPAB decouples trajectory integration from parameter dimension by accumulating cell-local moments and projecting them onto the CPA basis only at the end. Its gradient computation has O(TN + dC) cost, with checkpointing reducing trajectory storage through replay.

3 Method

  • FG-CPAB reduces backward complexity from O(dTN) to O(TN + dC) by accumulating parameter-independent cell-wise moments and projecting once onto the linear CPA basis.A sparse basis formulation improves practical runtime and memory use in both forward and backward passes.Source: §3 · p. 4

3.1 Trajectory integration and gradient computation

  • The forward affine update is exact while a trajectory stays within a cell; at cell crossings, the method assumes a sufficiently small step and uses a midpoint rule for gradient computation.The backward pass uses a discretize-then-optimize treatment because nonsmooth cell boundaries can make an optimize-then-discretize adjoint inconsistent.Source: §3.1 · pp. 4–5

3.2 Adjoint factorization via moment accumulation

  • Summing sensitivity contributions across timesteps cancels propagated terms, leaving local contributions that can be grouped by cell and factored from the basis.The accumulated moments are independent of d, and the final sparse basis projection recovers all parameter gradients.Source: §3.2 · pp. 5–6

3.3 Memory-efficient adjoint via segmented checkpointing

  • Segmented checkpointing stores ⌈T/S⌉ states per point and replays each segment during the reverse sweep, reducing storage to O(N(T/S + S)) at constant-factor recomputation cost.The implementation uses S = 4.Source: §3.3 · p. 6

3.4 Gradient computation algorithm

  • The algorithm checkpoints the forward trajectory, performs a reverse adjoint sweep with moment accumulation, then projects through the sparse basis; total gradient cost is O(TN + dC).Integration and adjoint propagation are independent of d, while projection is independent of N and T.Source: §3.4 · p. 6

4 Results

The results evaluate FG-CPAB’s efficiency, expressiveness, registration performance, and use as a geometry code. Across these settings, factorized gradients and sparse bases make finer CPAB tessellations practical.

4.1 Acceleration and memory savings over the standard CPAB implementation

  • In 2D, FG-CPAB achieves up to 104× faster gradient computation and over 500× lower peak memory than standard CPAB; in 3D, backward acceleration reaches 4,477×.The standard implementation does not run beyond Ntess = 12 under the benchmark setup, while FG-CPAB scales to finer resolutions.Source: §4.1 · pp. 6–7

4.2 Forward and gradient computation runtime

  • FG-CPAB reduces total runtime by 2.4× versus VoxelMorph and 21.4× versus torch-lddmm in 2D, with 2.3× and 21.2× speedups in 3D.The comparison measures forward transformation and parameter-gradient latency on standard 2D and 3D grids.Source: §4.2 · p. 7, Table 1 · p. 7

4.3 Expressiveness and optimization stability compared with dense SVFs

  • Under matched optimization budgets, FG-CPAB achieves lower residual error, faster convergence, and fewer topology violations than dense SVF.Dense SVFs can eventually approach or surpass CPAB reconstruction error with heavy regularization, but only after orders of magnitude more optimization steps.Source: §4.3 · p. 7, Table 2 · p. 7, Figure 2 · p. 8

4.4 CPAB for 3D deformable registration

  • On IXI, a single CPAB reaches 0.760 ± 0.023 Dice with fewer than 0.0001% non-positive Jacobians, while per-pair optimization takes approximately 5 seconds.Composing four transformations raises Dice to 0.765 while the non-positive Jacobian fraction remains below 0.0001.Source: §4.4 · p. 8, Table 3 · p. 8

4.5 CPAB as a low-dimensional geometry code for neural operator learning

  • Using CPAB parameters as geometry codes yields the lowest test relative L2 error among the compared encodings under the same operator backbone.The CPAB code is tied directly to the transport map, with dimension controlled by the tessellation.Source: §4.5 · p. 9, Table 4 · p. 9

4.6 CPAB composition for non-stationary velocity fields

  • Composing CPAB transformations reduces error for crossing-trajectory deformations that a single stationary CPAB cannot model accurately.The composition remains a diffeomorphism, although CPAB transformations are not closed under composition.Source: §4.6 · p. 9, Figure 4 · p. 10

4.7 Ablations and hyper-parameter analysis

  • Finer tessellations improve alignment quality, sparse bases save memory and increasingly benefit runtime at larger tessellations, and accuracy largely plateaus around 50 solver steps.All paper results use 50 CPAB solver steps unless noted otherwise; runtime continues to increase with more steps.Source: Figure 3 · p. 8, §4.7 · pp. 9–10

5 Discussion

FG-CPAB makes CPAB more practical by separating trajectory integration from parameter-space projection, while enabling finer tessellations that improve deformation quality. Its scalability gains are complemented by a sparse basis that chiefly reduces memory use.

5 Discussion

  • The discussion characterizes improved scalability as supporting greater practical expressiveness, competitive registration accuracy, and strong topology preservation in 2D and 3D.The paper’s performance evaluation compares against standard libcpab in 2D and 3D and additionally against DIFW in 1D.Source: §5 · p. 10, A Performance · p. 14

A.1 Why the standard backward stops scaling

  • FG-CPAB removes parameter dependence from integration, accumulating cell-wise moments before a separate projection so the reference backward pass can be avoided at finer tessellations.The reference implementation’s parameter-dependent kernel grid reaches a device-dependent launch ceiling as tessellation grows.Source: §A.1 · p. 14

A Performance

  • 4477× speedup in 3D gradient computation and over 180× lower peak memory are reported for FG-CPAB, which also scales beyond tessellation sizes supported by the standard implementation.The comparison uses a single RTX A6000 GPU; the figure reports up to 60× forward-pass speedup as well.Source: Figure 5 · p. 14

B.1 Tessellation resolution hyper-parameter analysis

  • Increasing tessellation resolution improves alignment quality through the largest resolutions tested, making finer parameterizations useful for modeling complex deformations.The ablation evaluates real and synthetic registration settings; prior implementations make these larger tessellations too slow or memory-intensive to optimize reliably, especially in 3D.Source: §B.1 · p. 16, Figure 8 · p. 16

B.2 Sparse basis ablation

  • At tessellation size 20, the sparse basis reduces peak memory from 37678 MiB to 474 MiB, while also lowering forward latency from 85.0 ms to 23.5 ms and backward latency from 140.3 ms to 83.4 ms.The sparse formulation mainly improves memory use; its additional runtime benefits become more pronounced at larger tessellations.Source: §B.2 · pp. 16–17, Table 5 · p. 17

B.3 Solver step hyper-parameter analysis

The solver-step sweep examines the runtime–accuracy trade-off for 2D CPAB integration. Accuracy improves substantially at low step counts and then plateaus, supporting 50 steps as a robust default.

B.3 Solver step hyper-parameter analysis

  • Increasing solver steps from 5 to 20 reduces mean MSE from 1.42 × 10−2 to 1.04 × 10−4, while accuracy largely plateaus beyond 20 steps.The sweep varies only CPAB solver steps in a 2D synthetic-registration setup with a 30 × 30 tessellation.Source: §B.3 · p. 17, Table 6 · p. 18
  • At 100 steps, mean MSE is 8.31 × 10−5, while mean fitting time rises from 0.287 s to 0.468 s and backward latency from 0.415 ms to 0.720 ms.Source: §B.3 · p. 17
  • 50 solver steps are a robust operating point in the accuracy plateau, capturing most attainable improvement without the noticeably higher cost of longer integrations.The paper uses 50 steps unless otherwise noted.Source: §B.3 · p. 17

C.1 Optimization Stability and Practical Expressiveness Compared with Dense SVFs

In the evaluated registration tasks, FG-CPAB optimizes more reliably and quickly than dense SVF under a matched budget, with fewer topology violations. Extended-budget tests show that regularized dense SVF can eventually attain lower 3D error, while segmented checkpointing reduces trajectory storage at the cost of recomputation.

C.1 Optimization Stability and Practical Expressiveness Compared with Dense SVFs

  • Under a matched optimization budget, FG-CPAB achieves lower residual error, shorter fitting times, and fewer topology violations than dense SVF.Results aggregate 100 independently generated target registrations per dimension; both methods use the same reconstruction objective and Adam optimization.Source: §C.1 · p. 19
  • After approximately 50,000–60,000 iterations, regularized dense SVF attains lower final MSE than CPAB in 3D, but not in 2D within the tested budget.The comparison uses 100,000 iterations; regularization improves dense SVF stability, but convergence remains much slower than CPAB.Source: §C.1 · p. 19, Figure 10 · p. 20

D Segmented Checkpointing

  • Segmented checkpointing replaces O(T) per-thread trajectory storage with O(T/S) checkpoint states and O(S) temporary states, while adding a constant-factor recomputation overhead.The forward trajectory is replayed segment by segment during the reverse sweep; the factorized gradient itself is unchanged.Source: D Segmented Checkpointing · p. 21

E Application-Specific Experimental Details

The application details describe CPAB evaluation for 3D atlas-to-patient registration and as a geometry code for neural operator learning. A nodal construction provides a sparse CPA basis while preserving the globally continuous velocity fields of the direct algebraic construction.

E.1 CPAB for 3D deformable registration

  • FG-CPAB is evaluated on all 115 IXI test subjects using atlas-to-patient registration, with anatomical labels reserved for evaluation.The data comprise 576 preprocessed T1-weighted MRI volumes; CPAB uses a 20 × 20 × 20 tessellation and fits each test subject independently for 200 Adam iterations.Source: §E.1 · p. 23

E.2 CPAB as a low-dimensional geometry code for operator learning

  • For DIMON operator learning, CPAB parameters serve as the geometry code and are compared with 122 PCA coefficients, keeping both code dimensions equal.The DIMON data generation, transport pipeline, and operator architecture are unchanged.Source: §E.2 · p. 23

F Nodal Sparse Basis

  • The nodal construction maps vertex velocities to cell-wise affine coefficients and yields a sparse basis that preserves the same globally continuous CPA velocity fields as the direct algebraic construction.On a standard simplex mesh, continuity across shared cell boundaries is automatic; the remaining active constraints are global boundary conditions.Source: F Nodal Sparse Basis · pp. 24–25

G Proof of the Telescoping Property of the Discrete Adjoint

The proof shows that the discrete-adjoint sensitivity terms telescope, leaving the direct parameter contributions that form cell-wise moments. It also shows that translation does not affect the midpoint step’s state Jacobian.

G Proof of the Telescoping Property of the Discrete Adjoint

  • Summing the discrete-adjoint identity over timesteps cancels propagated sensitivity terms, leaving the accumulated direct contributions.The proof uses the fixed, parameter-independent initial point and the loss-gradient definition of the final adjoint as boundary conditions.Source: G Proof of the Telescoping Property of… · p. 26
  • The remaining direct contributions assemble into cell-wise moment matrices, from which the parameter gradient is recovered by projection onto the basis.The cancellation is algebraic for the exact discrete adjoint of the midpoint integrator and holds to machine precision.Source: G Proof of the Telescoping Property of… · p. 26

H Why Translation Does Not Enter the State Jacobian

  • The midpoint step’s state Jacobian depends only on the affine velocity’s linear block; its translation contributes only additive terms independent of the state.Because the translation terms do not depend on the state, their derivatives with respect to the state vanish.Source: H Why Translation Does Not Enter the… · p. 27

I Integrator mismatch

The forward matrix-exponential update and midpoint-based backward pass are not algebraically identical, so the computed gradient approximates the forward computation’s derivative. The paper analyzes the mismatch and validates gradients in its experimental regime.

Source of the mismatch and relation to prior CPAB implementations

  • The forward pass uses a matrix-exponential update, while the backward pass differentiates midpoint-discretized dynamics, making the gradient approximate for the forward computation.The paper retains this pairing, also used by reference CPAB implementations, for a fair comparison with prior methods.Source: I Integrator mismatch · p. 28

Effect of inconsistent cell assignments and different integrators within a fixed cell

  • Inconsistent cell assignments can cause an O(h) local discrepancy, whereas the fixed-cell integrator mismatch has O(h^3) local and O(h^2) global error.The integrator-error comparison assumes a fixed cell matrix over a finite integration interval; the cell-assignment result assumes a bounded state.Source: I Integrator mismatch · p. 28

Empirical gradient validation

  • With 50 solver steps, cosine similarity exceeded 0.998 and relative error was below 0.05 against the finite-difference reference gradient.The authors report that the approximate gradient is closely aligned with the reference in the regime used in their experiments.Source: I Integrator mismatch · p. 28

J Why Gradient-Based Optimization Is Needed: Comparison with MCMC

The paper argues that fine-tessellation CPAB needs gradient-based optimization for end-to-end backpropagation and high-dimensional parameter spaces. In its IXI registration benchmark, FG-CPAB outperformed MCMC under the reported time budgets.

End-to-end differentiability

  • MCMC does not provide the vector–Jacobian products needed to backpropagate through CPAB parameters into preceding components of a larger model.The passage distinguishes inference over CPAB parameters from the gradient propagation required for end-to-end training.Source: J Why Gradient-Based Optimization Is… · p. 29

Scaling with the parameter dimension

  • The paper reports 842 parameters in 2D, 8,232 in 3D, and 20,577 per transformation for IXI, and says MCMC is uncompetitive in this fine-tessellation regime.The original CPAB landmark experiment used 38 parameters; the authors caution that their findings do not make MCMC unsuitable for every CPAB application.Source: J Why Gradient-Based Optimization Is… · p. 29

Empirical comparison

  • With 100× more optimization time, MCMC achieved a Dice score of 0.453, versus 0.760 for gradient-based FG-CPAB on IXI atlas-to-patient registration.The comparison used the official CPAB MCMC implementation and the same registration setup as FG-CPAB.Source: J Why Gradient-Based Optimization Is… · p. 29

K.1 Public datasets and benchmark data

The two tables document the paper’s benchmark-data and external-software assets.

K.1 Public datasets and benchmark data

  • Table 8 covers datasets and benchmark data used in the paper.The caption does not identify individual datasets.Source: Table 8 · p. 30

K.2 External code, software, and baseline implementations

  • Table 9 covers external codebases, libraries, and baseline implementations used by the project.The caption does not name the individual codebases or baselines.Source: Table 9 · p. 30

Related papers

Loading 2610.34766v1…