research: Evaluate JAX-ReaxFF integration as a q2mm backend

[ericchansen]

Summary

Evaluate integrating JAX-ReaxFF as a q2mm backend or adopting its architecture for q2mm's own engines. JAX-ReaxFF reports 10–100× GPU speedup for reactive force field fitting using a vmap-batched, JIT-compiled, gradient-based pipeline.

Motivation

GPU benchmark analysis identifies two non-exclusive paths to GPU acceleration:

1. Restructure q2mm's own objective into an end-to-end JIT-compiled params → loss function (tracked in perf: Restructure JAX objective as end-to-end differentiable loss (JAX-ReaxFF pattern) #176)
2. Integrate JAX-ReaxFF as a backend, gaining access to a mature GPU-optimised reactive force field engine

This issue tracks option 2. These are not mutually exclusive — JAX-ReaxFF integration could coexist with q2mm's own JIT-compiled loss pipeline.

What JAX-ReaxFF provides

• Reactive force field (ReaxFF) energy evaluation in pure JAX
• vmap over all training geometries in a single kernel launch
• Analytical gradients via jax.grad through the full energy computation
• Gradient-based L-BFGS optimisation inside JAX's traced computation graph
• float32 sufficient for their workloads
• Proven GPU speedups (10–100× over serial implementations)

Sources:

• JAX-ReaxFF paper (ChemRxiv)
• JAX-ReaxFF poster (Zenodo)
• Source code

Open questions

• Should JAX-ReaxFF be a new engine alongside JaxEngine, or should its architecture inform how we restructure q2mm's existing engines?
• JAX-ReaxFF only needs 1st-order gradients (no Hessians) — would q2mm's frequency fitting (which needs Hessians) still benefit from their architecture?
• Is ReaxFF force field fitting a use case q2mm should support? Or is the value mainly in learning from their architecture?
• How would this interact with the cycling loop (GRAD + SIMP)?

Related

• perf: Restructure JAX objective as end-to-end differentiable loss (JAX-ReaxFF pattern) #176 (restructure q2mm's own objective — the alternative/complementary approach)
• perf: Float32 viability experiment for JaxEngine harmonic force fields #178 (float32 viability — relevant since JAX-ReaxFF uses float32)
• perf: GPU acceleration for MM backends and optimization #155 (GPU acceleration umbrella)
• feat: MM3 functional forms in JAX for differentiable MM3 energy evaluation #91 (MM3 functional forms in JAX)

[ericchansen]

JAX-ReaxFF Evaluation

Summary

After evaluating JAX-ReaxFF's architecture, source code, and papers against q2mm's engine interface and use cases, the recommendation is: learn from the architecture, don't integrate as a backend.

What JAX-ReaxFF does well

JAX-ReaxFF achieves 10–100× GPU speedup through three key architectural decisions:

1. vmap over all geometries — all training structures evaluated in a single kernel launch
2. End-to-end JIT — one compiled function from parameters → loss, enabling full XLA fusion
3. jax.grad for analytical parameter gradients — eliminates finite-difference overhead

These are exactly the patterns q2mm needs to adopt (tracked in #176).

Why NOT to integrate as a backend

1. ReaxFF is a different problem domain. JAX-ReaxFF fits reactive force field parameters using energy/force matching. q2mm fits MM force field parameters (bonds, angles, torsions, vdW) using frequency matching — which requires Hessians (2nd derivatives), not just gradients. JAX-ReaxFF's vmap(calculate_energy_and_charges) doesn't compute Hessians at all.

2. No frequency fitting pathway. q2mm's core value proposition is Seminario estimation + frequency-based objective functions. A ReaxFF backend wouldn't support engine.hessian() or engine.frequencies() without major new development — essentially reimplementing ReaxFF's energy in a way that supports jax.hessian.

3. Force field format incompatibility. ReaxFF parameters (~100s of terms: bond order, overcoordination, lone pair, conjugation, hydrogen bonding, etc.) are fundamentally different from q2mm's ForceField model (bonds, angles, torsions, vdW with MM3/harmonic forms). Supporting ReaxFF would require a parallel parameter system.

4. Maintenance burden vs. user base. ReaxFF is a niche use case for q2mm's user base (organometallic chemistry with MM3/AMBER-style force fields). Adding a ReaxFF backend would increase codebase complexity without serving existing users.

What to adopt from JAX-ReaxFF's architecture

These patterns should be applied to q2mm's existing JAX engines:

Pattern	JAX-ReaxFF	q2mm today	q2mm target
Molecule batching	vmap over all geometries	Python loop	vmap over padded molecules (#180)
Loss compilation	Single jax.jit params→loss	Python loop + numpy eigendecomp	End-to-end JIT (#176 Phase 2)
Parameter gradients	jax.grad(loss)	Finite-difference / per-evaluator	jax.grad through eigendecomp (#176 Phase 2)
Optimizer	JAX-native L-BFGS	Scipy L-BFGS-B (black-box)	JAXopt L-BFGS (#176 Phase 3)
Precision	float32	float64 (required — see #178)	float64 (float32 not viable for >30 atoms)

Relationship to other issues

• perf: Restructure JAX objective as end-to-end differentiable loss (JAX-ReaxFF pattern) #176 already tracks the concrete implementation of JAX-ReaxFF patterns in q2mm's own engines
• perf: vmap over molecules in objective function #180 tracks the vmap over molecules specifically
• perf: Float32 viability experiment for JaxEngine harmonic force fields #178 confirmed float32 isn't viable, which differs from JAX-ReaxFF's experience (they don't need Hessian eigendecomposition)

Recommendation: Close as research-complete

The valuable insights from JAX-ReaxFF are already captured in #176's phased plan. No separate integration work is needed — the architectural patterns should be applied to q2mm's existing JAX/JAX-MD engines instead.

[ericchansen]

Closing as research-complete. Key findings captured above. Actionable patterns (vmap, end-to-end JIT, jax.grad) are tracked in #176. No separate ReaxFF integration needed.

[ericchansen]

Reopening — closing was premature. JAX-ReaxFF integration is a future goal. Research findings above document current architectural gaps; the issue should remain open for tracking.

[ericchansen]

Status update — April 2026

The research findings from this issue have been adopted into the codebase. Updating with current state:

Architectural patterns adopted from JAX-ReaxFF

Pattern	JAX-ReaxFF	q2mm status
jax.grad for parameter gradients	jax.grad(loss)	✅ energy_and_param_grad(), hessian_and_param_jacobian() via jax.jacrev(jax.hessian)
jax.vmap over geometries	vmap(calculate_energy)	✅ batched_hessians() in q2mm/backends/mm/batched.py (#180 closed)
jax.jit compilation	Single JIT params→loss	✅ Per-function JIT (_energy_fn, _coord_hess_fn, _hess_param_jac_fn); ❌ not yet end-to-end
Analytical parameter gradients	jax.grad(loss)	✅ All evaluators (frequency, eigenmatrix, hessian_element, energy) have gradient() methods
float32	Default	❌ Not viable for >30 atoms (#178 closed)
JAX-native optimizer	JAXopt L-BFGS	❌ Still using scipy

What this means for #179

The research recommendation — "learn from the architecture, don’t integrate as a backend" — was correct and has been executed. The three key patterns (vmap, analytical jax.grad, jax.jit) are all implemented at the component level. Remaining work (end-to-end JIT, JAX-native optimizer) is tracked in #176 and #152.