perf: Restructure JAX objective as end-to-end differentiable loss (JAX-ReaxFF pattern)

[ericchansen]

Summary

Restructure the JAX objective evaluation to follow the JAX-ReaxFF pattern: a single jax.jit-compiled, jax.vmap-vectorized loss function that evaluates all molecules in one kernel launch, enabling end-to-end differentiation from parameters → loss.

Motivation

GPU benchmarks show CPU is 1.6x faster than GPU for q2mm's current workloads (see GPU Acceleration docs). Research into how JAX-ReaxFF achieves 10--100x GPU speedup identified three architectural differences that explain the gap:

Current q2mm architecture (GPU-unfriendly)

1. Per-molecule Python loop -- ObjectiveFunction._compute_residuals() iterates over molecules in Python, calling engine.hessian() once per molecule. Each call is a separate GPU kernel launch.
2. No cross-molecule batching -- Hessians for 9 structures are computed in 9 separate kernels rather than one vmapped call.
3. Scipy drives the optimizer -- Each scipy.optimize.minimize call invokes the objective as a black-box function, preventing JAX from tracing through the full optimization step.

JAX-ReaxFF architecture (GPU-friendly)

1. vmap over all geometries -- jax.vmap(calculate_energy_and_charges, in_axes=(0,0,0,None)) evaluates all molecules in a single kernel launch.
2. End-to-end differentiable loss -- One JIT-compiled function from parameters to total loss, enabling jax.grad(loss)(params) for analytical gradients.
3. JAX-native optimization -- Gradient-based optimizers (L-BFGS) run inside JAX's traced computation graph.

Sources:

• JAX-ReaxFF paper (ChemRxiv)
• JAX-ReaxFF optimizer.py
• JAX GPU Performance Tips
• JAX Benchmarking Guide -- shows GPU is 10x slower than CPU for small matrices

Proposed Changes

Phase 0: Investigate float32 viability (quick win)

The current code forces float64 globally (_jax_common.py sets jax_enable_x64=True), which limits the RTX 5090 to 1.6 TFLOPS instead of its full 104.8 TFLOPS -- a 64x penalty.

Analysis shows float64 may not be necessary for the current harmonic-only JaxEngine:

• The unit conversion (kcal/mol/A^2 to Hartree/Bohr^2) is linear -- it does not change relative precision. Float32's ~7 significant digits give the same accuracy at any scale.
• For harmonic terms (E = k(r-r0)^2), the Hessian is exactly 2k -- no catastrophic cancellation in the autodiff tape.
• To resolve 1 cm^-1 at a 1000 cm^-1 mode, you need ~8% relative precision in the Hessian element. Float32 provides 1.2e-7 relative precision -- 5 orders of magnitude more than needed.

Float64 would become necessary for:

• VdW terms (1/r^12 - 1/r^6 near-cancellation in autodiff)
• Morse potentials, cross-terms
• Very soft modes (~10 cm^-1) near zero eigenvalues

Action: Run CH3F and rh-enamide benchmarks with jax_enable_x64=False and compare frequency accuracy against float64 baseline. If frequencies agree to <0.1 cm^-1, float32 is viable and unlocks 64x more GPU throughput on consumer hardware.

Phase 1: Batch Hessians across molecules

• Pad molecules to uniform atom count
• Use jax.vmap(hessian_fn) to compute all Hessians in one kernel launch
• Estimated improvement: 2--3x for multi-molecule systems

Phase 2: End-to-end differentiable frequency loss

• Create a single JIT-compiled function: params -> energy_fn -> hessian -> eigenvalues -> frequencies -> residuals -> loss
• Use jax.grad(loss) for analytical parameter gradients (replacing finite-difference in GRAD steps)
• This eliminates the Python loop and enables full GPU kernel fusion

Phase 3: JAX-native optimizer

• Replace scipy.optimize.minimize with a JAX-traced optimizer (e.g., JAXopt L-BFGS)
• The entire optimization loop becomes a single compiled computation
• This is the step that enables true GPU acceleration

Additional Considerations

