Towards Learned Simulators for Cell Migration
Anonymous Author(s)
AffiliationAddress email
Abstract
Simulators driven by deep learning are gaining popularity as a tool for efficientlyemulating accurate but expensive numerical simulators.
Successful applicationsof such neural simulators can be found in the domains of physics, chemistry,and structural biology, amongst others.
Likewise, a neural simulator for cellulardynamics can augment lab experiments and traditional computational methods toenhance our understanding of a cell’s interaction with its physical environment.
Inthis work, we propose an autoregressive probabilistic model that can reproducespatiotemporal dynamics of single cell migration, traditionally simulated with the
Cellular Potts model.
We observe that standard single-step training methods donot only lead to inconsistent rollout stability, but also fail to accurately capture thestochastic aspects of the dynamics, and we propose training strategies to mitigatethese issues.
Our evaluation on two proof-of-concept experimental scenariosshows that neural methods have the potential to faithfully simulate stochasticcellular dynamics at least an order of magnitude faster than a state-of-the-artimplementation of the Cellular Potts model.
Introduction
Studying the variety of mechanisms through which cells migrate and interact with their physicalenvironment is of crucial importance for our understanding of cell biology.
For example, cellmigration plays a key role in the interaction between the immune system and implant surfaces [7, 30],the development of embryos [24] and the progression of cancer [27, 17].
As experimental capacityin the lab is inherently limited, computational methods have emerged as a tool to investigate thestochastic and dynamic movement and shape of cells.
However, these methods can be computationallydemanding.
This is especially restrictive in scenarios requiring many simulations, for example due tosubstantial stochasticity or in the case of inverse design where parameters are optimized iterativelybased on the simulator’s output.
Moreover, parameterizing such models to realistically simulate cellscan be a difficult task, requiring careful design and expert knowledge.
On the other hand, deep learning has been gaining traction as a tool for learning fast approximatesimulators.
For example, for continuous-time and continuous-space systems defined with partialdifferential equations (PDEs), neural solvers learn to emulate a system’s dynamics from a dataset ofsimulations generated by a more computationally demanding solver [18, 28, 4, 12, 34].
In this setting,a large computational cost is paid up front to generate the training set, but once trained, approximatesolutions can be generated at a fraction of the original cost.
Moreover, learned simulators hold thepromise of emulating systems for which the laws governing the dynamics are not known, by insteadtraining on experimental observations.
Based on the above considerations, we propose to use neural simulators to simulate cellular dynamics.
More specifically, we consider the scenario where both the movement and shape of the cell showstochastic aspects and are highly dynamic, which is typically modeled in the Cellular Potts modelingframework, proposed in [10].
Given their various successful applications in modeling spatiotemporaldata, we hypothesize that neural simulators are capable of faithfully emulating the ground truthdynamics, while accelerating the simulation process.
Our contributions are summarized as follows:
• We propose a neural simulation model to simulate stochastic single-cell dynamics similar tothose generated by the Cellular Potts model;
• We develop and evaluate autoregressive training strategies, with the aim to improve themodel’s rollout performance and its ability to capture stochastic dynamics;
• We observe that our method has the capacity to faithfully emulate the cellular dynamics of
the Cellular Potts model, while generating simulations an order of magnitude faster.
Background and Related Work2.1 Cellular Potts Model
The Cellular Potts (CP) model is a computational modeling framework for simulating cellulardynamics and the dynamic and fluctuating morphology of cells on a lattice [10, 23, 1].
The CP modelhas gained prominence due to its flexibility in modeling cell shape and movement, the interactionbetween multiple cells, stochastic aspects of cell behavior, and multiscale mechanisms [25, 20, 11].
In the CP framework, the system is modeled as a Euclidean lattice L and Hamiltonian H .
Thefunction x : L → S maps each lattice site l i ∈ L to its state x ( l i ) ∈ S , where S is the set of allcells and materials that can be present in the system.
Note that in the CP literature x is commonlyreferred to as σ ; we deviate from this to stick to machine learning convention.
To evolve the system,a Markov-Chain Monte Carlo sampling algorithm is used.
At every iteration, a lattice site l i is chosenat random.
Then, a proposal is made to modify x such that state x ( l i ) is changed to x ( l j ) , where l j isa site adjacent to l i .
Finally, the difference in energy ∆ H is calculated between the proposed andcurrent system state.
If ∆ H ≤ 0 , the proposed state is accepted as the new system state; if ∆ H > 0 ,it is accepted with probability e − ∆ HT , with T being the temperature parameter of the model.
The Hamiltonian H itself differs per application, but typically consists of at least contact energy andvolume preservation terms, as originally proposed in [10]:
[Equation]
where N ( L ) is the set of all pairs of neighboring lattice sites in L , J ( x ( l i ) , x ( l j )) is the contactenergy between cells and/or materials x ( l i ) and x ( l j ) , and δ x,y is the Kronecker delta.
Furthermore, Cis the set of all cells in the system, V ( c ) is the number of lattice sites occupied by cell c (from here onreferred to as the volume of cell c ), V ∗ ( c ) is the target volume of cell c , and λ V is a Lagrange multiplier.
H other can consist of many extensions and modifications of the original Hamiltonian, for exampletaking into account cellular dynamics induced by forces, gradients in chemical concentrations, cellsurface area constraints, and many more biological concepts.
The specific Hamiltonians used forsimulating our data can be found in Appendix A.
2.2 Neural Simulators
Neural networks have been employed for simulation in many domains [19, 5], often by eithercombining ML models with existing numerical solvers [32, 15] or by using ML models to simulatedynamics in their entirety [4, 18, 28].
The latter, which we refer to as neural simulators , encompassthe type of model proposed in this work, as we seek to emulate the CP simulations as a whole.
Of particular interest are autoregressive methods operating on a spatial grid, as these fit both the temporaland spatial component of the CP simulations.
This setup generally comes with challenges of ensuringprediction quality and stability over longer rollout trajectories.
Common approaches to address thisinclude injecting noise and incorporating model rollouts in the training procedure [28, 4].
In the context of cellular dynamics, neural networks have been used in various settings to aid insimulation.
TrajectoryNet [31] utilizes optimal transport in combination with continuous normalizingflows to interpolate time-evolving gene expression data based on population measurements at fixedtimepoints.
In a similar spirit, CellBox [35, 13] simulates molecular biological processes with neural
ODEs [6].
LEUP [2] models cells as atomic particles moving in two-dimensional space to investigatecollective cell migration.
However, these methods do not address simulation of spatiotemporaldynamics of cells on a grid.
Another branch of related research focuses on grid-based generativemodeling of cell morphology and subcellular organization [8, 9, 22].
Very recently, Wiesner et al.
[33]proposed a method to model time evolving cell shapes.
Although these methods can impressivelyperform (conditional) cell image generation, they do not address migration dynamics.
Method3.1 Machine Learning Formulation
We consider the problem of learning an autoregressive probabilistic mapping to simulate spatiotemporal cellular dynamics.
More formally, at each time t the system is described by its state x t , acategorically-valued function on a fixed grid (see also Section 2.1).
Correspondingly, a system evolution is specified by a sequence of states x 0: t .
We postulate a ground-truth probability distributionp ∗ ( x 0: t ) over system evolutions x 0: t from which we can sample using the CP simulator.
Hence, ouraim is to learn the parameters θ of the model p θ such that
[Equation]
is maximized.
By the Markov property of the data generating process of the CP model configurationsthat we consider, and with p ∗ ( x 0 ) a specified distribution of initial conditions, this is equivalent to
[Figure]
maximizing Equation 3 with respect to θ :
[Figure]
[Equation]
Consequently, we aim to model p ∗ ( x 0: t ) by learning p θ ( x t +1 | x t ) and applying it autoregressively.
[Figure]
3.2 Model Design
[Figure]
The design choices for modeling p θ ( x t +1 | x t ) are driven by the goals to produce a model that is capa-
[Figure]
ble of producing realistic trajectories with good sample efficiency.
To achieve high sample efficiency,
[Figure]
we introduce latent variable z , following the structure of a Conditional Variational Autoencoder(CVAE)
In our case, latent variable z is conditioned on previous state x t according to conditionalprior p θ ( z | x t ) , and the subsequent state is specified by the output distribution p θ ( x t +1 | x t , z ) .
[Figure]
As our model is an autoregressive latent variable model, it consists of a forward model that evolves
[Figure]
the state of the system to the next state.
However, rather than directly producing the parameters
[Figure] [Figure]
of a distribution that we can sample the next state from, the state representation produced by theforward model is used to first condition the prior distribution of the latent variable.
Then, to produce asample of the next state, we sample the latent variable from this prior and combine it with the forwardrepresentation, which is then decoded into the next state.
This procedure is depicted in Figure 1 (left).
To enable the model to produce realistic trajectories, we align the model’s structure with that of thedata.
Consequently, we maintain the geometry of the system in the forward representations, which inpractice means using a fully convolutional architecture in the forward model as well as the decoder.
Specifically, the forward model is implemented using the U-net architecture [21], but it could inprinciple be substituted for another fully convolutional architecture.
Additionally, we do not push allinformation through a latent variable that has no geometrical structure, but allow the decoder to alsoaccess the intermediate, geometrically meaningful representation produced by the forward model.
The use of a U-net is further motivated by the desire to model dynamics on various spatial scales.
Other architectural details are provided in Appendix B. Note that, as mentioned before, latent variablez does not share the system geometry.
To sidestep this issue z is broadcasted to the entire domainbefore being concatenated to the internal representation, as done in a probabilistic U-net [16].
3.3 Training and Rollout Stability
As per the VAE framework, an approximate posterior distribution q ϕ ( z | x t , x t +1 ) is learned, alsoreferred to as the inference network.
Note that this distribution is inferred from both the currentand next timestep.
Then, while training, we optimize the Evidence Lower Bound (ELBO) on thelog-likelihood as follows:
[Equation]
The inference procedure is depicted in Figure 1 (right).
During trajectory generation the subsequenttimestep is not available and we infer latent variable z only from the prior distribution p θ ( z | x t ) .
Bythen sampling ˜ x t +1 ∼ p θ ( x t +1 | x t , z ) and repeating the entire process we can simulate over longertimeframes.
One of the major challenges of autoregressive neural simulators is the instability of long rollouts.
As the model is applied iteratively, errors accumulate, causing p θ ( x t + r | x t ) to stray from the datadistribution for sufficiently large rollout length r ∈ N + .
To mitigate this issue, various methods havebeen proposed, typically involving some form of noise injection or rollout training.
We consider twotraining methods.
The Multi-step training procedure applies the model for r iterations and calculates the lossfor each iteration.
At each step t , a reconstruction ˆ x t +1 is created by sampling fromp θ (ˆ x t +1 | x t , z ) q ϕ ( z | x t , x t +1 ) .
ˆ x t +1 then serves as the next input for the model, which is usedto create a reconstruction ˆ x t +2 , and so on.
Note that, since ˆ x t +1 is discrete, the gradients are onlybackpropagated a single step.
The Pushforward training method, adapted from [4], is similar innature.
We sample r ∼ U [1 , r max ] , rollout the model r times as described for multi-step training,but now only calculate and backpropagate the loss at the final iteration.
In this way, the modellearns to correct its own rollout errors and to map back to the data distribution.
From a probabilisticperspective, the rollout errors can be seen as adding noise to the data, where the noise is sampledfrom the model’s error distribution itself.
Orthogonal to the training procedure, we also consider two ways to sample from p θ to improvestability.
Both have in common that z is sampled from the posterior (when training) or the prior (whengenerating new simulations), but differ in how ˜ x t +1 is sampled from p θ ( x t +1 |
z, x t ) .
With maximumlikelihood sampling, a single z is sampled, and ˜ x t +1 is discretized such that p θ ( x t +1 | z, x t ) ismaximized for that z .
However, this method does not exploit any domain knowledge about the system.
With volume preservation sampling, knowledge about the (approximately) preserved volume of thecell is injected in the method.
Here, ˜ x t +1 is also discretized to maximize p θ ( x t +1 | z, x t ) , but underthe additional constraint that the volume of the cell equals its target volume.
Experiments
Our evaluation consists of two experiments.
In the first, randomly scattered walls are placed on a grid,and a force pointing to the bottom center of the lattice is applied to the cell.
Despite the comparativelysimple setting, local stochasticity in the system already leads to interesting behavior such as cellshape fluctuations, and even emerging bifurcations of trajectories.
The second experiment consists of a system where more global stochasticity with bifurcating trajectories is simulated.
A downwards force is applied to the cell, such that it makes contact with pillarsplaced at fixed positions along its path.
Then, either a leftward or rightward force is applied to thecell with equal probability, such that it passes each pillar on the left or right.
The system always hasapproximately identical initial conditions, but a strongly multi-modal distribution over trajectories.
All Cellular Potts simulations are generated with CompuCell3D [29], the state-of-the-art softwarepackage for CP model simulations.
We implement python extensions in order to generate the desireddynamics.
In all experiments, one time unit corresponds to 500 Monte Carlo steps in CompuCell3D.
For each experiment, we use 1350 training, 150 validation and 150 testing trajectories.
We compare maximum likelihood and volume preservation sampling and investigate various trainingstrategies, see Section 3.3 for details on both.
All models are trained for 300 epochs.
We optimize the
ELBO with the reparameterization trick [14] and use a linear KL-annealing schedule [3].
We set themaximum rollout length r max to a third of the total trajectory length for pushforward and multi-steptraining.
4.1 Simple Dynamics
Quantitatively, we evaluate models with two metrics: the estimated ELBO on the log-likelihood(LL), and the unrolled reconstruction log-likelihood (URLL).
LL is calculated as the sum of thelower bounds on the one-step conditional log-likelihoods (see also Equations 3 and 4).
URLLis a heuristic metric for assessing rollout stability.
Starting from initial state x 0 , we estimatep θ ( x 1 | x 0 ) , and at each subsequent step, the log-likelihood of x t is estimated from p θ ( x t | ˆ x t − 1 ) ,where ˆ x t − 1 ∼ p θ (ˆ x t −
| ˆ x t − 2 , z ) q ϕ
z | ˆ x t − 2 , x t − 1 ) is a sampled reconstruction of the model at thepreceding timestep.
A high URLL should correspond to a model that can reconstruct entire trajectorieswell, rather than only individual steps.
The results of simulators with varying training and sampling strategies are given in Table 1.
Asexpected, training strategies such as pushforward and multi-step training improve the URLL scorefor both sampling methods, compared to single-step training.
Interestingly, pushforward trainingis not always beneficial for single-step log-likelihoods, whereas multi-step training does show animprovement.
Another interesting result is that volume preservation sampling does not necessarilyimprove upon maximum likelihood sampling.
We expected that the former would net an advantageas domain knowledge is integrated in the sampling procedure, but the results are not conclusive.
Infact, multi-step training with maximum likelihood sampling performs best for both metrics, albeitwith a small margin.
Qualitatively, the simulators can capture the stochastic behavior of the system, both locally in the formof a fluctuating cell membrane, and globally, generating diverse but realistic trajectories from identical
[Table]
G r ound T r u t h
R e c on - s t r u c t i on
S a m p l e 3 S a m p l e 2 S a m p l e
G r ound T r u t h
R e c on - s t r u c t i on
S a m p l e 3 S a m p l e 2 S a m p l e initial conditions.
However, this behavior is not captured well when using single-step training.
Thisis exemplified in Figure 2, which shows reconstructed and sampled trajectories for models usingvolume preservation, trained with one-step and multi-step training.
Despite the probabilistic approach,one-step training results in samples that are almost identical.
In contrast, the same architecture trainedwith multi-step training is capable of generating a variety of realistic trajectories.
As a sanity check,we also tried replacing the CVAE by a deterministic model, which mapped the forward model’soutput directly to ˜ x t +1 .
As expected, such an approach could not capture any stochastic dynamics,resulting in a frozen cell and averaged out trajectories.
We also measured the time required for one model rollout of 60 steps (corresponding to [-
-]30000 Monte
Carlo steps for the CP model simulation) on commodity hardware (details in Appendix C).
Over100 repetitions, the average model rollout time was 2 .
16 seconds ( σ = 0 .
08 s ) on our CPU, and 0 .
56seconds ( σ = 0 .
01 s ) on our GPU.
Generating a simulation with CompuCell3D took 14 .
16 secondson average ( σ = 0 .
76 ) on the same CPU, and GPU acceleration is not supported by CompuCell3Dfor simulations not involving chemical diffusion.
Consequently, even for these simple dynamics, theneural simulator already provides an 85% speedup compared to the CP model on identical hardware.
Moreover, we envision speedups to become even larger when the dynamics are more involved, forexample, dynamics involving diffusion of chemicals or multi-cellular systems.
4.2 Bifurcating dynamics
Figure 3 shows trajectories for bifurcating dynamics with identical initial conditions.
Observe that themodel generates trajectories in which the cell follows varying realistic paths.
To quantitatively assessthe training strategies, we examine the position of the cell at the end of the simulation.
If the neuralsimulator works well, these aggregate statistics should match those of the trajectories generated by
G r ound T r u t h
R e c on - s t r u c t i on
S a m p l e 3 S a m p l e 2 S a m p l e
[Table]
Expected
[Table] [Table]
[Table]
the CP model.
Here, we present volume-preserving models as we found them to work better in thisscenario; results for models using maximum likelihood sampling are provided in Appendix D. Theend positions’ distributions are found in Table 2.
The expected distribution closely matches that ofthe CP model and our method trained with multi-step training.
However, when trained with one-stepor pushforward training, the distribution of the neural simulator differs substantially from the truedistribution.
This demonstrates that a probabilistic model in itself is not sufficient to capture global,long-term stochastic aspects of the evolution of the system, and training methods that go beyondsingle-step predictions are necessary.
Conclusion
In this work, we proposed a probabilistic neural simulation model for spatiotemporal cellular dynamics.
We adapted training strategies from autoregressive models and found that these improve rolloutquality and enable the model to accurately capture the stochastic dynamics of the system.
To evaluateour method, we generated data using the CP model and show that the learned simulator is capableof faithfully emulating the dynamics.
Furthermore, sampling from the learned simulator is aroundan order of magnitude faster than the CP simulations that generated the training data.
We concludethat neural simulators are a promising method for simulating spatiotemporal cellular dynamics, withmany interesting avenues for relevant research, for example simulating systems with multiple cells,or integrating neural cellular dynamics simulators and neural PDE solvers for multiscale modeling.
A Hamiltonians for Simple and Bifurcating Dynamics
The Hamiltonian for both simple and bifurcating dynamics (Sections 4.1 and 4.2) is as follows:
[Equation]
, where COM ( c ) is the center-of-mass vector of the cell c .
The values for each of the parameters for both simple and bifurcating dynamics are given in Table 3:
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Table]
[Equation]
[Table] [Equation] [Table] Here, F simple x = 50 −
[Table]
( c
[Table]
· 35 is continuously updated throughout the CP simulation.
For thebifurcating dynamics, if the cell is not in contact with a micropillar or if the cell reached the bottomof the grid, F bifurcate x = 0 .
If the cell comes into contact with a micropillar and F bifurcate x equals 0, itis set to a value with magnitude equal to the contact area with a micropillar, of which the sign isnegative or positive with equal probability.
Finally, if the cell is still in contact with a micropillarand F bifurcate x is nonzero, the horizontal force is rescaled such that its magnitude equals the contactarea of the micropillar.
We explicitly note that these parameters were not necessarily chosen to bebiologically plausible, but to generate stylized CP simulations that exhibit behavior which will berelevant for modeling more complex, realistic scenarios as well.
B Model Architecture Details
The details of all model components are given below.
All convolutional layers are 2D convolutionswith kernel size 3x3, unless otherwise mentioned.
• Forward Model:
The forward model starts with a linear convolutional layer to lift the inputto a higher dimensional representation.
This representation is processed by a U-net.
TheU-net has 4 contrastive blocks, one block that operates on the lowest level of the U-net,and 4 expansive blocks.
Each contrastive block consists of two convolutional layers with ReLU activation, followed by 2x2 maxpooling.
The block that operates on the lowest levelhas the same architecture, but does not do maxpooling.
Each expansive block consists of a2x2 upconvolution and the concatenation of the result with the output of the corresponding
contrastive block.
Then, two convolutional layers with ReLU activation are applied.
Eachcontrastive block doubles the channel dimension, while each expansive block halves thechannel dimension.
Finally, the output of the U-net is cropped from the center such that thedimensionality of the output is the same as the input, to which it is summed (i
e., a residualconnection).• Conditional Variational Autoencoder: The CVAE consists of a prior network, generativenetwork, and inference network.
For the CVAE, the kernel size of the convolutional layersis 5x5 instead of 3x3.
– Prior network p θ ( z | x t ) : first, the image channel corresponding to the one-hot encodingof the walls is concatenated with the forward model’s outputs.
Then, two convolutionallayers with ReLU activation are applied, followed by 2x2 maxpooling.
Subsequently,two repetitions of a convolutional layer with ReLU activation followed by maxpoolingare applied.
Finally, the output is flattened and mapped to a lower-dimensional spaceusing a linear layer with ReLU activation.
Two separate linear layers map this outputto the parameters of the normal distribution over the latent space µ and log σ .–
Generative network p θ ( x t +1 | z, x t ) :
z is concatenated channel-wise to the forwardmodel’s output, along with the wall channel of x t .
Two convolutional layers with ReLUactivation are applied, followed by one convolutional layer with sigmoid activation, toobtain the pixel-wise Bernoulli probabilities of the cell’s location.
As the cell cannotbe located on top of a wall, we explicitly set the decoder output probabilities for thesepixels to 0.
– Inference network q ϕ ( z | x t +1 , x t ) : the architecture of the inference network is identical to the prior network, with the exception that it takes an extra input channel,corresponding to the cell channel of the one-hot encoding of x t +1 .
C Hardware
Model training was done on a single Nvidia RTX 3060 GPU, and took at most three hours per model,depending on the dataset and training strategy.
For a fair comparison of the simulation times of the Cellular Potts model and our method, we resorted to local hardware, as CompuCell3D did anyhownot support GPU acceleration for simulations without chemical diffusion.
The CPU used was an Inteli7-9750H CPU clocked at 2.6GHz.
For the GPU time measurements, we used an Nvidia Quadro
P2000.
D Maximum Likelihood Sampling for Bifurcating Dynamics
To quantitatively evaluate maximum likelihood sampling we again investigate the cell’s position at theend of a simulation.
Table 4 contains the distributions of these end positions.
While the model trainedusing the multi-step method approaches the ground-truth distribution, there are clear deviations.
Thevolume-preserving model discussed in Section 4.2 much closer aligns with the Cellular Potts modeland the expected distribution.
[Table]