---
title: Differentiable Boltzmann Solver
url: https://www.emergentmind.com/topics/differentiable-boltzmann-solver
type: topic
---

# Differentiable Boltzmann Solver

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 [2311.03291] [2503.23657] [2010.13649] [2509.24740]. 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 [2311.03291]. 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 [2311.03291]. 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 [2509.24740].

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 [2503.23657]. 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 [2010.13649].

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 [2311.03291]. 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)$ [2311.03291]. 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 [2509.24740].

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 $T_L$, and with heat flux obtained from directional integrals over the energy distribution [2503.23657]. The implementation includes isothermal, diffusely reflecting, and specularly reflecting boundary conditions, enforced per angle and band in the discrete solver [2503.23657].

The neural-network-accelerated Boltzmann solver works with the classical gas-kinetic equation
$$
\partial_t f + \mathbf v\cdot \nabla_{\mathbf x} f = Q(f,f),
$$
but replaces direct evaluation of the fivefold collision integral by a universal Boltzmann equation,
$$
\partial_t f + \mathbf v\cdot\nabla_{\mathbf x} f
=
\nu(\mathcal M-f)+\mathrm{NN}_\theta(\mathcal M-f),
$$
where the neural network acts on the deviation from equilibrium [2010.13649]. In the reported homogeneous and shock-structure settings, the solver is formulated in reduced variables such as $h(t,u)$ or $(h,b)$ and trained on time series generated by reference solvers [2010.13649].

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 [2311.03291] [2503.23657] [2010.13649].

## 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 $k$-modes or cosmologies [2311.03291]. The stiff Einstein–Boltzmann system is solved with Diffrax’s implicit Kvaerno5 ESDIRK method, which requires Jacobians; automatic differentiation supplies these exactly and stably [2311.03291]. Because JAX sparse support is experimental, modes are currently integrated separately rather than assembled into a large block-diagonal system [2311.03291].

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 $J=\partial f/\partial u$, automatically separates background and perturbations, and prepares interpolants and line-of-sight integrals [2509.24740]. 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 [2509.24740]. 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 [2509.24740].

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 [2503.23657]. 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 $(\alpha,\lambda)$ [2503.23657]. 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 [2503.23657].

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 [2010.13649]. Here differentiability primarily serves supervised learning of the neural correction to the BGK-like baseline rather than parameter sensitivities of a fixed physical model [2010.13649].

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 [2311.03291] [2509.24740] [2503.23657] [2010.13649].

## 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 $P(k,a)$, velocity-divergence fields, redshift-space power spectrum multipoles, and correlation function multipoles, with derivatives taken with respect to
$$
\theta=(H_0,\Omega_m,\Omega_b,N_{\rm eff},m_\nu,T_{\rm CMB},Y_{\rm He},A_s,n_s,w_0,w_a,c_a^2)^\top
$$
[2311.03291]. The paper explicitly considers Jacobians such as $J_\theta:=\nabla_\theta \log P_m(k,a|\theta)$ at $a=1$ for all 12 parameters and across $(k,a)$ [2311.03291]. Because the solver directly evolves $\delta_m(k,a)$ and $\theta_m(k,a)$, 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 [2311.03291].

SymBoltz.jl exposes a similarly broad set of differentiable outputs. Any output $y(\theta)$—including background distances, transfer functions, $P(k)$, and $C_\ell$—is differentiable with respect to any set of inputs $\theta$, and the paper demonstrates forward-mode Jacobians of $C_\ell$ with respect to parameters such as $(\Omega_b,\Omega_c,h,A_s,n_s,\sum m_\nu,w_0,w_a,\dots)$ [2509.24740]. Fisher forecasts can therefore be assembled directly from exact derivatives without finite-difference step tuning [2509.24740].

JAX-BTE uses differentiability for inverse simulations. The objective can depend on the converged energy field and model parameters such as $\tau(\omega,p)$, $C(\omega,p)$, $v(\omega,p)$, boundary reflection coefficients, heat source intensity, or geometry parameters, and the gradient is obtained through an adjoint linear system [2503.23657]. 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 [2503.23657].

In the neural surrogate paper, the differentiable stack is used to train $\mathrm{NN}_\theta$ by minimizing supervised losses over time and phase-space grids [2010.13649]. 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 [2010.13649].

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 $\Lambda$CDM model extended with one massive neutrino and CPL dark energy [2311.03291]. With $\ell_{\max}=32$ and 5-point neutrino momentum integration, it matches CAMB and CLASS at the few per-mille level for baryons, CDM, $b+c$, and total matter over a wide $k$-range at $z=0$, and agreement remains excellent at $z=99$ [2311.03291]. Differences are larger for the extremely suppressed massive-neutrino spectra at $k\gtrsim 1\ {\rm Mpc}^{-1}$, but these do not materially affect total-matter statistics [2311.03291]. 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 $n_s$ that likely trace modeling differences [2311.03291].

SymBoltz.jl reports agreement with CLASS and CAMB for matter and CMB spectra at the $0.1$–$1\%$ level across scales, and derivatives $\partial C_\ell/\partial \theta$ also agree when CLASS precision and finite-difference steps are carefully tuned [2509.24740]. The package further emphasizes dramatic compression of model-definition complexity: the full $\Lambda$CDM model is reported as 277 lines in SymBoltz, compared to 27,721 lines across 10 files in CLASS [2509.24740].

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 [2503.23657]. In 3D FinFET examples, centerline temperatures agree with GiftBTE, and for a five-FinFET array the maximum temperature is reported as $492.9\ {\rm K}$ for JAX-BTE versus $493.4\ {\rm K}$ for GiftBTE, with approximately $1\%$ difference [2503.23657]. Performance benchmarks span roughly $1{\rm M}$ to $110{\rm M}$ degrees of freedom; on a single GPU, the solver reports 6–8× speedup at $1{\rm M}$ DoF and about 4× speedup at $25{\rm M}$ DoF relative to the compared CPU GiftBTE configurations [2503.23657].

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 [2010.13649]. It also reports approximately $33\times$ speedup and $97\%$ 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 [2010.13649].

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 [2311.03291]. 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 $O(1)$ speedups [2311.03291]. 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 [2311.03291].

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 [2509.24740]. 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 [2509.24740].

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 $100{\rm M}$ DoF, and transient simulations, multi-material interfaces, and more complex boundary physics are planned rather than implemented [2503.23657]. 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 [2503.23657].

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 [2010.13649]. Likewise, Galilean and rotational invariance are not enforced in the neural network, positivity is not imposed, and no discrete $H$-theorem is proved, although entropy decay is shown numerically for the homogeneous case [2010.13649].

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 [2010.13649] [2311.03291] [2503.23657] [2509.24740].

Source: https://www.emergentmind.com/topics/differentiable-boltzmann-solver