Federated Multi-Task Learning for Bladder Tumor Segmentation and MIBC Classification Using a Hybrid CNN-Transformer Architecture

Malhar Udmale; Divyanshu Dwivedi; Aarohi Dhand; Sachin Dudda Nagaraju; Mayank Rai; Bagesh Kumar

Source-linked AI summary

Federated Multi-Task Learning for Bladder Tumor Segmentation and MIBC Classification Using a Hybrid CNN-Transformer Architecture

Malhar Udmale, Divyanshu Dwivedi, Aarohi Dhand, Sachin Dudda Nagaraju, Mayank Rai, Bagesh Kumar

arXiv:2608.30458v1Aug 31, 2026cs.CV

TL;DR

Robust bladder tumor segmentation and MIBC classification across institutions are difficult because patient data cannot be centrally pooled and MRI characteristics vary across scanners. The paper addresses this with federated multi-task learning using a Swin Hybrid architecture and augmentation strategies. Geo+Elastic produced the strongest reported federated balance, while the framework supported collaborative analysis without centralizing patient data.

Problem
Privacy constraints and scanner heterogeneity limit robust multi-center bladder tumor segmentation and MIBC classification from T2-weighted MRI.
Method
A federated multi-task framework jointly trains segmentation and MIBC classification using a ResNet-34/Swin-Tiny hybrid with segmentation-guided classification.
Results
Geo+Elastic augmentation achieved a federated DSC of 0.8100, patient-level AUC of 0.8931, and highest combined score of 0.8474.
Takeaways & Limitations
Joint segmentation and classification can be performed across multiple institutions with federated training without centralizing patient data.

Abstract

from arXiv · show

Accurate bladder tumor segmentation and assessment of mus- cle invasion from T2-weighted MRI are important for treatment plan- ning, but developing robust models across institutions is challenging be- cause patient data cannot be centrally pooled and imaging characteristics vary across scanners and acquisition protocols. We propose a federated multi-task learning framework for joint bladder tumor segmentation and MIBC/NMIBC classification across four clinical centers. The proposed Swin Hybrid model combines a ResNet-34 branch for local texture and boundary information with a Swin-Tiny Transformer for global anatomi- cal context. A segmentation-guided classification mechanism further uses tumor localization information to support MIBC prediction. We also investigate several augmentation strategies under both centralized and federated training to improve robustness to multi-center variability. Ex- periments on the FedBCa dataset show that the Swin Hybrid provides the best overall balance between segmentation and classification among the evaluated architectures. Under federated training, Geo+Elastic aug- mentation achieved a DSC of 0.8100 and a patient-level AUC of 0.8931, yielding the highest combined score of 0.8474. These results demonstrate that joint segmentation and classification can be effectively performed across multiple institutions using federated training without centralizing patient data.

1 Introduction

Bladder cancer MRI analysis requires accurate tumor staging and segmentation, yet privacy constraints and scanner heterogeneity complicate robust multi-center modeling. The paper proposes federated multi-task learning with a hybrid CNN–Transformer model and evaluates augmentation strategies for cross-center robustness.

Motivation: Accurate NMIBC/MIBC staging and tumor segmentation support treatment planning, monitoring, and quantitative assessment from T2-weighted MRI.MIBC typically requires more aggressive treatment than NMIBC, while manual delineation and staging are labor-intensive and variable across observers.
Challenge: Federated learning enables collaborative training across hospitals without transferring patient data, but scanner and protocol differences create non-IID domain shifts.These shifts can reduce model robustness and generalizability across clinical sites.
Approach: The proposed framework jointly optimizes bladder tumor segmentation and MIBC classification across four heterogeneous clinical centers using FedAvg.The framework is designed for privacy-preserving collaboration across institutions.
Approach: The Swin Hybrid combines a ResNet-34 branch for local texture and boundary features with a Swin-Tiny Transformer branch for global anatomical context.Segmentation-guided attention connects the dual-stream encoder to classification.
Results: Geo+Elastic augmentation achieved a DSC of 0.8100 and AUC of 0.8931 under federated learning, identifying it as the evaluated strategy with the strongest reported federated result.The study systematically evaluates six augmentation strategies under centralized and federated training.

2 Related Works

The related-work discussion motivates multi-task learning for bladder cancer MRI by linking tumor localization with muscle-invasion assessment. The proposed formulation shares representations and lets segmentation features guide MIBC classification while incorporating global anatomical context.

Multi-task learning: Multi-task learning jointly optimizes related objectives in a shared representation space so information from one task can support the other.This contrasts with treating segmentation and classification as independent single-task problems.
Prior work: Prior medical-imaging studies report that cross-task feature exchange and multi-scale multi-task architectures can improve task-specific representations and feature utilization.These studies provide precedent for modeling interactions between related prediction tasks.
Clinical relationship: Bladder tumor localization and MIBC assessment are anatomically linked because invasion depends on tumor relationships with the muscular bladder wall.Segmentation delineates tumor extent and boundaries, while MIBC assessment evaluates disruption or extension into muscle.
Proposed formulation: The framework uses shared multi-task representations and segmentation-guided feature interaction, allowing tumor localization to guide the staging decision.The classification branch combines segmentation-derived features with global anatomical context.

3 Methodology

The methodology formulates joint bladder tumor segmentation and MIBC classification as a privacy-preserving federated learning problem across four hospitals. Each client trains a hybrid CNN–Transformer model locally, applies standardized preprocessing and augmentation, and shares only parameters for FedAvg aggregation.

