Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Steven Korevaar, Ruwan Tennakoon, Alireza Bab-Hadiashar

TL;DR

The paper questions whether large convolutional networks are required for robust medical image segmentation under distribution shifts. It compares a lightweight Neural Cellular Automata (NCA) against U‑Net variants on RETOUCH and MNM2, using a consistent training protocol and Dice loss. The findings show that NCA achieves superior out-of-distribution generalization despite far fewer parameters (about 20,608) compared with U‑Net's ~14.8 million, though IID performance varies by dataset; stability issues emerge on some domains. This work suggests NCAs offer a promising, deployable approach for domain-generalizable segmentation, warranting further exploration of stability improvements and extensions to other architectures like GNNs and transformers.

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.

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

TL;DR

The paper questions whether large convolutional networks are required for robust medical image segmentation under distribution shifts. It compares a lightweight Neural Cellular Automata (NCA) against U‑Net variants on RETOUCH and MNM2, using a consistent training protocol and Dice loss. The findings show that NCA achieves superior out-of-distribution generalization despite far fewer parameters (about 20,608) compared with U‑Net's ~14.8 million, though IID performance varies by dataset; stability issues emerge on some domains. This work suggests NCAs offer a promising, deployable approach for domain-generalizable segmentation, warranting further exploration of stability improvements and extensions to other architectures like GNNs and transformers.

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.
Paper Structure (10 sections, 3 equations, 4 figures, 2 tables)

This paper contains 10 sections, 3 equations, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Methodology
  4. Network Architectures
  5. Datasets
  6. Results
  7. Discussion & Future Work
  8. Conclusion
  9. Acknowledgments.
  10. Disclosure of Interests.

Figures (4)

  • Figure 1: The Cell update model of the Neural Cellular Automata used in this work.
  • Figure 2: Sample images from the RETOUCH dataset: one input example (top row) and one segmentation map (bottom row) for each type of scanner in the dataset. Intraretinal fluid (red), Subretinal fluid (green), and Pigment Epithelial Detachment (blue).
  • Figure 3: Sample images from the MnM2 dataset: one input example (top row) and one segmentation map (bottom row) for each type of scanner in the dataset. The left (LV) and right ventricle (RV) blood pools are blue and red, respectively, and the left ventricular myocardium (MYO) is green.
  • Figure 4: Randomly sample outputs from the out-of-distribution domain during testing. The MNM2 sample comes from the Philips scanner, and the RETOUCH sample from the Cirrus scanner. MNM2 labels: left and right ventricle blood pools are blue and red, and the left ventricular myocardium is green. RETOUCH labels: Intraretinal fluid is red, Subretinal fluid is green, and Pigment Epithelial Detachment is blue.