Coordinate-Free Automatic Differentiation

Alex Guerra; Joshua Aduol; Yaqian Tang; Daniel Williams; Ryan P. Adams; Peter Orbanz

cs.LGOct 8, 2026

Coordinate-Free Automatic Differentiation

Authors: Alex Guerra, Joshua Aduol, Yaqian Tang, Daniel Williams, Ryan P. Adams, Peter Orbanz

Organizations: Department of Computer Science, Princeton University, Princeton, NJ 08540 · Gatsby Computational Neuroscience Unit, University College London

Abstract

Automatic differentiation (AD) is a cornerstone of modern machine learning. Since all AD frameworks implicitly fix a basis of Euclidean space to define gradients and Jacobians, they natively process mathematical objects that are arrays of a fixed set of coordinates. This makes it awkward to express objects such as rotations, volume forms, and curvatures that arise, for example, in applications of AI to science. Bespoke workarounds lack reusability. We show that (i) a single, basis-independent derivative operator can represent many geometric differential operations, and (ii) this operator can be expressed entirely in terms of existing AD primitives. This reusable transformation can be composed with JAX batching, compilation, and hardware acceleration. Geometry enters as algebraic data. We experimentally demonstrate its reuse across different manifold operators, its composition through a spacetime inverse-problem pipeline, and its use in optimization over different rotor subgroups.

Figures & tables

Figure 3: Flat square (black) flowed under curvature loss (red) vs. with the added Gauss–Bonnet term (orange).

Panel of Figure 4 — Figure 4: Learned minimal surfaces for trefoil (left) and a saddle (right).

Figure 7: Recovered trajectory ${r(t)}$ (red) against the target helix, with intermediate training trajectories also plotted.

Figure 8: Projections of ${F}$ field at ${t{=}2}$ .

Figure 9: Vectorized optimization over different subgroups: Given an object (gray) and a target conformal transformation (red), we simultanously learn 4 rotors in different rotor groups corresponding to maskings of the 10-dimensional bivector space of

{\mathrm{CGA}(3)}

: translation (3D), rigid (6D), similarity (7D), and conformal (10D).

Figure 10: Recovering a Lorentz rotor from field measurements. Left: a ship crosses the lab frame at uniform velocity (dashed), past a static point charge, measuring the Faraday bivector

\tilde{F_{k}}

in its own frame at events along its path,

{\tilde{F}_{k}=R^{\star}F_{k}\widetilde{R^{\star}}+\epsilon_{k}}

. In the ship frame the pure-electric Coulomb field tilts and compresses (electric part, blue) and acquires a circulating magnetic part (red). Right: training loss for recovering

{R^{\star}}

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 11: Initial tilings. Left: the icosidodecahedron on

S^{2}

. Right: the rectified

\{3,7\}

tiling of the Poincaré disk up to hyperbolic radius

3.5

Figure 12: The tilings under the learned divergence-free flow at

t=0,0.25,0.5,0.75,1

. Top: pentagons (blue) deformed into a chiral

\mathbb{Z}_{5}

pinwheel on

S^{2}

, under the icosahedral rotation group. Bottom: heptagons (blue) deformed into a

D_{7}

gear on

H^{2}

, under the

(2,3,7)

Coxeter group. In both, the triangles (orange) flow under the same trained velocity field, and the dashed curve at

t=1

is the target. The rightmost panels show the training loss (log scale) against the training step.

Figure 13: The square under the curvature loss alone (

\lambda_{\mathrm{GB}}=0

). Top: the learned flow at

t=0,\tfrac{1}{3},\tfrac{2}{3},1

, which collapses the square into a small spherical cap. Bottom: training loss (left) and the Monte Carlo Euler characteristic estimate

\hat{\chi}_{\theta}

(right) over training; without the topological term,

\hat{\chi}_{\theta}

stays far from the sphere’s value of 2.

Figure 14: The square under the curvature loss with the Gauss–Bonnet term (

\lambda_{\mathrm{GB}}=1.5

). Top: the learned flow at

t=0,\tfrac{1}{3},\tfrac{2}{3},1

, which closes the square into a sphere. Bottom: training loss (left) and the Monte Carlo Euler characteristic estimate

\hat{\chi}_{\theta}

(right) over training, which approaches 2.

boundary	Chamfer	Hausdorff (95%)	Hausdorff (max)	coverage	spurious	area (exact)
torus knot	0.0099	0.014	0.050	94.8%	0.13%	5.53 (5.44)
polynomial graph	0.0150	0.028	0.046	99.96%	1.51%	5.12 (5.21)
graph $w=z^{2}$	0.0123	0.023	0.041	100.0%	0.42%	4.69 (4.58)
annulus $zw=c$	0.0083	0.016	0.034	100.0%	0.08%	3.24 (3.08)
two disks $w=(c/\rho^{2})\bar{z}$	0.0126	0.021	0.048	97.4%	1.04%	5.80 (5.75)