Federated learning framework: Each hospital trains on private T2-weighted MRI data while transmitting only updated model parameters for sample-size-weighted FedAvg aggregation.The global model is redistributed and training repeats for 50 communication rounds with all four hospitals participating.
Federated learning framework: The shared model jointly predicts tumor segmentation masks and MIBC/NMIBC labels using a sample-size-weighted objective.The framework treats segmentation and classification as coupled outputs of one shared model.
Data preprocessing and augmentation: The preprocessing pipeline standardizes heterogeneous MRI inputs through bias correction, intensity clipping, 2.5D five-slice extraction, resampling, and fixed-size cropping or padding.Clients also apply Geo+Elastic augmentation combining flipping, rotation, and elastic deformation to model positioning and morphology variation.
Multi-task hybrid CNN–Transformer model: The Swin Hybrid encoder combines ResNet-34 local texture and boundary features with Swin-Tiny global context, followed by task-specific heads.An attention-gated decoder produces the tumor mask, while classification uses segmentation-guided pooling of Transformer features.
Multi-task hybrid CNN–Transformer model: Segmentation-guided attention uses the predicted tumor probability map to pool deep Transformer features before MIBC classification.The segmentation map is detached during pooling, preventing classification gradients from propagating through the segmentation prediction.

4 Experiments

Experiments use the heterogeneous FedBCa dataset and consistent center-wise patient partitions to compare architectures, augmentation strategies, and centralized versus federated training. Evaluation reports segmentation and patient-level classification metrics separately, with a combined score used for model selection.

Dataset and splits: The FedBCa dataset contains 275 T2-weighted MRI scans from four medical centers, each paired with tumor masks and pathological MIBC/NMIBC labels.Scanner and acquisition differences provide natural cross-center heterogeneity for evaluating federated learning.
Dataset and splits: Approximately 71% of patients were assigned to training, 10% to validation, and 19% to testing through independent stratified splits within each center.The same center-wise partitions were used in centralized and federated experiments without patient overlap between splits.
Evaluation metrics: Segmentation was evaluated with DSC on tumor-containing slices, while classification used patient-level AUC after averaging slice-level MIBC probabilities.The combined validation score was used only for checkpoint selection; DSC and AUC remained the primary reported metrics.
Experimental protocol: Five candidate architectures were compared under centralized training using identical patient partitions and evaluation metrics before selecting a model for augmentation experiments.The candidates included three CNN-based alternatives, a multi-task EfficientNet-B0, and the proposed Swin Hybrid.
Experimental protocol: Seven augmentation configurations were evaluated with the selected Swin Hybrid under centralized training and then repeated with FedAvg across all four centers.The same architecture, partitions, and evaluation protocol enabled direct augmentation comparisons across training settings.

5 Results

The proposed Swin Hybrid offered the strongest balance between bladder tumor segmentation and MIBC classification, and Geo+Elastic was the best-balanced augmentation under federated training. Federated performance remained competitive despite modest reductions from centralized training.

Architecture comparison: 0.8214 patient-level AUC and 0.8138 combined score made the multi-task Swin Hybrid the best overall-balanced architecture under centralized training.ResNet50 U-Net achieved the highest single-task segmentation DSC, but Swin Hybrid was selected for subsequent augmentation experiments.
Centralized augmentation: 0.8487 DSC was highest for Augmentation Stack under centralized training, while Geo+Elastic achieved the highest combined score at 0.8688.No Augmentation achieved the highest AUC of 0.9286, but Geo+Elastic provided the strongest balance between segmentation and classification.
Federated augmentation: 0.8474 combined score and 0.8931 AUC made Geo+Elastic the strongest federated augmentation, despite Augmentation Stack achieving the highest DSC of 0.8293.Geo+Elastic was selected for the final federated Swin Hybrid model because its lower DSC was offset by stronger patient-level classification.
Centralized versus federated performance: 0.8100 DSC and 0.8931 AUC under federated training were modestly below Geo+Elastic’s centralized values of 0.8272 and 0.9196, respectively.Geo+Elastic nevertheless retained the highest combined performance across both settings, with scores of 0.8688 centrally and 0.8474 federally.
Discussion and impact of multi-task learning: The Swin Hybrid addresses the CNN-versus-global-context trade-off by combining ResNet-34 local features with Swin-Tiny anatomical modeling and tumor-focused attention.Its attention-gated pooling restricts classification to predicted tumor regions rather than irrelevant bladder-wall background.
Discussion and impact of multi-task learning: 0.8474 combined score with 0.8931 AUC showed that Geo+Elastic was optimally balanced for federated settings despite the Augmentation Stack’s higher centralized DSC of 0.8487.Elastic deformation simulated bladder-volume and irregular-tumor variations and was described as a regularizer against non-IID client drift.

6 Conclusion

The study introduced a federated multi-task framework for joint bladder tumor segmentation and MIBC classification from multi-center T2-weighted MRI. The Swin Hybrid and Geo+Elastic augmentation supported balanced performance across centralized and federated settings without centralizing patient data.

Conclusion: The framework jointly performs bladder tumor segmentation and MIBC classification from multi-center T2-weighted MRI.It is designed for collaborative analysis across institutions.
Conclusion: The Swin Hybrid combines CNN-based local feature extraction with Transformer-based global contextual modeling for balanced segmentation and classification.The conclusion describes this as a strong balance between the two clinical tasks.
Conclusion: Geo+Elastic was the most effective augmentation strategy across centralized and federated settings, particularly with heterogeneous multi-center data.The framework maintained competitive performance on both clinical tasks while patient data remained decentralized.