Variational Neural Cellular Automata

1 INTRODUCTION . The process of cellular growth and differentiation is capable of creating an astonishing breadth of form and function , filling every conceivable niche in our ecosystem . Organisms range from complex multi-cellular organisms like trees , birds , and humans to tiny microorganisms living near hydrothermal vents in hot , toxic water at immense pressure . All these incredibly diverse organisms are the result of the same generative process of cellular growth and differentiation . Cellular Automata ( CA ) are computational systems inspired by this process of cellular growth and differentiation , where ” cells ” iteratively update their state based on the state of their neighbor cells and the rules of the CA . Even with simple rules , CA can exhibit complex and surprising behavior — with just four simple rules , Conway ’ s Game of Life is Turing complete . Neural Cellular Automata ( NCA ) are CA where the cell states are vectors , and the rules of the CA are parameterized and learned by neural networks ( Mordvintsev et al. , 2020 ; Nichele et al. , 2017 ; Stanley & Miikkulainen , 2003 ; Wulff & Hertz , 1992 ) . NCAs have been shown to learn to generate images , 3D structures , and even functional artifacts that are capable of regenerating when damaged ( Mordvintsev et al. , 2020 ; Sudhakaran et al. , 2021 ) . While these results are impressive , the NCAs can only generate ( and regenerate ) the single artifact it is trained on , lacking the diverse generative properties of current probabilistic generative models . In this paper , we introduce the VNCA , an NCA based architecture that addresses these limitation while only relying on local communication and self-organization . Generative modeling is a fundamental task in machine learning , and several classes of methods have been proposed . Probabilistic generative models are an attractive class of models that directly define a distribution over the data p ( x ) , which enable sampling and computing ( at least approximately ) likelihoods of observed data . The average likelihood on test data can then be compared between models , which enables fair comparison between generative models . The Variational Auto-Encoder ( VAE ) is a seminal probabilistic generative model , which models the data using a latent variable model , such that p ( x ) = ∫ z pθ ( x|z ) p ( z ) , where p ( z ) is a fixed prior and pθ ( x|z ) is a learned decoder ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . The parameters of the model is learned by maximizing the Evidence Lower Bound ( ELBO ) , a lower bound on log p ( x ) , which can be efficiently computed using amortized variational inference ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . On a high level , our proposed generative model combines NCAs with VAEs , by using a NCA as the decoder in a VAE . We perform experiments with a standard NCA decoder and a novel NCA architecture , which duplicates all cells every M steps , inspired by cell mitosis . This results in better generative modeling performance and has the computational benefit that the VNCA only computes cell updates for the living cells at each step , which are relatively few early in growth . While our model is inspired by the biological process of cellular growth and differentiation , it is not intended to be an accurate model of this process and is naturally constrained by what we can efficiently compute . Other works have explored similar directions . The VAE-NCA proposed in Chen & Wang ( 2020 ) , is actually , despite the name , not a VAE since it does not define a generative model or use variational inference . Rather it is an auto-encoder that decodes a set of weights , which parameterize a NCA that reconstructs images . NCAM ( Ruiz et al. , 2020 ) similarly uses an auto-encoder to decode weights for a NCA , which then generates images . In contrast to both of these , our method is a generative model , uses a single learned NCA , and samples the initial state of the seed cells from a prior p ( z ) . StampCA ( Frans , 2021 ) also learns an auto-encoder but similarly encodes the initial hidden NCA state of a central cell . None of these are generative models , in the sense that it ’ s not possible to sample data from them , nor do they assign likelihoods to samples . While previous works all show impressive reconstructions , it is impossible to evaluate how good generative models they really are . In contrast , we aim to offer a fair and honest evaluation of the VNCA as a generative model . Graph Neural Networks ( GNNs ) can be seen as a generalization of the local-only communication in NCA , where the grid neighborhood is relaxed into an arbitrary graph neighborhood ( Grattarola et al. , 2021 ) . GNNs compute vector-valued messages along the edges of the graph that are accumulated at the nodes , using a shared function among all nodes and edges . GNNs have been successfully used to model molecular properties ( Gilmer et al. , 2017 ) , reasoning about objects and their interactions ( Santoro et al. , 2017 ) , and learning to control self-assembling morphologies ( Pathak et al. , 2019 ) . Powerful generative image models have been proposed based on the variational auto-encoder framework . A common distinction is whether the decoder is auto-regressive or not , i.e. , whether the decoder conditions its outputs on previous outputs , usually those above and to the left of the current pixel ( Van den Oord et al. , 2016 ; Salimans et al. , 2017 ) . Auto-regressive decoders are so effective at modelling p ( x ) directly that they have the tendency to ignore the latent codes thus failing to reconstruct images well ( Alemi et al. , 2018 ) . Additionally , they are expensive to sample from since sampling must be done one pixel at a time . Our decoder is not auto-regressive , thus cheaper to sample from , and more similar to e.g. , BIVA ( Maaløe et al. , 2019 ) or NVAE ( Vahdat & Kautz , 2020 ) . Contrary to BIVA , NVAE and other similar deep convolution-based VAEs , our VNCA defines a relatively simple decoder function that only consider the immediate neighborhood of a cell ( pixel ) at each step and is applied iteratively over a number of steps . The VNCA thus aims to learns a self-organising generative process . NCAs and self-organizing systems in general , represent an exciting research direction . The repeated application of a simple function that only relies on local communication lends itself well to a distributed leaderless computation paradigm , useful in , for instance , swarm robotics ( Rubenstein et al. , 2012 ) . Additionally , such self-organizing systems are often inherently robust to perturbations ( Mordvintsev et al. , 2020 ; Najarro & Risi , 2020 ; Tang & Ha , 2021 ; Sudhakaran et al. , 2021 ) . In the VNCA , we show that this robustness allows us to set a large fraction of the cells ’ states to zero and recover almost perfectly by additional iterations of growth ( Fig . 1 , right ) . By introducing an NCA that is a proper generative probabilistic model , we hope to further spur the development of methods at the intersection of artificial life and generative models . 2 VARIATIONAL NEURAL CELLULAR AUTOMATA . The VNCA defines a generative latent variable model with z ∼ p ( z ) and x ∼ pθ ( x|z ) , where p ( z ) = N ( z|µ = 0 , Σ = I ) is a Gaussian with zero mean and diagonal co-variance and pθ ( x|z ) is a learned decoder , with parameters θ . The model is trained by maximizing the evidence lower bound ( ELBO ) , log p ( x ) −DKL ( qφ ( z|x ) ‖ p ( z|x ) ) = −DKL ( qφ ( z|x ) ‖ p ( z ) ) + Ez∼qφ ( z|x ) log pθ ( x|z ) , where , qφ ( z|x ) is a learned encoder , with parameters φ , which learns to perform amortized variational inference . Since DKL ≥ 0 , the right hand side is a lower bound on log p ( x ) and increasing it either increases log p ( x ) or decreases DKL ( qφ ( z|x ) ‖ p ( z|x ) ) , which is how closely the amortized variational inference matches the true unknown posterior p ( z|x ) . By using the reparameterization trick , the sampling , z ∼ qφ ( z|x ) , is differentiable , and both φ and θ can be learned efficiently with stochastic gradient ascent ( Kingma & Welling , 2013 ) . The DKL term on the right hand side can be weighted with a parameter β to control whether the VAE should prioritize better reconstructions ( β < 1 ) , or better samples ( β > 1 ) ( Higgins et al. , 2016 ) . For β > 1 this is still a ( looser ) lower bound on log p ( x ) . The decoder of the generative process , pθ ( x|z ) , is based on a NCA ( Fig . 2 ) . The NCA defines a recurrent computation over vector-valued ” cells ” in a grid . At step t of the NCA , the cells compute an additive update to their state , based only on the state of their neighbor cells in a 3× 3 neighborhood , zi , t = zi , t−1 + uθ ( { zj , t−1 } j∈N ( i ) ) , ( 1 ) where uθ is a learned update function , index i denote a single cell and N ( i ) denotes the indexes of the 3 × 3 neighborhood of cell i , including i . All cells are updated in parallel . In practice , uθ is efficiently implemented using a Convolutional Neural Network ( CNN ) ( Gilpin , 2019 ) with a single 3 × 3 filter followed by a number of 1 × 1 filters and non-linearities . The initial grid z0 consists of a single sample from the latent space repeated in a 2× 2 grid . A traditional NCA defines whether a cell is alive or dead based on the cell state and its neighbors ’ cell states ( Mordvintsev et al. , 2020 ) . This allows the NCA to “ grow ” from a small initial seed of alive cells . Our NCA uses a different approach inspired by cell mitosis . Every M steps of regular NCA updates all the cells duplicate by creating new cells initialized to their own current state to the right and below themselves , “ pushing ” the existing cells out of the way as they do so . This is efficiently implemented using a “ nearest ” type image rescaling operation ( Fig . 2 , right ) . After K doubling and M ( K + 1 ) steps of NCA updates , the final cell states zT condition the parameters of the likelihood distribution . For RGB images , we use a discrete logistic likelihood with L mixtures ( Salimans et al. , 2017 ) , such that the 10L first channels are the parameters of the discrete logistic mixture1 and for binary images a Bernoulli likelihood with the first channel of each cell being log p. The update function is a single 3×3 convolution , followed by four residual blocks ( He et al. , 2016 ) . Each block consists of a 1 × 1 convolution , Exponential Linear Unit ( ELU ) non-linearity ( Clevert et al. , 2015 ) and another 1× 1 convolution . Finally , a single 1× 1 convolution maps to the cell state size . This last convolution layers ’ weights and biases are initialized to zero to prevent the recurrent computation from overflowing . The number of channels for all convolutional layers , the size of the latent space , and the cell state are all 256 . The total number of parameters in the NCA decoder is approximately 1.2M . Since we are primarily interested in examining the properties of the NCA decoder , the encoder qφ ( z|x ) is a fairly standard deep convolutional network . It consists of a single layer of 5 × 5 con- 1We use the implementation from https : //github.com/vlievin/biva-pytorch volution with 32 channels followed by four layers of 5× 5 convolution with stride 2 , where each of these layers has twice the number of channels as the previous layer . Finally , the down-scaled image is flattened and passed through a dense layer which maps to twice the latent space dimensionality . The first half encodes the mean and the second half the log-variance of a Gaussian . All convolutional layers are followed by ELU non-linearities . For a detailed model definition , see the appendix .

Summary: This paper introduces a variational generative model based on Neural Cellular Automata (NCA). The model is called the Variational NCA (VNCA).  The VNCA is designed for images:   - The encoder is a typical convolutional neural network.  - The decoder is an NCA that iteratively refines the image, alternating convolutions and upsampling up to the desired image size. The upsampling process is loosely inspired by mitosis in living cells.   The authors perform morphogenesis experiments on three datasets: MNIST, Noto Emoji, and CelebA. The results are good on MNIST, less so on the other two datasets, although there is clear evidence that the model can learn to generate meaningful images.  The authors also perform an experiment to see if the VNCA is robust to perturbations (occlusions) and show that the model has a reasonable degree of robustness even without ever seeing any perturbations at training time.