AuraViT-FL: A Resource-Efficient 2D Hybrid Transformer Framework for Federated Lung Tumor Segmentation

Abstract

Accurate lung tumor segmentation using computed tomography (CT) scans is needed for efficient tumor treatment. However, the development of deep learning models is often constrained by strict patient privacy regulations that limit direct data sharing. This work presents a system that enables multi-institutional collaboration while training high-quality lung tumor segmentation models without requiring access to sensitive patient data. The proposed framework features the AuraViT suite, which includes the standard AuraViT—a hybrid model with 136 million parameters that combines a Vision Transformer (ViT) encoder, Atrous Spatial Pyramid Pooling (ASPP), and attention-gated residual connections—and the Lightweight AuraViT (LAURA) family (Small, Tiny, and Mobile). These variants are designed for resource-constrained environments and potential edge deployment scenarios. Training is conducted on publicly available datasets (MSD Lung and NSCLC) in a simulated five-client federated learning setup that emulates collaboration among institutions while ensuring patient privacy. The framework uses a federated learning setup with FedProx, adaptive weighted aggregation, and a dynamic virtual client strategy to handle data and system differences. The framework is further evaluated through ablation studies on model architecture and feature importance. The results show that the standard AuraViT-FL achieves a global mean Dice score of 80.81%, while maintaining performance close to centralized training. Additionally, the LAURA variations show a better trade-off between accuracy and efficiency. Notably, the Mobile variant with ∼5 M parameters reduces model complexity by over 96% while maintaining competitive performance (82.96% Dice on MSD Lung).