Differentiable Boltzmann Solver

• Differentiable Boltzmann solvers are computational frameworks that solve Boltzmann-type equations while intrinsically providing automatic sensitivity analysis via differentiation.
• They are implemented across diverse applications, including linear Einstein–Boltzmann cosmology, GPU-accelerated phonon transport, and neural-network-enhanced gas-kinetic models.
• These solvers enable efficient gradient-based inference, optimization, and model validation, integrating traditional numerical methods with modern machine learning techniques.

A differentiable Boltzmann solver is a computational framework for Boltzmann-type equations in which the forward solution and derivatives of outputs with respect to inputs are available through automatic differentiation or closely related adjoint constructions. In the supplied literature, the term covers at least three distinct but related implementations: differentiable linear Einstein–Boltzmann solvers for cosmology, including DISCO-DJ’s Einstein–Boltzmann module and SymBoltz.jl, and a differentiable solver for the phonon Boltzmann Transport Equation, JAX-BTE; it also includes a hybrid neural-network-enhanced Boltzmann equation in which the collision operator is augmented within a universal differential equation formulation (Hahn et al., 2023, Shang et al., 31 Mar 2025, Xiao et al., 2020, Sletmoen, 29 Sep 2025). Across these settings, the common objective is to retain a Boltzmann-level description while exposing exact or efficiently computed sensitivities for gradient-based inference, sampling, optimization, emulator training, Fisher forecasting, or inverse design.

1. Definition and conceptual scope

In cosmology, a differentiable Einstein–Boltzmann solver is a code that integrates the linearized Einstein–Boltzmann system for the background and linear perturbations of all relevant species, while simultaneously exposing automatic derivatives of all outputs with respect to all input parameters through automatic differentiation (Hahn et al., 2023). DISCO-DJ’s Einstein–Boltzmann module is implemented entirely in JAX, so any function of the solver’s outputs is differentiable with respect to cosmological parameters, and this provides, together with the forward solution, the tangent hyperplane in parameter space (Hahn et al., 2023). SymBoltz.jl presents the same idea in a symbolic-numeric Julia environment: any output, including transfer functions and spectra, is differentiable with respect to any input parameter through automatic differentiation (Sletmoen, 29 Sep 2025).

Outside cosmology, the same concept appears in transport theory. JAX-BTE is described as a GPU-accelerated, fully differentiable solver for the phonon BTE, implementing a steady-state, energy-based, single-mode relaxation time approximation formulation with band and angular discretization and finite-volume spatial discretization (Shang et al., 31 Mar 2025). The hybrid solver of “Using neural networks to accelerate the solution of the Boltzmann equation” is differentiable in a different sense: the transport solver, collision model, and neural network are embedded in a universal differential equation and trained through reverse-mode automatic differentiation via Zygote.jl and DiffEqFlux.jl (Xiao et al., 2020).

This suggests that “differentiable Boltzmann solver” is not a single algorithmic template. It is a class of solvers in which Boltzmann-type dynamics remain explicit, but sensitivities are elevated to first-class outputs.

2. Governing equations and problem classes

The cosmological implementations solve the linearized Einstein–Boltzmann equations in Fourier space. DISCO-DJ works in a spatially flat FLRW background with conformal time $\tau$, synchronous gauge, scalar metric perturbations parameterized by $h$ and $\eta$, and mode-by-mode integration on a $k$-grid after linearization (Hahn et al., 2023). Its system includes metric equations, CDM and baryon continuity and Euler equations, photon and neutrino Boltzmann hierarchies, a momentum-dependent hierarchy for massive neutrinos, and a clustering dark-energy fluid with CPL parameterization $w(a)=w_0+w_a(1-a)$ (Hahn et al., 2023). SymBoltz.jl instead currently uses the conformal Newtonian gauge, with scalar metric perturbations $\Phi$ and $\Psi$, and evolves the background with General Relativity while solving exact stiff ODEs for photons with polarization, massless neutrinos, massive neutrinos, CDM, baryons, and recombination microphysics (Sletmoen, 29 Sep 2025).

JAX-BTE addresses a different kinetic equation. It implements the steady-state energy-based phonon BTE under single-mode relaxation time approximation,

