[Re] Variational Neural Cellular Automata

Albert Sund Aillet, Simon Sondén

Keywords: Neural Cellular Automata, Cellular Automata, Self-Organization, Generative Models

TL;DR: We reproduce the results of a paper proposing the Variational Neural Cellular Automata, a probabilistic generative model based on neural cellular automata

Abstract:

The main claim of the paper being reproduced is that the proposed Variational Neural Cellular Automata (VNCA) architecture, composed of a convolutional encoder and a Neural Cellular Automata (NCA)-based decoder, is able to generate high-quality samples.
The paper presents two variants of this VNCA decoder: the doubling VNCA variant that is claimed to have a simple latent space, and the non-doubling VNCA variant that is claimed to be optimized for damage recovery and stability over many steps.

To reproduce the results, we re-implemented all of the VNCA models and a fully-convolutional baseline in JAX, by using the descriptions given in the paper. We then followed the same experimental setup and hyperparameter choices as in the original paper.

All of the models were trained on a TPU v3-8 provided by Kaggle, with a total budget of around 4 TPU hours, not counting unreported experiments.

All but one of the figures and results from the original study were possible to reproduce. The obtained Evidence Lower Bound (ELBO) of the doubling VNCA was within $0.3\%$ of the stated and for the non-doubling VNCA the ELBO was within $1.8\%$ and the observed damage recovery was similar. We were however not able to reproduce the t-SNE reduction experiment for the baseline and were therefore not able to show the VNCA decoder having a cleaner t-SNE separation than the baseline.

The implementation of the convolutional baseline and most parts of the NCA-based decoder were straightforward to re-implement based on the description provided in the paper.

One of the difficulties faced in the reproduction study was obtaining the labels of the binarized MNIST dataset used in the original paper since it officially is provided without labels. This made it unclear how the original paper got the labels for the latent space visualization.
Additionally, implementing the pool for the non-doubling version of the VNCA model with a distributed training setup was challenging as it required operations between TPU cores.

No communication was made with the original authors.

Paper Url: https://arxiv.org/abs/2201.12360

Paper Review Url: https://openreview.net/forum?id=7fFO4cMBx_9

Paper Venue: ICLR 2022

Confirmation: The report pdf is generated from the provided camera ready Google Colab script, The report metadata is verified from the camera ready Google Colab script, The report contains correct author information., The report contains link to code and SWH metadata., The report follows the ReScience latex style guides as in the Reproducibility Report Template (https://paperswithcode.com/rc2022/registration)., The report contains the Reproducibility Summary in the first page., The latex .zip file is verified from the camera ready Google Colab script

Journal: ReScience Volume 9 Issue 2 Article 31

Doi: https://www.doi.org/10.5281/zenodo.8173729

Code: https://archive.softwareheritage.org/swh:1:dir:3de126d09281c3dab4749dc8ac9330e53a22c558