State-dependent temperature control in Langevin diffusions using numerical exploratory Hamiltonian-Jacobi-Bellman equations

TL;DR

This paper develops a neural network-based method to solve Hamilton-Jacobi-Bellman equations for state-dependent noise control in Langevin dynamics, improving optimization in high-dimensional problems.

Contribution

It introduces principled control bounds and a physics-informed neural network framework to stabilize and accurately solve exploratory HJB equations in high dimensions.

Findings

Method remains robust in high-dimensional optimization tasks.

Stable and accurate estimation of state-dependent noise achieved.

Outperforms existing approaches in challenging problems.

Abstract

Choosing how much noise to add in Langevin dynamics is essential for making these algorithms effective in challenging optimization problems. One promising approach is to determine this noise by solving Hamilton-Jacobi-Bellman (HJB) equations and their exploratory variants. Though these ideas have been demonstrated to work well in one dimension, extension to high-dimensional minimization has been limited by two unresolved numerical challenges: setting reliable control bounds and stably computing the second-order information (Hessians) required by the equations. These issues and the broader impact of HJB parameters have not been systematically examined. This work provides the first such investigation. We introduce principled control bounds and develop a physics-informed neural network framework that embeds the structure of exploratory HJB equations directly into training, stabilizing computation, and enabling accurate estimation of state-dependent noise in high-dimensional problems. Numerical experiments demonstrate that the resulting method remains robust and effective well beyond low-dimensional test cases.