OctreeNCA: Scalable Medical Segmentation

• OctreeNCA is a segmentation architecture that extends Neural Cellular Automata with a multiscale octree to process large medical images and videos in a single pass.
• The fused CUDA inference kernel achieves significant memory efficiency, using 90% less VRAM than UNet and enabling segmentation of 184.9 MP images and minute-long videos.
• The design trades slight accuracy limits on small structures for remarkable VRAM reduction, speed gains, and potential for extensions such as learned pooling or attention mechanisms.

OctreeNCA is a segmentation architecture for very large medical images and videos that extends the Neural Cellular Automaton (NCA) with a multiscale octree and a fused CUDA inference kernel. Its stated purpose is to preserve full spatial or temporal context by inferring large inputs at once rather than through patch-wise or frame-wise decomposition, while avoiding the VRAM scaling behavior of UNets and Vision Transformers. In the reported configuration, OctreeNCA performs single-pass segmentation of 184.9 Megapixel pathology images and 76.2-second surgical videos on a single NVIDIA RTX 4090 24 GiB, with evaluation-time VRAM occupancy described as 90% less than a UNet (Lemke et al., 9 Aug 2025).

1. Problem setting and design rationale

The motivating problem is segmentation of large inputs such as prostate MRIs, pathology slices, and videos of surgery. The central constraint is GPU VRAM: if a model is asked to segment a large image volume or a long video sequence in a single pass, architectures such as UNets or Vision Transformers scale poorly in memory and often force patch-wise or frame-wise inference. The paper frames that compromise as detrimental to both global consistency and inference speed.

OctreeNCA starts from a different premise. A standard NCA is lightweight and, by construction, size-invariant, so it does not require a fixed input resolution. However, the same locality that makes an NCA scalable also limits it: with only local communication rules, it lacks global knowledge. OctreeNCA addresses that limitation by generalizing the neighborhood definition with an octree data structure, so that long-range context can be traversed efficiently at coarse scales and then propagated back to fine scales.

A recurring misconception is that size-invariance alone is sufficient for globally coherent inference. In the formulation summarized here, that is explicitly not the case: the base NCA is size-invariant, but still local-only. The octree is introduced precisely to add a mechanism for global context without abandoning the NCA update paradigm (Lemke et al., 9 Aug 2025).

2. Standard NCA formulation

At each pixel or voxel $i$, the model maintains a state vector $s_i^{(t)}\in\mathbb{R}^c$, with $c=16$. The first channels are the input image and the remaining channels are hidden state. The local neighborhood is the $3\times 3$ grid in 2D or the $3\times 3\times 3$ grid in 3D.

Perception is implemented by a learnable convolution applied per-channel:

$p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$

The update rule concatenates the current cell state and its perception, then applies a two-layer MLP with a ReLU nonlinearity in between, producing an additive update $u_i^{(t)}\in\mathbb{R}^c$. The paper gives the full update as

$$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$

$$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$

with $W_1\in\mathbb{R}^{64\times(c+c)}$ and $s_i^{(t)}\in\mathbb{R}^c$0. The asynchronous masking is described as a fire-rate mechanism: each element of the update is zeroed out independently with probability $s_i^{(t)}\in\mathbb{R}^c$1.

This formulation preserves the canonical NCA property that the same local rule is applied over a grid of arbitrary size. The architectural intervention in OctreeNCA is therefore not a replacement of the update dynamics, but a restructuring of how those dynamics are deployed across scales.

3. Octree generalization and multiscale information flow

The octree is fixed to $s_i^{(t)}\in\mathbb{R}^c$2 levels. Level $s_i^{(t)}\in\mathbb{R}^c$3 is the full-resolution input, and level $s_i^{(t)}\in\mathbb{R}^c$4 is a coarsening of the original image by a factor of $s_i^{(t)}\in\mathbb{R}^c$5 in each spatial dimension. Formally, if $s_i^{(t)}\in\mathbb{R}^c$6 is the original image of size $s_i^{(t)}\in\mathbb{R}^c$7, then

$s_i^{(t)}\in\mathbb{R}^c$8

Each level $s_i^{(t)}\in\mathbb{R}^c$9 has a separate NCA with its own weights and a state grid $c=16$0. Global context is built by starting at the coarsest level, where the grid is small enough that a modest number of local NCA updates can diffuse information over very long effective distances. Hidden channels from that coarse level are then upsampled by nearest-neighbor with factor $c=16$1 and inserted into the next finer level, whose color channels are set from the image pyramid.

