Variational Neural Cellular Automata

TL;DR

The paper addresses the need for probabilistic, self-organizing generative processes by merging neural cellular automata with variational inference to form Variational Neural Cellular Automata (VNCA). The approach uses a VAE framework with a 3×3 neighborhood NCA decoder and an innovative mitosis-inspired doubling mechanism, trained via ELBO and optionally beta-weighted KL terms. Key contributions include establishing VNCA as a proper generative model with a learnable encoder, demonstrating self-organizing growth and damage-resilient attractors on MNIST and CelebA, and introducing pool-and-damage training alongside a non-doubling variant to study resilience. The findings show VNCA can learn diverse outputs from a single latent code and recover from significant damage, although it trails state-of-the-art generative models in raw sample quality, highlighting a promising direction for robust, emergent generation.

Abstract

In nature, the process of cellular growth and differentiation has lead to an amazing diversity of organisms -- algae, starfish, giant sequoia, tardigrades, and orcas are all created by the same generative process. Inspired by the incredible diversity of this biological generative process, we propose a generative model, the Variational Neural Cellular Automata (VNCA), which is loosely inspired by the biological processes of cellular growth and differentiation. Unlike previous related works, the VNCA is a proper probabilistic generative model, and we evaluate it according to best practices. We find that the VNCA learns to reconstruct samples well and that despite its relatively few parameters and simple local-only communication, the VNCA can learn to generate a large variety of output from information encoded in a common vector format. While there is a significant gap to the current state-of-the-art in terms of generative modeling performance, we show that the VNCA can learn a purely self-organizing generative process of data. Additionally, we show that the VNCA can learn a distribution of stable attractors that can recover from significant damage.

Figures (17)

• Figure 1: VNCA self-organized growth and damage recovery. Left: The VNCA has learned a self-organising generative process that generates faces, from an initial random seed of cells from $\mathcal{N}(0,I)$. Time goes top-down and left-to-right. Every eight steps all the cells duplicate, which can be seen on the diagonal. Right: A damage recovery sequence for an MNIST digit.
• Figure 2: VNCA overview. Left: Generative process. $z_0$ is sampled from $p(z)$ and acts as an initial $2\times2$ seed of cells, which grows through a series of K doubling NCA steps. $|Z|$ denotes the size of the latent space, which is identical to the size of the cell states. Finally, the last cell hidden state conditions the parameters of the likelihood distribution $p(x|z)$. Middle: Each doubling NCA step consists of a number of NCA steps followed by a doubling operation. At each NCA step, the cells can only communicate with their immediate neighbors. Right: The doubling operator doubles the cell grid as depicted, where the color indicates the state vector of each cell.
• Figure 3: Left: Test set reconstructions. Right: Unconditional samples from the prior. The VNCA achieves $\log p(x) \geq -84.23$ nats evaluated with 128 importance weighted samples.
• Figure 4: Self-organising MNIST, with the seed cell states sampled from the prior. Each $32 \times 32$ sub-image shows the digit after 1 step of NCA or doubling. Time flows top-to-bottom and left-to-right. The doubling steps can be seen on the diagonal. Note: this shows sample averages for clarity.
• Figure 5: Exploring the latent space of VNCA. Fig. \ref{['fig:latent_space_exploration:interpolations']} shows several linear interpolations between random digits in the latent space of a VNCA trained on binarized MNIST. Figs. \ref{['fig:latent_space_exploration:tSNE-vaenca']} and \ref{['fig:latent_space_exploration:tSNE-baseline']} show the result of reducing the dimensionality of 5000 digits chosen at random for the VNCA and for the deep convolutional baseline respectively. Notice how the latent space of VNCA has more t-SNE structure and cleaner separation of digit encodings. Note: this shows sample averages for clarity.