BraiNCA - brain-inspired neural cellular automata and applications to morphogenesis and motor control

year:
paper: brainca-brain-inspired-neural-cellular-automata-and-applications-to-morphogenesis-and-motor-control
website:
code:
connections: NCA, Benedikt Hartl, levin

---

Notation

${\mathbf{c}}_i^t\in {\mathbb{R}}^C$ … cell state of cell $i$ at time $t$
${\rho }_i^t$ … optional conditioning input (positional encoding, prev action, …)
${\mathbf{s}}_i^t=[{\mathbf{c}}_i^t;{\rho }_i^t]\in {\mathbb{R}}^S$ … extended state
${\mathcal{N}}_i$ … local neighborhood indices
${\mathcal{L}}_i$ … long-range neighborhood indices
$a_{ij}^t\in {\mathbb{R}}^A$ … raw local attention score
$b_{ij}^t\in {\mathbb{R}}^A$ … raw long-range attention score
${\alpha }_{ij}^t$ … normalized local attention weight
${\beta }_{ij}^t$ … normalized long-range attention weight
$f_{\text{attn}}^{(\mathcal{N})}$ … local attention MLP, ${\mathbb{R}}^{2S}\to {\mathbb{R}}^A$
$f_{\text{attn}}^{(\mathcal{L})}$ … long-range attention MLP, ${\mathbb{R}}^{2S}\to {\mathbb{R}}^A$
${\mathbf{h}}_k^t\in {\mathbb{R}}^C$ … GRU hidden state (values aggregated by attention)
${\mathbf{n}}_i^t$ … aggregated local signal
${\mathbf{l}}_i^t$ … aggregated long-range signal
${\mathbf{z}}_i^t$ … interaction vector
$f_Z$ … composition operator for $\mathbf{z}$
${\mathbf{m}}_i^t\in {\mathbb{R}}^M$ … message MLP output, GRU input
$f_{\text{msg}}$ … message MLP
${\mathbf{u}}_t$ … system-level feedback (e.g. one-hot prev action)
${\mathbf{o}}_t$ … external observation
$f_{\text{GRU}}$ … GRU cell
${\mathbf{r}}_i^t\in {\mathbb{R}}^C$ … refinement residual
$f_{\text{refine}}$ … refinement MLP
${\mathbf{\ell }}_i^t$ … per-cell action logits (lunar lander)
$w_i^t$ … per-cell fire probability
${\mathcal{R}}_r$ … cells in action region $r$
$L_r^t$ … region-level action logit
$\pi (a_t=r\mid {\mathbf{o}}_t)$ … action policy
$N$ … number of cells
$C$ … cell state dim
$S$ … extended state dim
$A$ … attention score dim
$M$ … message dim
$T$ … final timestep

Limitations / obvious next steps

• GRU, not QKV-attention-memory
• With graph attn there’s also no any2any between the msgs per node (not sure if it matters - prlly not if u factor efficiency and diffusion over time)
• Hidden state is public / the same as the public message, not compressed / …
• Long-range connectivity is fixed (Zipf’s law)
• All cells get (the same) observations…
• Arrangement affects performance and learning speed - let the cells arrange themselves
• Growth instead of fixed structure path: weights are the genome, evolved slowly over many generations of growth, growth happening during rollouts (i.e. genome only meta learns). This prlly (definitely) needs some pre-training to have any chance of working. Perhaps… pre-pre-training the NCA on NCAs…

Unclarities

• Why does the morpho experiment use additive fusion $[\mathbf{c};\mathbf{n}+\mathbf{l}]$ for the interaction vector while lunar lander concatenates $[\mathbf{c};\mathbf{n};\mathbf{l}]$? Not motivated in the paper.
• why dont they do all to all comms in the attn agg? efficiency? unnecessary? (same info is spread out through timesteps) - maybe full attn only worth it for internal states
• Lunar lander readout: paper prose says ”$f_{\text{refine}}$ output is split into two heads” but the equations write both heads as reading from ${\mathbf{h}}_i^t$ directly. $f_{\text{act}}$: $M\to M\to 2$, but ${\mathbf{h}}_i^t\in {\mathbb{R}}^C$, $C\ne M$ suggests the prose is right and the equations are wrong. And it also makes more sense that way…