$\mathbf v_g \cdot \nabla e = \frac{e-e_0}{\tau}+Q,$

with energy conservation coupling the energy distribution $e$ and equilibrium energy $e_0$ through the lattice temperature $h$0, and with heat flux obtained from directional integrals over the energy distribution (Shang et al., 31 Mar 2025). The implementation includes isothermal, diffusely reflecting, and specularly reflecting boundary conditions, enforced per angle and band in the discrete solver (Shang et al., 31 Mar 2025).

The neural-network-accelerated Boltzmann solver works with the classical gas-kinetic equation

$h$1

but replaces direct evaluation of the fivefold collision integral by a universal Boltzmann equation,

$h$2

where the neural network acts on the deviation from equilibrium (Xiao et al., 2020). In the reported homogeneous and shock-structure settings, the solver is formulated in reduced variables such as $h$3 or $h$4 and trained on time series generated by reference solvers (Xiao et al., 2020).

A common misconception is that differentiability requires the same physical content across applications. The literature instead shows differentiable solvers for linear cosmological perturbations, steady-state phonon transport, and hybrid gas-kinetic collision modeling, all under the same computational principle of end-to-end gradient access (Hahn et al., 2023, Shang et al., 31 Mar 2025, Xiao et al., 2020).

3. Numerical realization and automatic differentiation

DISCO-DJ realizes differentiability through JAX and Diffrax. Both forward- and reverse-mode differentiation are available through Diffrax’s DirectAdjoint and BacksolveAdjoint options, while JAX’s XLA compilation, jit, and vmap are used to compile kernels and batch computations over $h$5-modes or cosmologies (Hahn et al., 2023). The stiff Einstein–Boltzmann system is solved with Diffrax’s implicit Kvaerno5 ESDIRK method, which requires Jacobians; automatic differentiation supplies these exactly and stably (Hahn et al., 2023). Because JAX sparse support is experimental, modes are currently integrated separately rather than assembled into a large block-diagonal system (Hahn et al., 2023).

SymBoltz.jl approaches the same problem from a symbolic-numeric direction. The modeler writes equations symbolically via ModelingToolkit/Symbolics, which are compiled into efficient numerical ODE functions; the compiler analytically generates the Jacobian $h$6, automatically separates background and perturbations, and prepares interpolants and line-of-sight integrals (Sletmoen, 29 Sep 2025). The package is described as approximation-free because it solves the full stiff Einstein–Boltzmann system at all times with implicit integrators, without switching approximations such as tight coupling approximation, ultra-relativistic fluid approximation, radiation streaming approximation, or Saha-switching in recombination (Sletmoen, 29 Sep 2025). By default it uses Rodas4P for the background and KenCarp4 for perturbations, with Newton iterations and Jacobian and LU-factorization reuse over several time steps (Sletmoen, 29 Sep 2025).

JAX-BTE uses a finite-volume method with upwind fluxes, Green–Gauss gradients, and a pseudo-time incremental iteration to stabilize the nonlinear coupling through the equilibrium temperature (Shang et al., 31 Mar 2025). After band and angular discretization, one linear system per band-angle pair is solved with a Bi-Conjugate Gradient method and a Jacobi preconditioner, batched across $h$7 (Shang et al., 31 Mar 2025). Its differentiability is not based on naive backpropagation through iterative solvers. Instead, explicit parts are differentiated with automatic differentiation, while the implicit linear solves use a discrete adjoint implemented as custom VJP/JVP rules (Shang et al., 31 Mar 2025).

The universal differential equation solver for the gas Boltzmann equation is fully differentiable through the time integrator and the neural network. Reverse-mode automatic differentiation is handled by Zygote.jl and DiffEqFlux.jl, and training-time integration uses solvers from DifferentialEquations.jl such as Tsitouras 5/4 (Xiao et al., 2020). Here differentiability primarily serves supervised learning of the neural correction to the BGK-like baseline rather than parameter sensitivities of a fixed physical model (Xiao et al., 2020).

These implementations show two main architectures. One is “differentiate the physical solver,” as in DISCO-DJ, SymBoltz.jl, and JAX-BTE. The other is “embed learnable structure in the kinetic equation,” as in the universal Boltzmann equation of the neural surrogate paper (Hahn et al., 2023, Sletmoen, 29 Sep 2025, Shang et al., 31 Mar 2025, Xiao et al., 2020).