• Molecule padding overhead -- vmap requires uniform array shapes. Padding small molecules (5 atoms) to match large ones (62 atoms) wastes compute. Grouping molecules by size or using JAX's dynamic shapes could mitigate this.
• Backward compatibility -- The existing ObjectiveFunction API should continue to work for non-JAX engines (OpenMM, Tinker). The JAX-native path would be an alternative code path activated when using JaxEngine.
• Datacenter GPUs -- Even with all optimizations, an A100 (FP64:FP32 = 1:2, 9.7 TFLOPS FP64) would be far more suitable than the RTX 5090 (1:64) for float64-heavy workloads.

Related Issues

• Evaluate Pint for unit handling: performance impact assessment #161 (Evaluate Pint for unit handling) -- Pint lives at the FF/engine boundary and cannot change what happens inside JAX-traced code. GPU performance work should be done first; Pint is orthogonal.

Success Criteria

• Float32 viability tested and documented (Phase 0)
• Multi-molecule frequency evaluation uses vmap (no Python loop)
• Analytical parameter gradients via jax.grad for frequency objective
• GPU benchmark shows speedup over CPU for rh-enamide (9 molecules, 94 params)
• All existing tests continue to pass
• Benchmark results documented with reproducible commands

[ericchansen]

Update: Phase 0 (float32) complete — not viable

PR #186 confirms float32 is not viable for molecules >30 atoms (Hessian eigendecomposition precision). See #178 for full results.

Remaining phases for GPU acceleration:

• Phase 1: Batch Hessians with jax.vmap (perf: vmap over molecules in objective function #180) — pad molecules to uniform atom count, consolidate N kernel launches into 1
• Phase 2: End-to-end JIT loss — single compiled function from parameters → objective score
• Phase 3: JAX-native optimizer — replace scipy with JAXopt for on-device optimization

Phase 1 is the most impactful next step. JAX-ReaxFF research (#179) confirms these patterns yield 10–100× GPU speedup.

[ericchansen]

Status update — April 2026

Updating with current implementation state. Several phases are further along than the issue body reflects.

Phase status

Phase 0 — Float32 viability: ✅ Complete. #178 closed. Not viable for molecules >30 atoms (Hessian eigendecomposition precision). Documented in docs/benchmarks/gpu.md.

Phase 1 — Batch Hessians across molecules: ✅ Complete. #180 closed. Implemented in q2mm/backends/mm/batched.py:

• group_by_topology() groups molecules by connectivity
• batched_hessians() uses jax.jit(jax.vmap(hess_fn)) over coordinate batches
• batched_frequencies() built on top
• ObjectiveFunction._precompute_batched_hessians() integrates it into the evaluation pipeline
• Unit tests in test/test_batched_eval.py

Phase 2 — End-to-end differentiable frequency loss: 🟡 Partially complete. The analytical gradient chain works but is NOT a single JIT-compiled function:

• FrequencyEvaluator.gradient() computes d(freq_score)/d(params) via eigenvalue sensitivity + hessian_and_param_jacobian()
• EigenmatrixEvaluator.gradient(), HessianElementEvaluator.gradient(), EnergyEvaluator.gradient() — all operational
• ObjectiveFunction.gradient() orchestrates per-molecule, per-evaluator analytical gradients
• ScipyOptimizer supports jac=‘analytical’ to use this gradient pipeline with L-BFGS-B
• What’s missing: the objective is still a Python function calling into JAX per molecule — not a single jax.jit(params → loss) function

Phase 3 — JAX-native optimizer: ❌ Not started. Still using scipy.optimize.minimize.

Success criteria checklist

• Float32 viability tested and documented (Phase 0) — perf: Float32 viability experiment for JaxEngine harmonic force fields #178
• Multi-molecule frequency evaluation uses vmap (no Python loop) — batched_hessians() via jax.vmap
• Analytical parameter gradients via jax.grad for frequency objective — FrequencyEvaluator.gradient()
• GPU benchmark shows speedup over CPU for rh-enamide — 5.6× JAX-MD, 2.1× JAX harmonic (docs/benchmarks/gpu.md)
• All existing tests continue to pass
• Benchmark results documented with reproducible commands — docs/benchmarks/gpu.md

Remaining work

The listed success criteria are met. Remaining architectural improvements:

1. Compose the analytical gradient chain into a single jax.jit-compiled params → loss function (Phase 2 completion)
2. Replace scipy optimizer with JAXopt L-BFGS (Phase 3)

These are significant but incremental — the component-level primitives all exist.