Appendix

Table 1: Extracted codimension-two surfaces against the exact minimal surfaces. Chamfer is the symmetric mean nearest-neighbour distance; Hausdorff distances are one-sided (extraction to reference); coverage is the fraction of the reference within

0.03

of the extraction, and spurious the fraction of extracted points farther than

0.03

from the reference. Area is that of the render mesh.

Panel of Figure 15 — Figure 15: The two codimension-two surfaces of Figs. 6 and 6 at four projection angles $\theta=0,\pi/4,\pi/2,3\pi/4$ . Each block shows the exact minimal surface (top row) and the surface extracted from the learned field (bottom row).

Panel of Figure 16 — Figure 16: The three further codimension-two boundaries of Table 1 at four projection angles $\theta=0,\pi/4,\pi/2,3\pi/4$ . Each block shows the exact minimal surface (top row) and the surface extracted from the learned field (bottom row), rendered as in Fig. 15 .

Panel of Figure 17 — Figure 17: Field of the learned worldline during training, at $t=2$ : $|E|$ with in-plane $E$ streamlines (top) and the out-of-plane component of $B$ with in-plane $B$ streamlines (bottom), on the $xy$ , $xz$ and $yz$ planes through the target particle. The learned worldline is drawn solid and the target dashed.

	group	dim	multiplicative (ours)	static chart	BFGS	closed form
$\sigma=0$	$T(3)$	3	2.7797	2.7797	2.7797	2.7797
	$SE(3)$	6	1.5897	1.5897	1.5897	1.5897
	$\mathrm{Sim}(3)$	7	0.39614	0.39614	0.39614	0.39614
	conformal	10	$1.7\cdot 10^{-21}$	$8.1\cdot 10^{-4}$	$2.5\cdot 10^{-21}$	–
$\sigma=0.05$	$T(3)$	3	2.9138	2.9138	2.9138	2.9138
	$SE(3)$	6	1.6620	1.6620	1.6620	1.6620

Appendix

Table 2: Final losses for the four subgroups of the conformal group of Fig. 9 , fitted to one shared conformal target, noise-free (top) and with noise

\sigma=0.05

(bottom).

Figure 18: Training loss on the full conformal group under the multiplicative update (ours), the static chart, and the corrected chart. Left: from the default initialization near the identity. Right: from an initialization at a rotation of 4 rad, where the descent path crosses the conjugate point

g=-1

\exp

method	loss $/\sigma^{2}$	violation	$\Lambda$ error	rapidity error
multiplicative (ours)	1.03	$7.3\cdot 10^{-16}$	$3.9\cdot 10^{-3}$	$1.4\cdot 10^{-3}$
static chart	1.03	$1.0\cdot 10^{-15}$	$3.9\cdot 10^{-3}$	$1.4\cdot 10^{-3}$
corrected chart	1.03	$6.1\cdot 10^{-16}$	$3.9\cdot 10^{-3}$	$1.4\cdot 10^{-3}$
chart, BFGS (40 steps)	1.03	–	$3.9\cdot 10^{-3}$	$1.4\cdot 10^{-3}$
even multivector + renormalize	1.06	$2.0\cdot 10^{-15}$	$1.0\cdot 10^{-2}$	$1.1\cdot 10^{-2}$
matrix + penalty ( $\lambda=0.1$ )	1.03	$1.9\cdot 10^{-3}$	$3.9\cdot 10^{-3}$	$1.7\cdot 10^{-3}$

Appendix

Table 3: Lorentz frame recovery from 48 noisy Faraday bivectors in

\mathrm{Cl}(1,3)

(400 Adam steps; noise floor

\sigma^{2}=3.36\cdot 10^{-4}

). Violation is

\lVert\Lambda^{\top}\eta\Lambda-\eta\rVert_{F}

for the vector representation

\Lambda

of the recovered transformation;

\Lambda

error is

\min_{\pm}\lVert\Lambda\mp\Lambda^{\star}\rVert_{F}/\lVert\Lambda^{\star}\rVert_{F}

; rapidity error is against the ground truth

0.7957

Automatic Differentiation

Automatic differentiation computes derivatives by tracking how numerical operations depend on their inputs, enabling efficient gradient-based learning and scientific inference.

See all 33 papers

Oct 8, 2026

asdex: Automatic Sparse Differentiation in JAX

Many tasks in scientific computing and machine learning require the Jacobian or Hessian matrix of a function. Automatic differentiation (AD) computes these derivatives to machine precision, but materializing a dense

m \times n

