Generalization Capabilities of Neural Cellular Automata for Medical Image Segmentation: A Robust and Lightweight Approach

TL;DR

This paper investigates Neural Cellular Automata (NCA) as a lightweight alternative to U-Net for medical image segmentation, demonstrating superior generalization to out-of-distribution data while maintaining competitive in-domain performance.

Contribution

It introduces the use of NCA models, significantly smaller than U-Net, showing they can achieve better generalization in medical image segmentation tasks.

Findings

NCA outperforms traditional models on out-of-distribution data.

NCA maintains good in-domain performance despite smaller size.

NCA achieves larger receptive fields through recursive processes.

Abstract

In the field of medical imaging, the U-Net architecture, along with its variants, has established itself as a cornerstone for image segmentation tasks, particularly due to its strong performance when trained on limited datasets. Despite its impressive performance on identically distributed (in-domain) data, U-Nets exhibit a significant decline in performance when tested on data that deviates from the training distribution, out-of-distribution (out-of-domain) data. Current methodologies predominantly address this issue by employing generalization techniques that hinge on various forms of regularization, which have demonstrated moderate success in specific scenarios. This paper, however, ventures into uncharted territory by investigating the implications of utilizing models that are smaller by three orders of magnitude (i.e., x1000) compared to a conventional U-Net. A reduction of this size in U-net parameters typically adversely affects both in-domain and out-of-domain performance, possibly due to a significantly reduced receptive field. To circumvent this issue, we explore the concept of Neural Cellular Automata (NCA), which, despite its simpler model structure, can attain larger receptive fields through recursive processes. Experimental results on two distinct datasets reveal that NCA outperforms traditional methods in terms of generalization, while still maintaining a commendable IID performance.