Generative diffusion model surrogates for mechanistic agent-based biological models

TL;DR

This paper tackles the computational burden of stochastic Cellular-Potts models by building a class-conditional DDPM surrogate for a vasculogenesis CPM and verifying generated configurations with an image classifier. The approach enables long-horizon inference ($$20{,}000\mathrm{MCS}$$) across a discretized 25-class parameter space, achieving about a $$22\times$$ reduction in runtime. The authors demonstrate surrogate selection via classifier-driven validation and morphology-aware metrics (EMD), obtaining representative configurations while noting limitations such as mode collapse and between-class interpolation. This work advances digital-twin-inspired surrogates for stochastic biological systems and points toward future enhancements like diffusion-guided generation and online updating.

Abstract

Mechanistic, multicellular, agent-based models are commonly used to investigate tissue, organ, and organism-scale biology at single-cell resolution. The Cellular-Potts Model (CPM) is a powerful and popular framework for developing and interrogating these models. CPMs become computationally expensive at large space- and time- scales making application and investigation of developed models difficult. Surrogate models may allow for the accelerated evaluation of CPMs of complex biological systems. However, the stochastic nature of these models means each set of parameters may give rise to different model configurations, complicating surrogate model development. In this work, we leverage denoising diffusion probabilistic models to train a generative AI surrogate of a CPM used to investigate in vitro vasculogenesis. We describe the use of an image classifier to learn the characteristics that define unique areas of a 2-dimensional parameter space. We then apply this classifier to aid in surrogate model selection and verification. Our CPM model surrogate generates model configurations 20,000 timesteps ahead of a reference configuration and demonstrates approximately a 22x reduction in computational time as compared to native code execution. Our work represents a step towards the implementation of DDPMs to develop digital twins of stochastic biological systems.

Figures (7)

• Figure 1: Cellular-Potts Model (CPM) parameters dictate model configuration over long timescales. (A) Reference configuration from a CPM assumes distinct configurations (B) due to changes in mechanistic model parameters over a long timescale of 20,000 simulation timesteps (Monte-Carlo steps, MCS). Memory of the initial configuration is lost but distinct patterns emerge unique to the area of the parameter space.
• Figure 2: Denoising Diffusion Probabilistic Model (DDPM) training and sample generation. A sample Cellular-Potts model configuration is used for training a DDPM model (top). The trained model may be used to quickly generate a representative model configuration sampled from the same class, or area of the parameter space (bottom).
• Figure 3: Image classifier identifies mode collapse in trained candidate DDPM surrogate models. DDPM surrogate model performance is sensitive to training hyperparameters P_mean and P_std. Each candidate model was used to generate 20,000 configurations uniformly distributed among the 25 classes, or discrete areas of the CPM parameter space. (A) Generated outputs were classified using a trained image classifier and sample counts per class are displayed as a heatmap in the equivalent 5x5 grid layout shown previously in Figure \ref{['fig:params']}B. Sample counts per class (or discrete parameter set) are displayed within each corresponding cell of each heatmap subplot. KL divergence between generated sample distributions and a uniform distribution is displayed for each DDPM model to provide a quantitative measure of uniformity. KL divergence values close to 0 indicate a more uniform distribution among classes in generated samples. (B) Confusion matrix plots demonstrating classifier identified class (detected class) compared with the conditioning label for generation (intended class) for each model. Generation accuracy is displayed for each DDPM model.
• Figure 4: DDPM surrogate generates representative CPM configurations. The selected DDPM model surrogate generates visually representative CPM configurations throughout the parameter space of interest. Ground truth from representative CPM simulations at t = 20,000 MCS for all areas of the defined parameter space (left) is compared to surrogate model generated configurations (right).
• Figure 5: Surrogate model evaluates significantly faster than native CPM code execution. P=6.31e-39 by independent samples t-test for n=25 simulation configurations each. Mean DDPM surrogate model evaluations were 20.6s (std 4.0s) as compared with a mean of 447.6s (std 51.1s) for native CompuCell3D (CC3D) code execution. The surrogate model achieved approximately a 22x speed increase. Surrogate model evaluations were performed on a single node of 4 V100 GPUs, time recorded included loading the model into GPU memory and saving generated samples. CC3D simulation execution was performed on an HPC node consisting of AMD EPYC 9495 architecture.