Unsupervised Training of Neural Cellular Automata on Edge Devices

TL;DR

The paper tackles the global inequity in access to ML-assisted medical imaging by enabling on-device training of Neural Cellular Automata (NCA) for X-ray lung segmentation. It introduces Variance-Weighted Segmentation Loss (VWSL), an unsupervised adaptation method that leverages unlabeled data to fine-tune NCAs on local devices, with a formulation that weights per-pixel losses by variance and enforces consistency between dual outputs. Across three multisite datasets (Padchest, ChestX-ray8, MIMIC-III) the approach yields Dice-score improvements of 0.7 to 2.8% over Med-NCA and, notably, 5 to 20% gains for smartphone-captured X-ray images, while remaining feasible for on-device training. The authors provide an open-source framework and demonstrate deployment on five Android devices using TensorFlow Lite, illustrating a practical path to robust, low-resource lung diagnostics in remote settings.

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.

Figures (7)

• Figure 1: Neural Cellular Automata enable primary care professionals with data but no access to infrastructure, to train their models or fine-tune a pre-trained model on a smartphone.
• Figure 2: Our variance-based fine-tuning VWSL leverages mean predictions from the pre-trained Med-NCA, along with the corresponding variances. This process enables us to refine the model by targeting these predictions and modulating their impact based on variance. Such an approach ensures the model's adaptation to new domains while preserving crucial information.
• Figure 3: Variance of Med-NCA trained on ChestX-ray8 and evaluated on MIMIC and Padchest for segmentation of both lungs.
• Figure 4: Med-NCA trained on the ChestX8 dataset before and after being fine-tuned on the MIMIC dataset.
• Figure 5: Med-NCA trained on the Padchest dataset before and after being fine-tuned on pictures taken by a Pixel XL.