BraiNCA

Init: ${\mathbf{c}}_i^0\sim \mathcal{N}(0,0.1\cdot I)$ (morpho) or ${\mathbf{c}}_i^0=0$ (lunar). ${\mathbf{o}}_t,{\mathbf{u}}_{t-1}$ broadcast to all cells.
Repeat for $T$ steps (morpho: $T=35$ total; lunar: $n_{\text{steps}}=3$ per env step), $i\in \{1,\dots ,N\}$ cells ($N=256$):

$$\begin{array}{rl}{\mathbf{s}}_i^t& =[{\mathbf{c}}_i^t;{\rho }_i^t]\\ a_{ij}^t& =f_{\text{attn}}^{(\mathcal{N})}([{\mathbf{s}}_i^t;{\mathbf{s}}_j^t]),{\alpha }_{ij}^t={\text{softmax}}_j(a_{ij}^t),{\mathbf{n}}_i^t={\sum }_{k\in {\mathcal{N}}_i}{\alpha }_{ik}^t{\mathbf{h}}_k^t\\ b_{ij}^t& =f_{\text{attn}}^{(\mathcal{L})}([{\mathbf{s}}_i^t;{\mathbf{s}}_j^t]),{\beta }_{ij}^t={\text{softmax}}_j(b_{ij}^t),{\mathbf{l}}_i^t={\sum }_{k^{\prime }\in {\mathcal{L}}_i}{\beta }_{i{k}^{\prime }}^t{\mathbf{h}}_{k^{\prime }}^t\\ {\mathbf{z}}_i^t& =f_Z({\mathbf{s}}_i^t,{\mathbf{n}}_i^t,{\mathbf{l}}_i^t)\\ {\mathbf{m}}_i^t& =f_{\text{msg}}({\mathbf{z}}_i^t;{\mathbf{u}}_t,{\mathbf{o}}_t)\\ {\mathbf{h}}_i^t& =f_{\text{GRU}}({\mathbf{m}}_i^t,{\mathbf{h}}_i^{t-1})\\ {\tilde{\mathbf{h}}}_i^t& ={\mathbf{h}}_i^t+f_{\text{refine}}({\mathbf{h}}_i^t)\\ {\mathbf{c}}_i^{t+1}& =\{\begin{array}{ll}{\tilde{\mathbf{h}}}_i^t& \text{(morphogenesis)}\\ W_C{\tilde{\mathbf{h}}}_i^t& \text{(lunar lander)}\end{array}\\ {\mathbf{\ell }}_i^t& =f_{\text{act}}({\tilde{\mathbf{h}}}_i^t)\text{(lunar lander only)}\end{array}$$

Approach to actions, lunar lander

• divided into regions
• softmax over logits $({\ell }_{\text{noop}},{\ell }_{\text{fire}})$ → fire probability $w_i$
• majority vote, confidence weighted
• Within each region, compute a weighted average of fire logits , weighted by the fire probability/confidence $w_i$:

$$L_r=\frac{{\sum }_{i\in \mathcal{R}r}{w}_i\cdot {\ell }_{i,\text{fire}}}{{\sum }_{i\in {\mathcal{R}}_r}{w}_i}=L_{\text{NOOP}},L_{\text{LEFT}},L_{\text{MAIN}},L_{\text{RIGHT}}$$

• softmax over $L_r$ to get actiono dist to sample from

Detailed

Extended state ${\mathbf{s}}_i^t=[{\mathbf{c}}_i^t;{\rho }_i^t]$ ${\rho }_i^t$ bundles all non-cell-state inputs. What it contains is task-dependent: Morphogenesis: $\rho =$ nothing $\mathbf{s}=\mathbf{c}$, $S=C=9$ … 3 visible channels (cell type logits) + 6 hidden channels. Lunar lander: ${\rho }_i^t=[{\ell }_i^{t-1};{\rho }_i]$ where ${\ell }_i^{t-1}\in \{0,1\}$ is the previous fire/noop decision, and ${\rho }_i\in {\mathbb{R}}^2$ is normalized grid position ¹ (static). $S=15$ … extended state. $C+1+2=12+1+2$. The 1 is the previous fire/noop decision, the 2 is the positional encoding. $C=12$ … latent dim (no "visible" channels, the readout is via a separate action head).

Neighbor scoring ( graph attention) $f_{\text{attn}}$: Linear($2S$, 64) → GELU → Linear(64, $A$), applied to concatenated pair $[{\mathbf{s}}_i^t;{\mathbf{s}}_j^t]$. Each cell scores each of its neighbors independently. Local and long-range use separate $f_{\text{attn}}$:

\begin{align*}
a_{ij}^t &= f_{\text{attn}}^{(\mathcal{N})}([\mathbf{s}i^t;\mathbf{s}j^t]), \quad \alpha{ij}^t = \text{softmax}j(a{ij}^t), \quad \mathbf{n}i^t = \sum{k\in\mathcal{N}i}\alpha{ik}^t,\mathbf{h}k^t \
b{ij}^t &= f{\text{attn}}^{(\mathcal{L})}([\mathbf{s}i^t;\mathbf{s}j^t]), \quad \beta{ij}^t = \text{softmax}j(b{ij}^t), \quad \mathbf{l}i^t = \sum{k’\in\mathcal{L}i}\beta{ik’}^t,\mathbf{h}{k’}^t
\end{align*}

Scores use extended states $\mathbf{s}$, but values aggregated are GRU hidden states $\mathbf{h}$.

"Interaction vector" (cell input)

\mathbf{z}{i}^{t}=f{Z}(\mathbf{s}{i}^{t},\mathbf{n}{i}^{t},\mathbf{l}_{i}^{t})

Morphogenesis vanilla: ${\mathbf{z}}_i^t=[{\mathbf{c}}_i^t;{\mathbf{n}}_i^t]$
Morphogenesis long-range: ${\mathbf{z}}_i^t=[{\mathbf{c}}_i^t;{\mathbf{n}}_i^t+{\mathbf{l}}_i^t]$
Lunar lander long-range: ${\mathbf{z}}_i^t=[{\mathbf{c}}_i^t;{\mathbf{n}}_i^t;{\mathbf{l}}_i^t]$

"Message MLP"

$${\mathbf{m}}_i^t=f_{\text{msg}}({\mathbf{z}}_i^t;{\mathbf{u}}_t,{\mathbf{o}}_t)$$

In morphogenesis, ${\mathbf{u}}_t$ and ${\mathbf{o}}_t$ are empty so it’s just $f_{\text{msg}}({\mathbf{z}}_i^t)$.
In lunar lander, they’re simply concatted to ${\mathbf{z}}_i^t$, all cells receive the same obs.

GRU update + readout

\begin{align*}
\mathbf{h}{i}^{t} &= f{\textnormal{GRU}}(\mathbf{m}{i}^{t},\mathbf{h}{i}^{t-1}) \
\tilde{\mathbf{h}}_i^t &= \mathbf{h}i^t + f{\textnormal{refine}}(\mathbf{h}_i^t)
\end{align*}

$f_\text{refine}$ … the "refinement MLP" (two dense layers, $C$ units, GELU) is just a [[ResNet|resblock]]. Morphogenesis: $\mathbf{c}_i^{t+1} = \tilde{\mathbf{h}}_i^t$ Lunar lander: $\mathbf{c}_i^{t+1} = W_C \, \tilde{\mathbf{h}}_i^t$ (state projection) and $\boldsymbol{\ell}_i^t = f_\text{act}(\tilde{\mathbf{h}}_i^t)$ (action logits). $f_\text{act}$: two-layer MLP (GELU), outputs 2D (noop, fire).