Compiling a Neural Net to C for a 1,744× speedup

Summary

Compiling a Neural Net to C for a 1,744× speedup 2025-05-27 · About 25 minutes long tl;dr: I trained a neural network (NN), with logic gates in the place of activation functions, to learn a 3×3 kernel function for Conway’s Game of Life. I wanted to see if I could speed up inference by extracting the learned logic circuit from the NN. So, I wrote some code to extract and compile the extracted logic circuit to bit-parallel C (with some optimizations to remove gates that don’t contribute to the output). I benchmarked the original NN against the extracted 300-line single-threaded C program.; compiling the NN to C resulted in a 1,744× speedup! Crazy, right? Here’s the repo: ~354 lines of Python/JAX, ~331 lines of C, if you want to reproduce it and/or mess around. The longer story While plumbing the intertubes (as one does), I came across this fun publication by the Self Organising Systems group at Google, about Differentiable Logic Cellular Automata. This research caught my attention (I mean, who doesn’t love a pretty picture), and as I read it, I realized the whole idea wouldn’t be too hard to replicate. (Which is crazy, because this is crazy. I mean, the authors cite creating computronium as a source of inspiration. Awesome!) To break things down a little bit: Differentiable Logic Cellular Automata is a quite a mouthful, I know, but the idea is the straightforward composition of two existing ideas: Cellular Automata (CA) are grids of cells, where each cell changes over time according to some local rule. The most famous cellular automata are Conway’s Game of Life and perhaps Rule 110. We call this local update rule a kernel, a function that looks at the local neighborhood around a cell to calculate the cell’s next state. By applying the kernel to each cell in our grid, we step the cellular automata forward through time. Simple rules can give rise to strikingly complex behaviour. Neural Cellular Automata (NCA) are a variant of CA that replace the kernel function with a n...