[Paper Review] FedDis: Disentangled Federated Learning for Unsupervised Brain Pathology Segmentation

Cosmin I. Bercea, Benedikt Wiestler+2 more|arXiv (Cornell University)|Mar 5, 2021

Medical Imaging and AnalysisEngineering25 citations

TL;DR

FedDis disentangles model parameters into shape and appearance, sharing only shape in federated learning to mitigate non-IID domain shifts, and achieves improved unsupervised brain pathology segmentation on multiple datasets.

ABSTRACT

In recent years, data-driven machine learning (ML) methods have revolutionized the computer vision community by providing novel efficient solutions to many unsolved (medical) image analysis problems. However, due to the increasing privacy concerns and data fragmentation on many different sites, existing medical data are not fully utilized, thus limiting the potential of ML. Federated learning (FL) enables multiple parties to collaboratively train a ML model without exchanging local data. However, data heterogeneity (non-IID) among the distributed clients is yet a challenge. To this end, we propose a novel federated method, denoted Federated Disentanglement (FedDis), to disentangle the parameter space into shape and appearance, and only share the shape parameter with the clients. FedDis is based on the assumption that the anatomical structure in brain MRI images is similar across multiple institutions, and sharing the shape knowledge would be beneficial in anomaly detection. In this paper, we leverage healthy brain scans of 623 subjects from multiple sites with real data (OASIS, ADNI) in a privacy-preserving fashion to learn a model of normal anatomy, that allows to segment abnormal structures. We demonstrate a superior performance of FedDis on real pathological databases containing 109 subjects; two publicly available MS Lesions (MSLUB, MSISBI), and an in-house database with MS and Glioblastoma (MSI and GBI). FedDis achieved an average dice performance of 0.38, outperforming the state-of-the-art (SOTA) auto-encoder by 42% and the SOTA federated method by 11%. Further, we illustrate that FedDis learns a shape embedding that is orthogonal to the appearance and consistent under different intensity augmentations.

Motivation & Objective

• Motivate privacy-preserving learning for brain pathology segmentation across multiple institutions.
• Propose FedDis to disentangle shape and appearance and share only shape parameters.
• Demonstrate improved anomaly segmentation performance on real pathological datasets.
• Show that FedDis learns a shape embedding orthogonal to appearance and robust to intensity shifts.

Proposed method

• Disentangled auto-encoder architecture separating shape and appearance parameters.
• Federated training where global aggregation averages only the shape parameters across clients.
• Latent contrastive loss LCL to enforce shape consistency (SCL) and latent orthogonality (LOL) between shape and appearance, via KL-divergence terms.
• Reconstruction-based unsupervised anomaly segmentation using residuals x - x_rec as anomaly evidence.
• Hyperparameters: alpha for reconstruction vs LCL balance; beta for weighting SCL vs LOL within LCL; gamma-shift augmentations for shape consistency.

Experimental results

Research questions

• RQ1 Can FedDis mitigate cross-site domain shifts in federated unsupervised brain anomaly segmentation by disentangling shape and appearance?
• RQ2 Does sharing only shape parameters improve generalization to unseen scanners and datasets compared to standard federated methods?
• RQ3 Do the latent representations exhibit a shape embedding orthogonal to appearance and robustness to intensity changes?
• RQ4 What is the impact of the latent contrastive loss components (SCL and LOL) on segmentation performance and disentanglement?

Key findings

• FedDis achieves an average Dice score of 0.38 on the main anomaly segmentation task, outperforming SOTA auto-encoder by 42% and the SOTA federated method by 11%.
• FedDis outperforms several baselines including FedAvg, FedVC, SiloBN and FedGN on real pathological datasets.
• FedDis learns a shape embedding orthogonal to appearance and stable under intensity shifts, as shown by latent analysis and ablations.
• Ablation indicates the latent contrastive loss is essential for effective disentanglement and performance gains.
• Compared to data-centralized training, FedDis offers strong segmentation performance with privacy-preserving distributed data.