4. Outputs, Jacobians, and scientific use

The defining feature of these solvers is that derivatives are operational outputs, not post-processing approximations. In DISCO-DJ, the main differentiated outputs include species transfer functions and power spectra $h$8, velocity-divergence fields, redshift-space power spectrum multipoles, and correlation function multipoles, with derivatives taken with respect to

$h$9

(Hahn et al., 2023). The paper explicitly considers Jacobians such as $\eta$0 at $\eta$1 for all 12 parameters and across $\eta$2 (Hahn et al., 2023). Because the solver directly evolves $\eta$3 and $\eta$4, it can be inserted into linear redshift-space distortion models, FFTLog-based correlation multipoles, and Alcock–Paczynski and Finger-of-God mappings in a fully differentiable pipeline (Hahn et al., 2023).

SymBoltz.jl exposes a similarly broad set of differentiable outputs. Any output $\eta$5—including background distances, transfer functions, $\eta$6, and $\eta$7—is differentiable with respect to any set of inputs $\eta$8, and the paper demonstrates forward-mode Jacobians of $\eta$9 with respect to parameters such as $k$0 (Sletmoen, 29 Sep 2025). Fisher forecasts can therefore be assembled directly from exact derivatives without finite-difference step tuning (Sletmoen, 29 Sep 2025).

JAX-BTE uses differentiability for inverse simulations. The objective can depend on the converged energy field and model parameters such as $k$1, $k$2, $k$3, boundary reflection coefficients, heat source intensity, or geometry parameters, and the gradient is obtained through an adjoint linear system (Shang et al., 31 Mar 2025). The reported inverse tasks include learning thin-film thickness from temperatures at 20 points and learning heat source intensity from temperature observations along a boundary (Shang et al., 31 Mar 2025).

In the neural surrogate paper, the differentiable stack is used to train $k$4 by minimizing supervised losses over time and phase-space grids (Xiao et al., 2020). The authors emphasize that the construction preserves the correct fluid dynamic limit at leading order, because the neural-network input vanishes at equilibrium, while the BGK part enforces mass, momentum, and energy conservation exactly (Xiao et al., 2020).

A plausible implication is that differentiable Boltzmann solvers redistribute computational effort. Instead of repeated finite-difference reruns for sensitivities, they compute forward states and tangent information in the same framework.

5. Validation, accuracy, and performance

DISCO-DJ is validated against CAMB and CLASS for a flat $k$5CDM model extended with one massive neutrino and CPL dark energy (Hahn et al., 2023). With $k$6 and 5-point neutrino momentum integration, it matches CAMB and CLASS at the few per-mille level for baryons, CDM, $k$7, and total matter over a wide $k$8-range at $k$9, and agreement remains excellent at $w(a)=w_0+w_a(1-a)$0 (Hahn et al., 2023). Differences are larger for the extremely suppressed massive-neutrino spectra at $w(a)=w_0+w_a(1-a)$1, but these do not materially affect total-matter statistics (Hahn et al., 2023). The paper also states that Disco-EB’s Fisher contours for the Euclid spectroscopic clustering case agree well with official forecasts, with minor discrepancies noted for $w(a)=w_0+w_a(1-a)$2 that likely trace modeling differences (Hahn et al., 2023).

SymBoltz.jl reports agreement with CLASS and CAMB for matter and CMB spectra at the $w(a)=w_0+w_a(1-a)$3–$w(a)=w_0+w_a(1-a)$4 level across scales, and derivatives $w(a)=w_0+w_a(1-a)$5 also agree when CLASS precision and finite-difference steps are carefully tuned (Sletmoen, 29 Sep 2025). The package further emphasizes dramatic compression of model-definition complexity: the full $w(a)=w_0+w_a(1-a)$6CDM model is reported as 277 lines in SymBoltz, compared to 27,721 lines across 10 files in CLASS (Sletmoen, 29 Sep 2025).

