相似度感知蒸馏的统一弱监督个性化联邦图像分割
pFedWSD: unified weakly supervised personalized federated image segmentation via similarity-aware distillation
- 2024年29卷第3期 页码:620-636
纸质出版日期: 2024-03-16
DOI: 10.11834/jig.230295
移动端阅览
浏览全部资源
扫码关注微信
纸质出版日期: 2024-03-16 ,
移动端阅览
潘建珊, 林立, 吴洁伟, 刘翼翔, 陈孝华, 林其友, 黄建业, 唐晓颖. 2024. 相似度感知蒸馏的统一弱监督个性化联邦图像分割. 中国图象图形学报, 29(03):0620-0636
Pan Jianshan, Lin Li, Wu Jiewei, Liu Yixiang, Chen Xiaohua, Lin Qiyou, Huang Jianye, Tang Xiaoying. 2024. pFedWSD: unified weakly supervised personalized federated image segmentation via similarity-aware distillation. Journal of Image and Graphics, 29(03):0620-0636
目的
2
联邦学习允许多个机构在不侵犯数据隐私、安全的前提下协作训练强大的深度模型。现有多数联邦范式在处理多中心不同数据分布时性能通常会下降,且弱监督条件下的联邦范式鲜有研究,特别是各站点数据采用不同形式稀疏标注的情况。针对该问题,提出一种站点分布相似度感知知识蒸馏的统一弱监督个性化联邦学习框架(unified weakly supervised personalized federated image segmentation via similarity-aware distillation,pFedWSD),以应对多中心数据分布和标注上的差异。
方法
2
所提出的pFedWSD通过循环知识蒸馏为每个站点训练个性化模型,包含动态循环公共知识积累及个性化两个阶段。第1阶段以不确定度感知方式动态地排序每轮训练中各站点模型性能,并以循环知识蒸馏的形式积累公共知识;第2阶段通过批标准化层的统计信息来度量各站点间相似性并聚合得到各站点教师模型并进行知识蒸馏。在弱监督方面,引入门控条件随机场损失和树能量损失相结合的训练目标,以产生更为精确的伪标注监督信号。
结果
2
在眼底视杯视盘分割和视网膜中心凹无血管区分割两项任务中,pFedWSD的Dice系数和HD95(95% Hausdorff distance)指标均优于多种中心式联邦和个性化联邦方法,在两项任务中,Dice系数分别为90.38%和93.12%,相比于较先进的方法FedAP(federated learning with adaptive batchnorm for personalized healthcare)和FedALA(adaptive local aggregation for personalized federated learning)分别提升了1.67%和6.56%,性能接近于全监督集中式训练所得的模型。
结论
2
本文提出的弱监督个性化联邦学习框架能有效统一不同形式稀疏标注数据并对不同分布的各站点数据训练得到个性化模型,使各站点分割性能均得到显著提升。
Objective
2
Federated learning (FL) allows multiple healthcare institutions to collaboratively train a powerful deep learning model without compromising data privacy and security (i.e., centralizing data). However, employing a single model to accommodate the diverse data distributions from different sites is extremely challenging. Performance degradation is common for existing approaches when huge distribution gaps exist across sites. Additionally, previous works paid little attention to FL under weak supervision, especially under the supervision of different sparsely grained forms (i.e., point-, bounding box-, scribble-, block-wise). Weakly supervised FL is clinically practical but challenging. To address this issue, we propose a unified and weakly-supervised personalized FL framework named pFedWSD, targeting medical image segmentation and based on similarity-aware knowledge distillation across multiple sites. We aim to accommodate the domain gaps and annotation drifts across multiple sites and enhance the segmentation model’s performance for each site.
Method
2
The proposed pFedWSD trains a personalized model for each site via cyclic knowledge distillation, which consists of two stages: uncertainty-aware dynamic and cyclic common knowledge accumulation and similarity-aware personalization. In the first stage, during each training round, the performance of each site’s model is dynamically ranked in an uncertainty-aware manner, and common knowledge is accumulated in the form of cyclic knowledge distillation. In the second stage, the similarity between two sites is measured and aggregated based on the statistics from the batch normalization layers to attain a teacher model for each site and perform knowledge distillation. As for weakly-supervised learning, a combination of partial cross-entropy loss, gated conditional random field (CRF) loss, and tree energy loss is employed. Specifically, the partial cross-entropy loss is employed for supervising the annotated regions, ensuring informative guidance. The tree energy loss establishes pairwise affinities on the basis of the preserved characteristics of high and low semantic spatial structures for the same object. This approach, in conjunction with the model’s predictions, generates soft pseudo-labels for the unlabeled regions. Through continuous online training and refinement, the model’s predictions and the delivered pseudo-annotations gradually improve over time. Furthermore, the gated CRF loss serves as a regularization term, effectively curbing the potential issues of excessive expansion or contraction of the target regions’ pseudo-labels that may arise from solely employing the tree energy loss. This approach adeptly consolidates diverse sparsely annotated data for training, facilitating real-time generations of additional pseudo proposals, and consequently attaining exceptional segmentation performance without requiring supplementary supervised data, iterative optimization, nor time-intensive post-processing. To the best of our knowledge, pFedWSD is a pioneering weakly supervised personalized federated learning approach for medical image segmentation and adeptly implemented under heterogeneous annotation settings on multiple client devices.
Result
2
We create two datasets (from multiple publicly available datasets), each with five subsets serving as five different sites, for optic/disc cup (OD/OC) segmentation and retinal foveal avascular zone (FAZ) segmentation, respectively. Quantitative and qualitative experimental results show that pFedWSD outperforms representative state-of-the-art (SOTA) centralized and personalized FL methods in terms of Dice coefficients and HD95 statistics. The proposed pFedWSD achieves an average Dice coefficient of 90.38% on the OD/OC segmentation task, exhibiting a remarkable improvement of 1.67% over the previous best-performing method. Moreover, pFedWSD demonstrates a marginal difference of only 0.58% compared with local training under full supervision and a slight gap of merely 1.23% from centralized training under full supervision. Regarding the FAZ segmentation task, the proposed method achieves an impressive average Dice coefficient of 93.12%, showcasing a substantial improvement of 6.56% over the previous state-of-the-art method. Furthermore, pFedWSD has a marginal difference of 0.5% from local training under full supervision and a mere 0.86% difference from centralized training under full supervision.
Conclusion
2
The proposed weakly-supervised and personalized FL framework (pFedWSD) can effectively unify different forms of sparsely labeled data and train personalized models that adapt well to different data distributions, with an established superior segmentation performance. Our pFedWSD demonstrates its effectiveness through achieving optimal performance on both OD/OC and FAZ segmentation tasks across datasets from multiple centers, with its overall performance closely approaching that of local or centralized training using fully supervised labels. Extensive ablation experiments demonstrate the importance and efficacy of each stage in pFedWSD and each component in the weakly supervised composite objective. Moreover, through site-ablation experiments, we analyze the contribution of each site to the federation, providing valuable guidance for medical institutions regarding the appropriate data volume and the sparse annotation form in federated learning. Future research directions include the further reduction of the communication and computation overhead and the integration of universal large model training paradigms, like prompt learning, to concurrently foster our proposed framework’s generalization performance and adaptive personalization capacity toward diverse data distributions.
相似度感知知识蒸馏弱监督学习个性化联邦学习医学图像分割
similarity-awareknowledge distillationweakly supervised learningpersonalized federated learningmedical image segmentation
Agarwal A, Balaji J, Raman R and Lakshminarayanan V. 2020. The foveal avascular zone image database (FAZID)//Proceedings of 2020 SPIE 11510, Applications of Digital Image Processing XLIII. Vitual: SPIE: 507-512 [DOI: 10.1117/12.2567580http://dx.doi.org/10.1117/12.2567580]
Chen Y Q, Lu W, Qin X, Wang J D and Xie X. 2023. MetaFed: federated learning among federations with cyclic knowledge distillation for personalized healthcare. IEEE Transactions on Neural Networks and Learning Systems [DOI: 10.1109/TNNLS.2023.3297103http://dx.doi.org/10.1109/TNNLS.2023.3297103]
Collins L, Hassani H, Mokhtari A and Shakkottai S. 2021. Exploiting shared representations for personalized federated learning//Proceedings of the 38th International Conference on Machine Learning [s.l.]: PMLR: 2089-2099
Collins L, Hassani H, Mokhtari A and Shakkottai S. 2022. Fedavg with fine tuning: local updates lead to representation learning//Proceedings of the 36th International Conference on Advances in Neural Information Processing Systems. New Orleans, USA: Curran Associates Inc.: 10572-10586
Fumero F, Alayón S, Sanchez J L, Sigut J and Gonzalez-Hernandez M. 2011. RIM-ONE: an open retinal image database for optic nerve evaluation//The 24th International Symposium on Computer-based Medical Systems. Bristol, UK: IEEE: 1-6 [DOI: 10.1109/cbms.2011.5999143http://dx.doi.org/10.1109/cbms.2011.5999143]
Gallager R G, Humblet P A and Spira P M. 1983. A distributed algorithm for minimum-weight spanning trees. ACM Transactions on Programming Languages and Systems, 5(1): 66-77 [DOI: 10.1145/357195.357200http://dx.doi.org/10.1145/357195.357200]
Kendall A and Gal Y. 2017. What uncertainties do we need in Bayesian deep learning for computer vision?//Proceedings of the 31st International Conference on Neural Information Processing Systems. Long Beach, USA: Curran Associates Inc.: 5580-5590
Li D Q, Kar A, Ravikumar N, Frangi A F and Fidler S. 2020a. Federated simulation for medical imaging//Proceedings of the 23rd International Conference on Medical Image Computing and Computer-assisted Intervention. Lima, Peru: Springer: 159-168 [DOI: 10.1007/978-3-030-59710-8_16http://dx.doi.org/10.1007/978-3-030-59710-8_16]
Li M C, Zhang Y H, Ji Z X, Xie K R, Yuan S T, Liu Q H and Chen Q. 2020b. IPN-V2 and OCTA-500: methodology and dataset for retinal image segmentation [EB/OL]. [2023-05-15]. https://arxiv.org/pdf/2012.07261v1.pdfhttps://arxiv.org/pdf/2012.07261v1.pdf
Li T, Sahu A K, Zaheer M, Sanjabi M, Talwalkar A and Smith V. 2020c. Federated optimization in heterogeneous networks//Proceedings of 2020 Machine Learning and Systems. Austin, USA: mlsys.org: 429-450
Li W Q, Milletarì F, Xu D G, Rieke N, Hancox J, Zhu W T, Baust M, Cheng Y, Ourselin S, Cardoso M J and Feng A. 2019. Privacy-preserving federated brain tumour segmentation//Proceedings of the 10th International Workshop on Machine Learning in Medical Imaging. Shenzhen, China: Springer: 133-141 [DOI: 10.1007/978-3-030-32692-0_16http://dx.doi.org/10.1007/978-3-030-32692-0_16]
Li X X, Jiang M R, Zhang X F, Kamp M and Dou Q. 2021. FedBN: federated learning on Non-IID features via local batch normalization//Proceedings of the 9th International Conference on Learning Representations. Vienna, Austria: OpenReview.net: 1-12
Liang Z Y, Wang T C, Zhang X Y, Sun J and Shen J B. 2022. Tree energy loss: towards sparsely annotated semantic segmentation//Proceedings of 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition. New Orleans, USA: IEEE: 16907-16916 [DOI: 10.1109/cvpr52688.2022.01640http://dx.doi.org/10.1109/cvpr52688.2022.01640]
Lin L, Wang Z H, Wu J W, Huang Y J, Lyu J Y, Cheng P J, Wu J and Tang X Y. 2021. BSDA-Net: a boundary shape and distance aware joint learning framework for segmenting and classifying OCTA images//Proceedings of the 24th International Conference on Medical Image Computing and Computer-Assisted Intervention. Strasbourg, France: Springer: 65-75 [DOI: 10.1007/978-3-030-87237-3_7http://dx.doi.org/10.1007/978-3-030-87237-3_7]
Lin L, Wu J W, Cheng P J, Wang K and Tang X Y. 2020. BLU-GAN: bi-directional convLSTM U-Net with generative adversarial training for retinal vessel segmentation//Proceedings of the 1st BenchCouncil International Federated Intelligent Computing and Block Chain Conferences. Qingdao, China: Springer: 3-13 [DOI: 10.1007/978-981-16-1160-5_1http://dx.doi.org/10.1007/978-981-16-1160-5_1]
Liu Q D, Chen C, Qin J, Dou Q and Heng P A. 2021. FedDG: federated domain generalization on medical image segmentation via episodic learning in continuous frequency space//Proceedings of 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Nashville, USA: IEEE: 1013-1023 [DOI: 10.1109/cvpr46437.2021.00107http://dx.doi.org/10.1109/cvpr46437.2021.00107]
Lu W, Wang J D, Chen Y Q, Qin X, Xu R J, Dimitriadis D and Qin T. 2022. Personalized federated learning with adaptive batchnorm for healthcare. IEEE Transactions on Big Data, 14(8): 1-12 [DOI: 10.1109/tbdata.2022.3177197http://dx.doi.org/10.1109/tbdata.2022.3177197]
Ma Y H, Hao H Y, Xie J Y, Fu H Z, Zhang J, Yang J L, Wang Z, Liu J, Zheng Y L and Zhao Y T. 2021. ROSE: a retinal OCT-angiography vessel segmentation dataset and new model. IEEE Transactions on Medical Imaging, 40(3): 928-939 [DOI: 10.1109/tmi.2020.3042802http://dx.doi.org/10.1109/tmi.2020.3042802]
McMahan B, Moore E, Ramage D, Hampson S and Arcas B A. 2017. Communication-efficient learning of deep networks from decentralized data//Proceedings of the 20th International Conference on Artificial Intelligence and Statistics. Fort Lauderdale, USA: PMLR: 1273-1282
Obukhov A, Georgoulis S, Dai D X and Van Gool L. 2019. Gated CRF loss for weakly supervised semantic image segmentation[EB/OL]. [2023-05-15]. https://arxiv.org/pdf/1906.04651v1.pdfhttps://arxiv.org/pdf/1906.04651v1.pdf
Orlando J I, Fu H Z, Breda J B, van Keer K, Bathula D R, Diaz-Pinto A, Fang R G, Heng P A, Kim J, Lee J, Lee J, Li X X, Liu P, Lu S, Murugesan B, Naranjo V, Phaye S S R, Shankaranarayana S M, Sikka A, Son J, Van Den hengel A, Wang S J, Wu J Y, Wu Z F, Xu G H, Xu Y L, Yin P S, Li F, Zhang X L, Xu Y W and Bogunović H. 2020. REFUGE Challenge: a unified framework for evaluating automated methods for glaucoma assessment from fundus photographs. Medical Image Analysis, 59: #101570 [DOI: 10.1016/j.media.2019.101570http://dx.doi.org/10.1016/j.media.2019.101570]
Ren D W, Wang Q L, Wei Y C, Meng D Y and Zuo W M. 2022. Progress in weakly supervised learning for visual understanding. Journal of Image and Graphics, 27(6): 1768-1798
任冬伟, 王旗龙, 魏云超, 孟德宇, 左旺孟. 2022. 视觉弱监督学习研究进展. 中国图象图形学报, 27(6): 1768-1798 [DOI: 10.11834/jig.220178http://dx.doi.org/10.11834/jig.220178]
Ronneberger O, Fischer P and Brox T. 2015. U-Net: convolutional networks for biomedical image segmentation//Proceedings of the 18th International Conference on Medical Image Computing and Computer-Assisted Intervention. Munich, Germany: Springer: 234-241 [DOI: 10.1007/978-3-319-24574-4_28http://dx.doi.org/10.1007/978-3-319-24574-4_28]
Sivaswamy J, Krishnadas S R, Joshi G D, Jain M and Tabish A U S. 2014. Drishti-GS: retinal image dataset for optic nerve head (ONH) segmentation//Proceedings of the 11th IEEE International Symposium on Biomedical Imaging. Beijing, China: IEEE: 53-56 [DOI: 10.1109/ISBI.2014.6867807http://dx.doi.org/10.1109/ISBI.2014.6867807]
Song L, Li Y W, Li Z M, Yu G, Sun H B, Sun J and Zheng N N. 2019. Learnable tree filter for structure-preserving feature transform//Proceedings of the 33rd International Conference on Neural Information Processing Systems. Vancouver, Canada: Curran Associates Inc.: 153
Sun Y H, Shi Y H, Li M, Yang R Z and Si P B. 2023. Personalized federated learning method based on collation game and knowledge distillation. Journal of Electronics and Information Technology, 45: 3702-3709
孙艳华, 史亚会, 李萌, 杨睿哲, 司鹏搏. 2023. 基于合作博弈和知识蒸馏的个性化联邦学习算法. 电子与信息学报, 45: 3702-3709 [DOI: 10.11999/JEIT221223http://dx.doi.org/10.11999/JEIT221223]
Tang M, Djelouah A, Perazzi F, Boykov Y and Schroers C. 2018. Normalized cut loss for weakly-supervised CNN segmentation//Proceedings of 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Salt Lake City, USA: IEEE: 1818-1827 [DOI: 10.1109/cvpr.2018.00195http://dx.doi.org/10.1109/cvpr.2018.00195]
Tian X, Wang L and Ding Q. 2019. Review of image semantic segmentation based on deep learning. Journal of Software, 30(2): 440-468
田萱, 王亮, 丁琪. 2019. 基于深度学习的图像语义分割方法综述. 软件学报, 30(2): 440-468 [DOI: 10.13328/j.cnki.jos.005659http://dx.doi.org/10.13328/j.cnki.jos.005659]
Voigt P and von dem Bussche A. 2017. The EU General Data Protection Regulation (GDPR). A Practical Guide. Cham: Springer [DOI: 10.1007/978-3-319-57959-7http://dx.doi.org/10.1007/978-3-319-57959-7]
Wang K K, Mathews R, Kiddon C, Eichner H, Beaufays F and Ramage D. 2019. Federated evaluation of on-device personalization [EB/OL]. [2023-05-15]. https://arxiv.org/pdf/1910.10252.pdfhttps://arxiv.org/pdf/1910.10252.pdf
Wang Y F, Shen Y Q, Yuan M, Xu J, Yang B, Liu C, Cai W J, Cheng W J, Wang W. 2021. A deep learning-based quality assessment and segmentation system with a large-scale benchmark dataset for optical coherence tomographic angiography image [EB/OL]. [2023-05-15]. https://arxiv.org/pdf/2107.10476.pdfhttps://arxiv.org/pdf/2107.10476.pdf
Wu J D, Fang H H, Li F, Fu H Z, Lin F B, Li J C, Huang Y, Yu Q J, Song S F, Xu X X, Xu Y Y, Wang W S, Wang L X, Lu S, Li H Q, Huang S H, Lu Z C, Ou C B, Wei X F, Liu B Y, Kobbi R, Tang X Y, Lin L, Zhou Q, Hu Q, Bogunović H, Orlando J I, Zhang X L and Xu Y W. 2023. GAMMA challenge: glaucoma grAding from multi-modality imAges. Medical Image Analysis, 90: #102938
Yang Q, Liu Y, Chen T J and Tong Y X. 2019. Federated machine learning: concept and applications. ACM Transactions on Intelligent Systems and Technology, 10(2): 1-19 [DOI: 10.1145/3298981http://dx.doi.org/10.1145/3298981]
Zhang J Q, Hua Y, Wang H, Song T, Xue Z G, Ma R H and Guan H B. 2023. FedALA: adaptive local aggregation for personalized federated learning//Proceedings of the 37th AAAI Conference on Artificial Intelligence. Washington, USA: AAAI: 11237-11244 [DOI: 10.1609/aaai.v37i9.26330http://dx.doi.org/10.1609/aaai.v37i9.26330]
Zhou C X, Sun Y, Wang D G and Ge H W. 2021. Survey of federated learning research. Chinese Journal of Network and Information Security, 7(5): 77-92
周传鑫, 孙奕, 汪德刚, 葛桦玮. 2021. 联邦学习研究综述. 网络与信息安全学报, 7(5): 77-92[DOI: 10.11959/j.issn.2096-109x.2021056http://dx.doi.org/10.11959/j.issn.2096-109x.2021056]
相关作者
相关机构