Factorized Gradients for Scalable Highly-expressive Parametric Diffeomorphisms

Diffeomorphisms provide a flexible, topology-preserving mathematical tool for modeling complex spatial deformations, and are used in various scientific fields. However, simultaneously achieving high expressivity and computational efficiency remains challenging. Continuous Piecewise-Affine Based (CPAB) transformations offer an attractive solution by parameterizing a family of diffeomorphisms via continuous velocity fields that are piecewise affine *w.r.t.* 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 computing gradients *w.r.t.* $\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 + d C))$, 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 huge speedups (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, for the first time, fine-tessellation CPAB practical in 2D and 3D, unlocking the potential of highly-expressive parametric diffeomorphisms.