JAX-BTE is validated against analytical solutions and GiftBTE. In a 1D cross-plane thin film with 100 cells and 128 directions, the temperature matches analytical solutions and captures boundary temperature slip at 100 nm while approaching Fourier behavior at 1000 nm (Shang et al., 31 Mar 2025). In 3D FinFET examples, centerline temperatures agree with GiftBTE, and for a five-FinFET array the maximum temperature is reported as $w(a)=w_0+w_a(1-a)$7 for JAX-BTE versus $w(a)=w_0+w_a(1-a)$8 for GiftBTE, with approximately $w(a)=w_0+w_a(1-a)$9 difference (Shang et al., 31 Mar 2025). Performance benchmarks span roughly $\Phi$0 to $\Phi$1 degrees of freedom; on a single GPU, the solver reports 6–8× speedup at $\Phi$2 DoF and about 4× speedup at $\Phi$3 DoF relative to the compared CPU GiftBTE configurations (Shang et al., 31 Mar 2025).

The neural surrogate paper reports that, in the homogeneous case, the universal Boltzmann equation reproduces fast spectral method distributions at training times, shows excellent agreement at interpolated and extrapolated times, and exhibits monotonically decreasing entropy matching fast spectral method trends (Xiao et al., 2020). It also reports approximately $\Phi$4 speedup and $\Phi$5 memory reduction relative to the fast spectral method in the homogeneous setting, while the shock-structure study shows density, velocity, temperature, stress, and heat-flux profiles matching the Shakhov reference (Xiao et al., 2020).

The literature therefore does not present differentiability as a replacement for validation. Each solver is benchmarked against established analytical or numerical references, and differentiability is added without abandoning conventional accuracy checks.

6. Limitations, misconceptions, and directions

Differentiability does not eliminate modeling assumptions. DISCO-DJ remains a linear-regime, large-scale-structure Einstein–Boltzmann solver in synchronous gauge, with truncated multipole hierarchies, simplified differentiable recombination through RECFASTmini, and current sparse-structure limitations in JAX/Diffrax (Hahn et al., 2023). Tight coupling, ultra-relativistic fluid approximation, and radiation streaming approximation are implemented but generally disabled in the first release because on GPUs they did not yield $\Phi$6 speedups (Hahn et al., 2023). Planned additions include HyRec/CosmoRec, modified gravity, additional dark energy and dark matter models, early dark energy, dark radiation, lensing potential, and coupling to differentiable LPT and N-body solvers (Hahn et al., 2023).

SymBoltz.jl’s approximation-free label can be misunderstood if read too broadly. In the paper it specifically means solving the full stiff equations at all times without switching approximations such as TCA, UFA, RSA, or Saha-switching in recombination (Sletmoen, 29 Sep 2025). The framework still involves hierarchy truncation, line-of-sight quadrature, adaptive time stepping, and current restrictions to linear theory; nonlinear corrections, CMB lensing, and reverse-mode AD are not yet included (Sletmoen, 29 Sep 2025).

JAX-BTE is limited to single-mode relaxation time approximation in the current implementation; the full scattering matrix is not included, high memory footprint limits single-GPU runs to approximately $\Phi$7 DoF, and transient simulations, multi-material interfaces, and more complex boundary physics are planned rather than implemented (Shang et al., 31 Mar 2025). The paper also notes that the Heaviside logic in upwinding yields piecewise-linear dependence, so care is needed if extending to fully nonlinear collision kernels (Shang et al., 31 Mar 2025).

The neural surrogate solver makes a different set of trade-offs. The BGK part enforces conservation exactly, but the paper does not provide an explicit analytic enforcement for the neural-network residual; conservation is instead observed numerically in the reported tests (Xiao et al., 2020). Likewise, Galilean and rotational invariance are not enforced in the neural network, positivity is not imposed, and no discrete $\Phi$8-theorem is proved, although entropy decay is shown numerically for the homogeneous case (Xiao et al., 2020).

A recurrent misconception is that differentiability automatically guarantees physical structure. The cited work shows the opposite: conservation, invariances, positivity, entropy decay, approximation switching, and collision modeling remain separate design decisions. Differentiability mainly changes how sensitivities are computed and how solvers interface with inference, optimization, and machine learning workflows (Xiao et al., 2020, Hahn et al., 2023, Shang et al., 31 Mar 2025, Sletmoen, 29 Sep 2025).