Latent Neural Cellular Automata for Resource-Efficient Image Restoration

TL;DR

This work introduces the Latent Neural Cellular Automata (LNCA) model, a novel architecture designed to address the resource limitations of neural cellular automata, and applies it in the context of image restoration, which aims to reconstruct high-quality images from their degraded versions.

Abstract

Neural cellular automata represent an evolution of the traditional cellular automata model, enhanced by the integration of a deep learning-based transition function. This shift from a manual to a data-driven approach significantly increases the adaptability of these models, enabling their application in diverse domains, including content generation and artificial life. However, their widespread application has been hampered by significant computational requirements. In this work, we introduce the Latent Neural Cellular Automata (LNCA) model, a novel architecture designed to address the resource limitations of neural cellular automata. Our approach shifts the computation from the conventional input space to a specially designed latent space, relying on a pre-trained autoencoder. We apply our model in the context of image restoration, which aims to reconstruct high-quality images from their degraded versions. This modification not only reduces the model's resource consumption but also maintains a flexible framework suitable for various applications. Our model achieves a significant reduction in computational requirements while maintaining high reconstruction fidelity. This increase in efficiency allows for inputs up to 16 times larger than current state-of-the-art neural cellular automata models, using the same resources.

Figures (4)

• Figure 1: LNCA structure. $+$ is an element-wise addition. $S$ is a channel switch. On the left, the Encoder, which is composed of adjusted convolution blocks, outputs the skip-connection $\bar{x}$ and the latent tensor $\hat{x}$. In the middle, the NCA block processes the latent tensor to remove the corruption, obtaining $\hat{x}^\prime$. On the right, the Decoder uses both the skip-connection and the latent tensor to reconstruct the final output $x^\prime$. The switch path around the NCA is used to bypass the NCA computation during the AE training, as described in Section \ref{['subsubse:autoencoder_training']}.
• Figure 2: Training memory trend of the tested models across the different configurations.
• Figure 3: Training latency trend of the tested models across the different configurations.
• Figure 4: Inference latency trend of the tested models across the different configurations.