Unsupervised Training of Neural Cellular Automata on Edge Devices

TL;DR

This paper demonstrates that neural cellular automata can be trained directly on smartphones for lung X-ray segmentation, using an unsupervised method that improves accuracy and robustness without requiring extensive labeled data or high computational resources.

Contribution

It introduces a novel unsupervised training approach for neural cellular automata on edge devices, enabling effective medical image segmentation in resource-constrained settings.

Findings

Achieved 0.7 to 2.8% improvement in Dice accuracy over traditional methods.

Enhanced robustness with 5-20% accuracy gains in low-quality image scenarios.

Validated on three diverse multisite X-ray datasets.

Abstract

The disparity in access to machine learning tools for medical imaging across different regions significantly limits the potential for universal healthcare innovation, particularly in remote areas. Our research addresses this issue by implementing Neural Cellular Automata (NCA) training directly on smartphones for accessible X-ray lung segmentation. We confirm the practicality and feasibility of deploying and training these advanced models on five Android devices, improving medical diagnostics accessibility and bridging the tech divide to extend machine learning benefits in medical imaging to low- and middle-income countries (LMICs). We further enhance this approach with an unsupervised adaptation method using the novel Variance-Weighted Segmentation Loss (VWSL), which efficiently learns from unlabeled data by minimizing the variance from multiple NCA predictions. This strategy notably improves model adaptability and performance across diverse medical imaging contexts without the need for extensive computational resources or labeled datasets, effectively lowering the participation threshold. Our methodology, tested on three multisite X-ray datasets -- Padchest, ChestX-ray8, and MIMIC-III -- demonstrates improvements in segmentation Dice accuracy by 0.7 to 2.8%, compared to the classic Med-NCA. Additionally, in extreme cases where no digital copy is available and images must be captured by a phone from an X-ray lightbox or monitor, VWSL enhances Dice accuracy by 5-20%, demonstrating the method's robustness even with suboptimal image sources.