Reversible random number generation for adjoint Monte Carlo simulation of the heat equation

TL;DR

The paper tackles the memory bottleneck in PDE-constrained optimization using adjoint methods when employing Monte Carlo solvers, where identical particle trajectories are needed for constraint and adjoint simulations. It introduces a reversible extension to the PCG pseudorandom number generator, with a strategy to reversibly generate distributions via accept-reject sampling and, when necessary, a secondary generator to supply additional random values without correlative bias. The approach is validated on a 1D heat equation, showing that constraint path recomputation can replace storing full trajectories while preserving similar computational costs, thereby mitigating memory blow-up. This reversible RNG framework offers a practical alternative to checkpointing, enabling scalable adjoint MC methods for larger and more complex PDEs in areas such as engineering design and fusion research.

Abstract

In PDE-constrained optimization, one aims to find design parameters that minimize some objective, subject to the satisfaction of a partial differential equation. A major challenges is computing gradients of the objective to the design parameters, as applying the chain rule requires computing the Jacobian of the design parameters to the PDE's state. The adjoint method avoids this Jacobian by computing partial derivatives of a Lagrangian. Evaluating these derivatives requires the solution of a second PDE with the adjoint differential operator to the constraint, resulting in a backwards-in-time simulation. Particle-based Monte Carlo solvers are often used to compute the solution to high-dimensional PDEs. However, such solvers have the drawback of introducing noise to the computed results, thus requiring stochastic optimization methods. To guarantee convergence in this setting, both the constraint and adjoint Monte Carlo simulations should simulate the same particle trajectories. For large simulations, storing full paths from the constraint equation for re-use in the adjoint equation becomes infeasible due to memory limitations. In this paper, we provide a reversible extension to the family of permuted congruential pseudorandom number generators (PCG). We then use such a generator to recompute these time-reversed paths for the heat equation, avoiding these memory issues.

Figures (5)

• Figure 1.1: A sketch of the Ziggurat accept-reject algorithm. Each horizontal light-blue region has the same area, while dark-blue regions mark the range of values for $\phi_k$ which are always accepted.
• Figure 1.2: Comparing forward and reverse mode timings for different random distributions. We see that both modes attain the same performance for all three considered distributions.
• Figure 1.3: The temperature profile and optimal cooling rate for the 1D heat equation problem.
• Figure 1.4: The increase in runtime and memory usage of a reversible and non-reversible constraint and corresponding adjoint simulations as the number of particles increases. As memory requirements approach the available RAM on the system (16GB), the non-reversible simulation's runtime increases drastically, while the reversible implementation maintains its linear scaling.
• Figure 1.5: Demonstration of the challenge of reflective boundaries. Both particles $X_p$ and $X_{p^\prime}$ have the same position at time $\tau+1$ after moving the same distance to the right, even though they have different positions at time $\tau$.