Factorized Gradients for Scalable Highly-expressive Parametric Diffeomorphisms
Authors: Amit Aflalo, Eran Treister, Chaim Baskin, Oren Freifeld
Organizations: The Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel · School of Electrical and Computer Engineering, Ben-Gurion University of the Negev, Israel
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(θ), the dimension of the parameter vector θ, 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 θ. 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 O(dTN) to 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., ≈104× in 2D), substantial memory savings (e.g., ≈500× 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.
Figures & tables
Figure 1: Quantitative comparison, in 2D, of performance scaling with tessellation size ( Ntess ), between the proposed method, called FG-CPAB , and libcpab [ 10 ] , the standard implementation of CPAB diffeomorphisms [ 12 ] . The evaluation was performed on a single RTX A6000 GPU with 125,000 2D points and a batch size of 16. (a) Forward pass latency. (b) Gradient computation latency. (c) Peak GPU memory allocation. In 2D, the domain is partitioned into C=4Ntess2 triangular cells. Notably, the standard implementation does not support Ntess>12 , whereas ours scales efficiently.
2D Integration ( 512×512 )
3D Integration ( 128×128×128 )
Method
Forward (ms)
Gradient (ms)
Peak Mem (MiB)
Forward (ms)
Gradient (ms)
Peak Mem (MiB)
FG-CPAB
1.69±0.06
8.64±0.02
197.97
31.10±0.18
117.50±0.34
2351.61
VoxelMorph SVF
11.41±0.02
13.42±0.08
1330.13
119.14±0.03
220.21±0.24
14432.13
torch-lddmm
51.23±0.41
169.75±0.06
3162.13
732.81±0.40
2414.12±0.53
43288.13
Table 1: Forward transformation and parameter gradient ( ∇θ ) latency, alongside peak memory footprint, on 2D and 3D data.
Dim
Params (C,S)
MSE C ↓
MSE S ↓
Fit C (s) ↓
Fit S (s) ↓
∣J∣≤0 (%) (C,S) ↓
2D
842 / 524288
8.853e-05 ± 3.9e-05
2.651e-03 ± 1.2e-03
0.525 ± 0.044
1.356 ± 0.007
0.028% / 0.178%
3D
8232 / 6291456
7.255e-05 ± 2.5e-05
1.676e-04 ± 4.1e-05
6.140 ± 0.021
12.906 ± 0.021
0.012% / 0.021%
Table 2: Quantitative synthetic registration results comparing FG-CPAB (C) and dense SVF (S) under the same optimization budget in 2D and 3D. Metrics report final reconstruction MSE, wall-clock fitting time, and the percentage of spatial locations with non-positive Jacobian determinant.
Figure 2: Qualitative comparison of CPAB and dense SVF on synthetic deformation fitting. Left to right: Source image, deformed target, absolute alignment error after optimization with CPAB, and absolute alignment error after optimization with dense SVF under the same optimization budget. CPAB achieves near-zero residual error despite its compact parameterization, whereas dense SVF more often exhibits structured residual artifacts under the matched budget.
Experiment
Geometry model
Test rel. L2 mean
Test rel. L2 std
FG-CPAB
FG-CPAB parameterization
0.052811
0.022436
LDDMM deformation PCA
PCA deformation code
0.082821
0.033168
Table 4: Reaction–diffusion operator-learning results using different geometry encodings. Lower is better for both metrics.
Figure 4: A composition of CPAB transformations can model deformations that a single stationary CPAB cannot. From left to right: source image, target image, absolute error for a single CPAB transformation, and absolute error for a composition of CPAB transformations. The target requires the two blobs to follow crossing trajectories, which cannot be captured well by a single stationary velocity field. While a single CPAB leaves substantial residual error, the composed model greatly reduces the error by realizing the deformation through multiple diffeomorphic stages.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Performance scaling with tessellation size ( Ntess ). Comparison of our optimized 3D CPAB vs. standard implementation on a single RTX A6000 GPU. (a) Forward pass latency. (b) Gradient computation latency. (c) Peak GPU memory allocation. Our approach yields up to a 60× speedup in the forward pass and a 4477× speedup in the gradient computation, while reducing peak memory by over 180× . Notably, the standard implementation does not support Ntess>5 , whereas ours scales efficiently.
Figure 6: Performance scaling with tessellation size ( Ntess ). Comparison of our optimized 1D CPAB vs. libcpab implementation on a single RTX A6000 GPU. (a) Forward pass latency. (b) Gradient computation latency. (c) Peak GPU memory allocation. Our approach yields up to a 7× speedup in the forward pass and a 301× speedup in the gradient computation, while reducing peak memory by 49× .
Figure 7: Performance scaling with tessellation size ( Ntess ). Comparison of our optimized 1D CPAB vs. the DIFW implementation on a single RTX A6000 GPU. (a) Forward pass latency. (b) Gradient computation latency. (c) Peak GPU memory allocation. Our approach yields up to a 2× speedup in the forward pass and a 104× speedup in the gradient computation, while reducing peak memory by 196× .
Figure 8: Effect of tessellation resolution on registration performance. Increasing the tessellation resolution improves deformation quality by providing a more expressive piecewise-affine parameterization. This trend is observed across both real and synthetic registration settings, supporting the claim that the scalability unlocked by gradient factorization is important not only for efficiency, but also for final task performance.
Sparse basis
Dense basis
Tess.
Fwd (ms)
Bwd (ms)
Peak (MiB)
Fwd (ms)
Bwd (ms)
Peak (MiB)
2
10.8
51.9
384
10.9
51.9
384
3
11.3
54.0
384
11.3
54.1
384
4
11.5
54.9
385
11.4
55.0
385
5
11.6
55.6
385
11.6
55.7
385
6
11.7
56.1
386
11.7
56.3
386
Appendix
Table 5: Sparse vs. Dense CPA basis across tessellation size. Forward/backward are mean runtime in ms (lower is better); peak is memory in MiB (lower is better).
nsolver
Forward (ms) ↓
Backward (ms) ↓
Fit time (s) ↓
Registration MSE ↓
5
0.1054
0.1822
0.2559
1.42×10−2
10
0.1068
0.1792
0.2490
2.27×10−3
15
0.1080
0.2080
0.2486
3.23×10−4
20
0.1138
0.2279
0.2438
1.04×10−4
25
0.1206
0.2523
0.2482
9.44×10−5
30
0.1248
0.2910
0.2483
9.18×10−5
Appendix
Table 6: Full CPAB ODE solver-step sweep over nsolver∈{5,10,…,100} . Lower values are better. Accuracy improves dramatically from the very low-step regime and then plateaus, while runtime continues to increase with the number of solver steps.
Figure 9: Convergence without regularization. Optimization trajectories for CPAB and dense SVF on the synthetic fitting task without additional regularization, shown for both 2D and 3D settings. In both cases, CPAB rapidly reaches a low reconstruction error within the first few hundred iterations, whereas dense SVF plateaus early at a substantially higher MSE. This highlights the optimization difficulty of the dense SVF parameterization under a matched iteration budget, despite its much larger number of degrees of freedom.
Figure 10: Convergence with regularization. Optimization trajectories for CPAB and dense SVF when dense SVF is equipped with standard smoothness regularization, shown for both 2D and 3D settings over an extended budget of 100,000 iterations. CPAB demonstrates extremely rapid initial convergence in both settings, stabilizing within the first few hundred iterations. In 2D (a), CPAB achieves and maintains a consistently lower fitting error. In 3D (b), regularized dense SVF converges much more slowly, but eventually surpasses CPAB to achieve a lower final mean squared error after approximately 50,000 to 60,000 iterations.
Ckpt seg.
Segments
Fwd (ms)
Bwd (ms)
2
25
0.898
4.161
3
17
0.885
3.977
4
13
0.882
3.860
5
10
0.885
3.905
6
9
0.891
3.899
7
8
0.889
4.302
Appendix
Table 7: Checkpoint segment size sweep. Forward/backward are mean runtime in ms (lower is better). Bold indicates the fastest value in each metric.
Asset
How used in this paper
Citation
Version
License
IXI brain MRI dataset
3D deformable registration benchmark
[ 21 ]
link 26.4.26
Creative Commons CC BY-SA 3.0 license
DIMON reaction–diffusion benchmark data / pipeline
Operator-learning experiment in § 4.5
[ 39 ]
Commit - a22ed33
MIT License
Appendix
Table 8: Datasets and benchmark data used in the paper.
Asset
How used in this paper
Citation
Version / commit
License
libcpab
Primary CPAB baseline and reference implementation
[ 10 ]
Commit - 22df6cd
MIT license
DIFW
1D performance baseline in Appendix A
[ 28 ]
Commit - f970dd3
MIT license
VoxelMorph
Dense SVF baseline / runtime comparison
[ 3 ]
Commit - 6fea3e8
Apache-2.0 license
torch-lddmm
LDDMM runtime baseline
[ 4 ]
Commit - 86d5c43
Apache-2.0 license
Appendix
Table 9: External codebases, libraries, and baseline implementations used by the project.
Continuous Piecewise-Affine Based (CPAB) transformations, introduced by Freifeld et al. (2015, 2017), are flows obtained by integrating continuous piecewise-affine (CPA) velocity fields. These expressive transformations have proven useful in applications including nonlinear temporal warping and trainable activation functions in 1D, as well as nonrigid spatial deformations in 2D and 3D. While computational costs were not an issue in 1D, they originally limited the practical expressiveness of CPAB transformations in higher dimensions, especially in 3D. Our recent work (Aflalo et al., 2026) removed this bottleneck, enabling efficient optimization with much finer tessellations and substantially greater expressiveness. There is, however, a mathematical subtlety that was overlooked in the original CPAB papers (Freifeld et al., 2015, 2017). CPAB transformations are always bi-Lipschitz homeomorphisms, implying that both the transformations and their inverses are differentiable almost everywhere. They are not, however, necessarily diffeomorphisms. This note remedies the omission by identifying sufficient regularity conditions under which a CPAB transformation is guaranteed to be a diffeomorphism. We explain how differentiability can fail and how the proposed conditions prevent this failure. In 1D and 2D, we show that these conditions can be easily enforced and, for commonly used tessellations, hold with probability one under standard sampling assumptions. We expect analogous behavior in 3D.
Amit Aflalo, Oren Freifeld
The Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel
Persistence-based topological optimization deforms a point cloud X⊂Rd by minimizing objectives of the form L(X)=ℓ(Dgm(X)), where Dgm(X) is a persistence diagram. In practice, optimization is limited by two coupled issues: persistent homology is typically computed on subsamples, and the resulting topological gradients are highly sparse, with only a few anchor points receiving nonzero updates. Motivated by diffeomorphic interpolation, which extends sparse gradients to smooth ambient vector fields via Reproducing Kernel Hilbert Space (RKHS) interpolation, we propose a more scalable pipeline that improves both subsampling and gradient extension. We introduce subsampling via random slicing, a lightweight scheme that promotes iteration-wise geometric coverage and mitigates density bias. We further replace the costly kernel solve with a fast Nadaraya-Watson (NW) Gaussian convolution, producing a globally defined smooth update field at a fraction of the computational cost, while being more suited for topological optimization tasks. We provide theoretical guarantees for NW smoothing, including anchor approximation bounds and global Lipschitz estimates. Experiments in 2D and 3D show that combining random slicing with NW smoothing yields consistent speedups and improved objective values over other baselines on common persistence losses.
Contact-rich inverse problems require efficient gradients and reliable enforcement of non-penetration. Existing differentiable methods provide fast simulation and rigid-deformable coupling, but contact treatments based on penalties or velocity-level constraints can permit penetration or positional drift. Incremental Potential Contact (IPC) maintains non-penetration through geometric barriers and {continuous collision detection (CCD)}. However, repeated forward and backward evaluations remain costly in full-space finite-element (FEM) implementations, particularly when nearly rigid components retain many vertex {degrees of freedom (DoFs)}. We present DiffGIPC, a GPU-accelerated differentiable IPC framework for optimizing interacting deformable and nearly rigid bodies. Our affine-body derivatives enable {rest-shape} design, initial-state optimization, and articulated control while retaining only twelve motion DoFs per affine body. These gradients account for how changes in geometry, placement, and actuation affect other objects through collision and friction. We preserve these interactions in one coupled FEM-affine adjoint system per frame, reducing the backward-system size without discarding the dependencies that drive the inverse problem. GPU-optimized backward computation further reduces the cost of repeated gradient evaluations. A reusable PyTorch interface connects supported physical parameters to user-defined parameter maps and losses across design, identification, motion, placement, and articulated control. Experiments demonstrate peak speedups over DiffIPC of approximately 150.0x with full-space FEM and 4133.3x with affine bodies for combined forward simulation and backward differentiation.
Kemeng Huang, Bowen Ren, Xinyu Lu +2
The University of Hong Kong, Hong Kong SAR, China · TransGP, Hong Kong SAR, China · Carnegie Mellon University, USA