The forward-pass logic is given procedurally. First the octree of images is built. Then the top-level state $c=16$2 is initialized by concatenating $c=16$3 with zero hidden channels. For each level from $c=16$4 down to $c=16$5, the model runs $c=16$6 NCA steps; if $c=16$7, the hidden part of $c=16$8 is upsampled and concatenated with $c=16$9 to initialize $3\times 3$0. The final segmentation is read from a segmentation channel in $3\times 3$1.

This suggests that OctreeNCA treats “neighborhood” in an effective rather than purely geometric sense. The per-level NCA still uses local kernels, but because coarse levels aggregate larger receptive fields, information can be transported globally before fine-scale refinement. The paper describes this as a generalized neighborhood definition that enables efficient traversal of global knowledge (Lemke et al., 9 Aug 2025).

4. CUDA realization and memory complexity

A central technical contribution is the CUDA implementation of NCA inference. The paper argues that standard deep learning frameworks are optimized for large multilayer networks and therefore do not fully exploit the advantages of NCAs. In a traditional PyTorch NCA implementation, intermediate tensors are allocated for padded inputs, convolution outputs, MLP activations, Bernoulli masks, ReLU outputs, and related temporaries. The reported empirical peak is approximately $3\times 3$2 channels times $3\times 3$3 pixels.

The OctreeNCA CUDA kernel fuses all per-cell operations—convolution, affine layers, ReLU, masking, and addition—into a single kernel. The only global memories are the input state and output state, each with $3\times 3$4 channels in the reported configuration. The $3\times 3$5-channel hidden activations of the MLP live in thread-local registers rather than global VRAM. The empirical peak is therefore $3\times 3$6 channels times $3\times 3$7 pixels.

The corresponding asymptotic memory factors are stated as follows:

• OctreeNCA, CUDA optimized: $3\times 3$8
• Naïve NCA: $3\times 3$9
• UNet with $3\times 3\times 3$0 resolutions and base channels $3\times 3\times 3$1: $3\times 3\times 3$2

The paper characterizes this as roughly a $3\times 3\times 3$3 larger factor for the UNet than for the optimized OctreeNCA. In practical terms, this memory behavior underlies the claim that OctreeNCA can segment 184 Megapixel pathology slices or 1-minute surgical videos at once, and that it occupies 90% less VRAM than a UNet during evaluation (Lemke et al., 9 Aug 2025).

5. Training regime and empirical results

The reported hyper-parameters are shared across tasks: fire-rate $3\times 3\times 3$4, $3\times 3\times 3$5 state channels, hidden MLP size $3\times 3\times 3$6, and kernel size $3\times 3\times 3$7 or $3\times 3\times 3$8 for 3D. Optimization uses Adam with $3\times 3\times 3$9, $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$0, and initial learning rate $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$1, followed by exponential decay $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$2 each step. An exponential moving average of weights is maintained according to $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$3, and evaluation uses $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$4. The loss is a combined cross-entropy and Dice objective,

$p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$5

Training data handling depends on modality. For radiology inputs of size $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$6 and surgical inputs of size $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$7, the full volume is too large during training, so random patches are extracted on the two coarsest levels, with examples given as $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$8 or $p_i^{(t)} = \mathrm{conv}_k\bigl(s^{(t)}\bigr)[i].$9. For pathology, approximately $u_i^{(t)}\in\mathbb{R}^c$0 images downsampled to $u_i^{(t)}\in\mathbb{R}^c$1 are used, and training is done on $u_i^{(t)}\in\mathbb{R}^c$2 patches containing epithelium.

On a single NVIDIA RTX 4090 24 GiB, the maximum single-pass inference and speed results are:

• UNet: 8.3 MP pathology, 9.8 MP/s, 3.8 s video, 112 fps
• Med/M3D-NCA: 28.7 MP pathology, 7.8 MP/s, 11.5 s video, 108 fps
• OctreeNCA (PyTorch): 28.7 MP pathology, 13.9 MP/s, 11.5 s video, 219 fps
• OctreeNCA-CUDA: 184.9 MP pathology, 77.4 MP/s, 76.2 s video, 453 fps

