I joined the Neural Cellular Automata team at Monash DeepNeuron in February 2025. The team builds the models from Mordvintsev et al. (2020) and runs them in the browser at neuralca.org. My part, from March to May, was the 3D version: could the same rule learn gastrulation from a recording of a real embryo. This is what I built and where it got to. There's a live 2D model you can cut on the project page.
The model#
An NCA is a grid of cells. Each cell holds 16 numbers, four of them colour and alpha, the other twelve whatever the network decides to use them for. Every step, each cell looks at its 3×3 neighbourhood through three fixed filters, runs that through a two-layer network with 128 hidden units, and adds the result to its own state. A random subset of cells updates each step. That's the whole model. Nothing has coordinates and nothing sees the whole grid, and it still grows a picture from one cell and regrows it when you cut a hole in it.
In 3D the neighbourhood is 3×3×3 and the grid is a voxel volume. The rule is otherwise the same.
Gastrulation#
Gastrulation is the part of embryo development where a ball of cells folds itself into layers and the body plan appears. No cell is in charge, each one reacts to its neighbours, which is the setup an NCA already has. So the question was whether the rule could learn it from data rather than from a target picture.
The data is a light-sheet recording of a mouse embryo from McDole et al. (2018), with every cell tracked through time as an x, y, z position. We turned the tracked positions into voxel density volumes, one per timepoint, and trained on transitions between timepoints instead of one final shape.
Getting the data in#
Most of the work was before any training happened. The recording comes as a database of tracked cells, one row per cell per timepoint, 265 timepoints from start to end. The model wants a voxel volume per timepoint with the embryo in the alpha channel. Getting from one to the other:
- Outliers. Drop any cell more than 1.5 IQR outside the middle on x, y or z. The tracking has stray points a long way from the embryo.
- Normalise. Centre each axis on its mean, then scale all three by the same global min and max into a 16×16×16 grid. Per-axis scaling is in there as an option and stayed commented out.
- Voxelise. Round every cell to a grid position and count. The count is the density, normalised by the biggest count across all frames so the frames are comparable.
- Crop. Find the bounding box of the non-zero voxels across every frame and crop the volumes to it, so the grid can be x×y×z instead of a cube.
- Reduce. A reduced version of the database with 283,454 records and 15 timepoints is committed to the repo, so training doesn't need the full recording.
Then the visualiser, because you can't tell if any of that worked from a tensor. A matplotlib 3D scatter per frame, coloured by density, stitched into a gif at 5 frames a second. One in preprocessing for the target data and one in training for the model output, so the two can be compared. The viewer above is the same idea in the browser.
The voxel trainer first#
Before the embryo there was a plain 3D trainer that grows .vox models, a donut, a potted plant, some trees, from one seed cell. My first commits, 25 March to 11 April, were on that:
- Floats. Getting it to run without a GPU. The code assumed float64 everywhere, so I made the dtype follow the hardware, float64 on CUDA and float32 on CPU, and moved the identity kernel and the seed to match.
- Logging. Loss saving, loss plotting, then persistent loss tracking across runs, so we could tell whether a run was learning at all instead of watching numbers scroll past.
- Loss. The MSE was scoring all 16 channels. Fixed it to score only RGB.
Training on the embryo#
The gastrulation commits run from 27 April to 14 May. The repo was restarted on 19 April, so the data work above only shows up in the history where I touched it again after that.
- Batches. Rewrote the training loop around batches that pair a volume with a later one. Each batch mixes three kinds of seed: the first frame, a random earlier frame, and the frame just before the target, each with its own number of update steps to reach the target. Frames go to float32 and onto the GPU when there is one. 100 epochs, where one epoch is a pass over every frame.
- IoU. Replaced the MSE with a soft IoU loss on the alpha channel,
1 - intersection / union, which keeps its gradients and gave better training feedback. Co-authored with Chloe Koe and Nathan Culshaw. - Defaults. Last commit was frame conversion and logging defaults so someone else could run it.
What didn't work#
The MSE on the embryo. It wasn't telling us much about whether the volume was in the right place, so the IoU replaced it.
Detaching the output before the loss. It cut the graph, so the model stopped learning without any error. Found it while wiring in the IoU.
The model itself, so far. The loss plot in the repo starts at 0.994, gets to 0.990 around epoch 40, and drifts back up to 0.993 by epoch 100. If that's the IoU loss, it's under 1% overlap with the target at its best. The 2D model works. The 3D one is where it was when I left it in May, and the viewer above draws the target data, not model output.
Now#
The code is at MonashDeepNeuron/3D-Neural-Cellular-Automata. The team is Afraz Gul, Chloe Koe, Nathan Culshaw, Angus Bosmans, Alexander Mai and me, with Keren Collins, Joshua Riantoputra and Nyan Knaw advising.