Solving stiff dark matter equations via Jacobian Normalization with Physics-Informed Neural Networks

TL;DR

This paper introduces a Jacobian normalization technique for Physics-Informed Neural Networks to effectively solve stiff differential equations, including complex dark matter equations, improving accuracy and convergence.

Contribution

The authors propose a Jacobian-based normalization method that enhances PINNs' ability to solve stiff equations without additional hyperparameters, validated on benchmark and real dark matter problems.

Findings

Achieves higher accuracy than previous methods on stiff ODEs.

Successfully solves the stiff Boltzmann equations for dark matter.

Accurately infers dark matter properties from limited data.

Abstract

Stiff differential equations pose a major challenge for Physics-Informed Neural Networks (PINNs), often causing poor convergence. We propose a simple, hyperparameter-free method to address stiffness by normalizing loss residuals with the Jacobian. We provide theoretical indications that Jacobian-based normalization can improve gradient descent and validate it on benchmark stiff ordinary differential equations. We then apply it to a realistic system: the stiff Boltzmann equations (BEs) governing weakly interacting massive particle (WIMP) dark matter (DM). Our approach achieves higher accuracy than attention mechanisms previously proposed for handling stiffness, recovering the full solution where prior methods fail. This is further demonstrated in an inverse problem with a single experimental data point - the observed DM relic density - where our inverse PINNs correctly infer the cross section that solves the BEs in both Standard and alternative cosmologies.