The pathology figure corresponds to the paper’s “single-pass 184 MP segmentation” claim. The 76.2-second video figure substantiates the statement that videos of more than one minute can be processed at once.

On a Raspberry Pi 4 with 2 GiB RAM, OctreeNCA is also reported to be feasible. For square images with side lengths $u_i^{(t)}\in\mathbb{R}^c$3, $u_i^{(t)}\in\mathbb{R}^c$4, $u_i^{(t)}\in\mathbb{R}^c$5, and $u_i^{(t)}\in\mathbb{R}^c$6 pixels, OctreeNCA inference times are $u_i^{(t)}\in\mathbb{R}^c$7 s, $u_i^{(t)}\in\mathbb{R}^c$8 s, $u_i^{(t)}\in\mathbb{R}^c$9 s, and $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$0 s, compared with $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$1 s, $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$2 s, $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$3 s, and $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$4 s for Med-NCA, and $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$5 s, $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$6 s, $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$7 s, and no reported $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$8-pixel result for UNet.

Accuracy and training VRAM usage reveal a different trade-off. On prostate MRI, OctreeNCA uses 2.18 GiB VRAM, has 17,600 parameters, and reaches Dice $$\eta_i^{(t)} \sim \mathrm{Bernoulli}(0.5)\in\{0,1\}^c$$9; M3D-NCA uses 14.34 GiB and reaches $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$0; UNet uses 7.79 GiB and reaches $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$1; nnUNet uses 7.35 GiB and reaches $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$2. On pathology epithelium segmentation, OctreeNCA uses 2.25 GiB, has 15,520 parameters, and reaches Dice $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$3; Med-NCA reaches $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$4, UNet $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$5, EfficientNet $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$6, and MobilenetV2 $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$7. On cholecystectomy video segmentation, measured as mean Dice across abdominal wall, liver, fat, grasper, and gallbladder, OctreeNCA uses 2.41 GiB, has 16,960 parameters, and reaches $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$8; M3D-NCA reaches $$s_i^{(t+1)} = s_i^{(t)} + \eta_i^{(t)}\odot \Bigl(W_2\;\mathrm{ReLU}\bigl(W_1\,[\,s_i^{(t)};\,p_i^{(t)}\,]\bigr)\Bigr),$$9; UNet $W_1\in\mathbb{R}^{64\times(c+c)}$0; ResNet18+U $W_1\in\mathbb{R}^{64\times(c+c)}$1; and prompt-based SAM2 $W_1\in\mathbb{R}^{64\times(c+c)}$2 (Lemke et al., 9 Aug 2025).

6. Limitations, misconceptions, and possible extensions

The paper lists several benefits directly: true size-invariance, processing of 184 MP images and more than 1-minute videos in one pass on a single 24 GiB GPU, very low VRAM of approximately 2 GiB, parameter count of approximately 16 k, inference speed of 77 MP/s on an RTX 4090, feasibility on a $50 Raspberry Pi, and built-in multi-scale global context through a simple octree of NCAs.

The limitations are equally explicit. Accuracy on small structures, with the grasper given as an example, still lags the best CNNs or promptable models. Per-level NCAs must be recompiled for each $W_1\in\mathbb{R}^{64\times(c+c)}$3 combination, although the paper describes this as minor. The absence of explicit normalization is associated with slightly lower top-end accuracy than large CNNs, with nnUNet on MRI given as the example.

These results help clarify another common misunderstanding: low memory usage and single-pass scalability do not imply state-of-the-art accuracy across all benchmarks. In the reported experiments, OctreeNCA is often competitive and sometimes better than UNet baselines under strong memory constraints, but it does not uniformly exceed larger CNN systems or promptable segmentation models. The contribution is therefore not only an accuracy claim; it is an operating-point claim about the trade-off among context, VRAM, speed, and parameter count.

The paper also names several possible extensions: incorporating attention mechanisms or learned pooling in the octree, jointly learning the downsampling and upsampling kernels instead of using fixed average pooling, extending the method to 4D octrees over space and time for very long videos, and combining the model with prompt signals as in SAM for interactive segmentation. A plausible implication is that the octree scaffold is intended as a general multiscale communication mechanism rather than a fixed endpoint, with the present formulation demonstrating that NCA-style dynamics can be scaled to consumer-hardware segmentation regimes that are difficult for more conventional architectures (Lemke et al., 9 Aug 2025).