Jacobian requires

n

forward-mode or

m

reverse-mode AD passes, one per column or row. For a large class of functions, each output depends on only a few inputs, making the derivative matrix sparse. Automatic sparse differentiation (ASD) exploits this structure in four steps: detection of the input-agnostic sparsity pattern, coloring of a graph to group columns or rows that can share an AD pass, compressed differentiation to compute a compressed derivative matrix with one AD pass per color, and finally decompression into the original sparsity pattern. The number of colors, and hence of AD passes, is often independent of the problem dimension: a banded Jacobian with

b

contiguous bands, for instance, only ever requires

b

colors, regardless of its size. asdex offers the first standalone ASD toolkit in the popular JAX ecosystem. With

\texttt{asdex.jacobian}

and

\texttt{asdex.hessian}

, it provides sparse drop-in replacements for

\texttt{jax.jacobian}

and

\texttt{jax.hessian}

Adrian Hill, Guillaume Dalle

BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany · Machine Learning Group, Technical University of Berlin, Berlin, Germany · LVMT, ENPC, Institut Polytechnique de Paris, Univ Gustave Eiffel, Marne-la-Vallée, France

Automatic Differentiation Structured Sparsity

Ask PDF

Mar 9, 2026

SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

Automatic differentiation (AD) frameworks such as JAX and PyTorch have enabled gradient-based optimization for a wide range of scientific fields. Yet, many "hard" primitives in these libraries such as thresholding, Boolean logic, discrete indexing, and sorting operations yield zero or undefined gradients that are not useful for optimization. While numerous "soft" relaxations have been proposed that provide informative gradients, the respective implementations are fragmented across projects, making them difficult to combine and compare. This work introduces SoftJAX and SoftTorch, open-source, feature-complete libraries for soft differentiable programming. These libraries provide a variety of soft functions as drop-in replacements for their hard JAX and PyTorch counterparts. This includes (i) elementwise operators such as clip or abs, (ii) utility methods for manipulating Booleans and indices via fuzzy logic, (iii) axiswise operators such as sort or rank -- based on optimal transport or permutahedron projections, and (iv) offer full support for straight-through gradient estimation. Overall, SoftJAX and SoftTorch make the toolbox of soft relaxations easily accessible to differentiable programming, as demonstrated through benchmarking and a practical case study. Code is available at github.com/a-paulus/softjax and github.com/a-paulus/softtorch.

Anselm Paulus, A. René Geist, Vít Musil +3

University of Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Masaryk University, Czechia +1

Straight-Through Estimator Differentiable Programming Automatic Differentiation

Ask PDF

May 4, 2026

Exact Higher-Order Derivatives for SE(3) via Analytical/AD Methods

Fast prototyping of new SE(3) estimation objectives remains awkward in practice. Modern Lie-group frameworks -- GTSAM, manif, Sophus, SymForce, Ceres -- target first-order workloads through different code-generation and automatic-differentiation strategies, each optimized for a particular seam between hand-derived geometry and generic differentiation. The remaining gap is a compact, AD-safe path from these first-order primitives to exact Hessians, observed-information matrices, and higher-order derivative tensors: the quantities needed for exact Newton steps, observed-information covariance estimates, and covariance correction. This paper presents a hybrid analytical/AD recipe for SE(3) negative log-likelihoods. The practitioner writes the NLL gradient once, generic over a scalar type, and places the analytical/AD seam at the point-action interface y = Tx. Closed-form Lie-group Jacobians are used up to this interface; AD is applied only beyond it. The same source is then instantiated with ordinary floating-point scalars for gradients, vector-seeded dual numbers for exact Hessians in a single forward-mode pass, and nested dual numbers for higher-order derivative tensors. On a representative 6-DoF, 5-landmark SE(3) NLL, the advocated seeded-Hessian path is approximately 5x faster than finite-differencing the AD gradient on this benchmark while matching a nested-AD oracle to machine precision. The implementation adds roughly 70 lines of analytical-Jacobian code over an AD-only baseline. We also identify and fix a removable singularity in the standard SO(3)/SE(3) scalar basis that would otherwise produce NaNs at the origin under seeded AD, and we audit which Lie-group derivative tensors require this stabilized basis. The result is a practical path from rapidly written SE(3) objectives to exact higher-order derivatives, with predictable runtime and no finite-difference tuning.

Frank O. Kuehnel

Sigma Point Labs, Denver, CO, USA

Automatic Differentiation Lie Group Methods Second-Order Optimization Differentiable Optimization

Ask PDF

Abstract

Automatic Differentiation

Explore similar work

asdex: Automatic Sparse Differentiation in JAX

SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

Exact Higher-Order Derivatives for SE(3) via Analytical/AD Methods