Factorized Gradients for Scalable Highly-expressive Parametric Diffeomorphisms

Amit Aflalo; Eran Treister; Chaim Baskin; Oren Freifeld

cs.CVOct 26, 2026

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

Abstract

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 $\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.

Figures & tables

Figure 1: Quantitative comparison, in 2D, of performance scaling with tessellation size (

N_{\text{tess}}

), 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=4N_{\text{tess}}^{2}

triangular cells. Notably, the standard implementation does not support

N_{\mathrm{tess}}>12

, whereas ours scales efficiently.

	2D Integration ( $512\times 512$ )			3D Integration ( $128\times 128\times 128$ )
Method	Forward (ms)	Gradient (ms)	Peak Mem (MiB)	Forward (ms)	Gradient (ms)	Peak Mem (MiB)
FG-CPAB	$\mathbf{1.69\pm 0.06}$	$\mathbf{8.64\pm 0.02}$	$\mathbf{197.97}$	$\mathbf{31.10\pm 0.18}$	$\mathbf{117.50\pm 0.34}$	$\mathbf{2351.61}$
VoxelMorph SVF	$11.41\pm 0.02$	$13.42\pm 0.08$	$1330.13$	$119.14\pm 0.03$	$220.21\pm 0.24$	$14432.13$
torch-lddmm	$51.23\pm 0.41$	$169.75\pm 0.06$	$3162.13$	$732.81\pm 0.40$	$2414.12\pm 0.53$	$43288.13$

Table 1: Forward transformation and parameter gradient (

\nabla_{\!\boldsymbol{\theta}}

) latency, alongside peak memory footprint, on 2D and 3D data.

Dim	Params (C,S)	MSE C $\downarrow$	MSE S $\downarrow$	Fit C (s) $\downarrow$	Fit S (s) $\downarrow$	$\|J\|\leq 0$ (%) (C,S) $\downarrow$
2D	842 / 524288	8.853e-05 $\pm$ 3.9e-05	2.651e-03 $\pm$ 1.2e-03	0.525 $\pm$ 0.044	1.356 $\pm$ 0.007	0.028% / 0.178%
3D	8232 / 6291456	7.255e-05 $\pm$ 2.5e-05	1.676e-04 $\pm$ 4.1e-05	6.140 $\pm$ 0.021	12.906 $\pm$ 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. $L^{2}$ mean	Test rel. $L^{2}$ 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 (

N_{\mathrm{tess}}

). 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\times

speedup in the forward pass and a

4477\times

speedup in the gradient computation, while reducing peak memory by over

180\times

. Notably, the standard implementation does not support

N_{\mathrm{tess}}>5

, whereas ours scales efficiently.

Figure 6: Performance scaling with tessellation size (

N_{\mathrm{tess}}

). 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\times

speedup in the forward pass and a

301\times

speedup in the gradient computation, while reducing peak memory by

49\times

Figure 7: Performance scaling with tessellation size (

N_{\mathrm{tess}}

). 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\times

speedup in the forward pass and a

104\times

speedup in the gradient computation, while reducing peak memory by

196\times

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).

$n_{\mathrm{solver}}$	Forward (ms) $\downarrow$	Backward (ms) $\downarrow$	Fit time (s) $\downarrow$	Registration MSE $\downarrow$
5	0.1054	0.1822	0.2559	$1.42\times 10^{-2}$
10	0.1068	0.1792	0.2490	$2.27\times 10^{-3}$
15	0.1080	0.2080	0.2486	$3.23\times 10^{-4}$
20	0.1138	0.2279	0.2438	$1.04\times 10^{-4}$
25	0.1206	0.2523	0.2482	$9.44\times 10^{-5}$
30	0.1248	0.2910	0.2483	$9.18\times 10^{-5}$

Appendix

Table 6: Full CPAB ODE solver-step sweep over

n_{\mathrm{solver}}\in\{5,10,\dots,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.

Gradient Projection

Gradient projection constrains optimization steps to feasible directions or subspaces, reducing interference while preserving desired properties.

See all 21 papers

A Regularity Condition Under Which a CPAB Transformation is a Diffeomorphism
Oct 23, 2026Amit Aflalo, Oren FreifeldDynamical Systems
The Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel
Towards Scalable Persistence-Based Topological Optimization
May 9, 2026Abderrahim Bendahi, Alexandre Duplessis, Arnaud FickingerPersistent Homology Topological Data Analysis
ENS Ulm, PSL Paris, France · UC Berkeley Berkeley, CA, USA
DiffGIPC: Differentiable Mixed-Coordinate Contact Dynamics on GPU
Oct 13, 2026Kemeng Huang, Bowen Ren, Xinyu Lu +2Differentiable Physics Inverse Problems
The University of Hong Kong, Hong Kong SAR, China · TransGP, Hong Kong SAR, China · Carnegie Mellon University, USA

Oct 23, 2026J/K move · Enter open · S save

A Regularity Condition Under Which a CPAB Transformation is a Diffeomorphism

Amit Aflalo, Oren Freifeld

The Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Israel

Dynamical Systems

Open paper AskPDF arXiv

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.

Abstract

Gradient Projection

Explore similar work

A Regularity Condition Under Which a CPAB Transformation is a Diffeomorphism

Towards Scalable Persistence-Based Topological Optimization

DiffGIPC: Differentiable Mixed-Coordinate Contact Dynamics on GPU