Solving High-Dimensional Population Balance Equations via Dynamics-Preserving Autoencoders

High-dimensional population balance equations (PBEs) provide a natural framework for modeling heterogeneous cell populations, but their direct numerical solution becomes computationally prohibitive when the internal state space contains many molecular variables. We propose a hybrid mechanistic-machine learning framework for reducing and simulating PBEs defined over high-dimensional intracellular coordinates. The cell population is described by a number density n(x, t), where x [isin] [R]N represents gene and protein states associated with macrophage activation. A dynamics-preserving autoencoder maps this state space to a low-dimensional latent coordinate z [isin] [R]d, with d << N, while retaining key qualitative features of the underlying gene regulatory network, including attractor structure and multistability. Mechanistic information from the original regulatory dynamics is used to construct interpretable drift and diffusion terms for the reduced latent-space PBE. The reduced PBE is solved using a stochastic Lagrangian particle representation, in which particles evolve according to stochastic differential equations (SDEs) corresponding to the latent drift and diffusion fields. The resulting latent-space solution is subsequently decoded and propagated back into the original state space to recover physically interpretable cellular dynamics. We demonstrate the framework on macrophage polarization under cytokine-dependent regulation, including gene knockout perturbations. Overall, the proposed framework provides a computationally tractable and mechanistically interpretable route for integrating single-cell genomic data with population balance models of cell-state dynamics.