[Paper Review] Fuzzy Attention Neural Network to Tackle Discontinuity in Airway Segmentation

Nan Yang, Javier Del Ser+8 more|arXiv (Cornell University)|Sep 5, 2022

COVID-19 diagnosis using AI4 citations

TL;DR

This paper proposes a Fuzzy Attention Neural Network (FANN) that integrates fuzzy logic with attention mechanisms to improve 3D airway segmentation in CT scans, particularly for small, discontinuous airway branches. By introducing a learnable Gaussian membership function and a novel Jaccard Continuity and Accumulation Mapping (JCAM) loss, FANN achieves superior continuity and completeness, outperforming state-of-the-art models on datasets with lung cancer, COVID-19, and pulmonary fibrosis.

ABSTRACT

Airway segmentation is crucial for the examination, diagnosis, and prognosis of lung diseases, while its manual delineation is unduly burdensome. To alleviate this time-consuming and potentially subjective manual procedure, researchers have proposed methods to automatically segment airways from computerized tomography (CT) images. However, some small-sized airway branches (e.g., bronchus and terminal bronchioles) significantly aggravate the difficulty of automatic segmentation by machine learning models. In particular, the variance of voxel values and the severe data imbalance in airway branches make the computational module prone to discontinuous and false-negative predictions. especially for cohorts with different lung diseases. Attention mechanism has shown the capacity to segment complex structures, while fuzzy logic can reduce the uncertainty in feature representations. Therefore, the integration of deep attention networks and fuzzy theory, given by the fuzzy attention layer, should be an escalated solution for better generalization and robustness. This paper presents an efficient method for airway segmentation, comprising a novel fuzzy attention neural network and a comprehensive loss function to enhance the spatial continuity of airway segmentation. The deep fuzzy set is formulated by a set of voxels in the feature map and a learnable Gaussian membership function. Different from the existing attention mechanism, the proposed channel-specific fuzzy attention addresses the issue of heterogeneous features in different channels. Furthermore, a novel evaluation metric is proposed to assess both the continuity and completeness of airway structures. The efficiency, generalization and robustness of the proposed method have been proved by training on normal lung disease while testing on datasets of lung cancer, COVID-19 and pulmonary fibrosis.

Motivation & Objective

• Address the challenge of discontinuous and incomplete segmentation of small airway branches (e.g., terminal bronchioles) in 3D CT images due to data imbalance and intensity homogeneity.
• Overcome the limitations of conventional evaluation metrics like IoU and Dice, which fail to capture segmentation continuity and completeness.
• Improve model generalization and robustness across diverse lung disease phenotypes, including fibrosis, COVID-19, and cancer.
• Develop a novel fuzzy attention mechanism that models feature uncertainty and enhances representation learning in heterogeneous feature channels.
• Introduce a new evaluation metric, CCF-score, to better guide model selection by assessing both continuity and completeness of airway structures.

Proposed method

• Propose a deep fuzzy set formulation using a set of voxels in the feature map and a learnable Gaussian membership function to model uncertainty in feature representations.
• Design a channel-specific fuzzy attention layer that adaptively weights features across different channels, improving robustness to heterogeneous feature distributions.
• Introduce a novel Jaccard Continuity and Accumulation Mapping (JCAM) loss function to optimize for spatial continuity and structural completeness of segmented airways.
• Develop a new evaluation metric, CCF-score (Continuity and Completeness Score), to assess segmentation quality beyond standard overlap metrics.
• Train the FANN model using a comprehensive loss combining cross-entropy, Dice, and the proposed JCAM loss to balance foreground-background imbalance and promote structural fidelity.
• Apply data augmentation and class-balanced sampling strategies to mitigate data imbalance in airway segmentation tasks.

Experimental results

Research questions

• RQ1 Can a fuzzy attention mechanism improve the robustness and generalization of deep neural networks in 3D airway segmentation, especially for small and discontinuous airway branches?
• RQ2 Does the integration of fuzzy logic with attention mechanisms lead to better feature representation learning under high uncertainty and data imbalance?
• RQ3 Can the proposed JCAM loss function effectively enhance the spatial continuity of airway segmentation compared to standard loss functions?
• RQ4 Does the novel CCF-score metric outperform traditional metrics like IoU and Dice in identifying models with superior segmentation continuity and completeness?
• RQ5 How does the FANN model generalize across diverse lung disease phenotypes, including fibrosis, COVID-19, and lung cancer, where airway structures are highly altered?

Key findings

• FANN achieved the highest CCF-score (0.9215) on the COVID-19 dataset and 0.7687 on fibrosis data, significantly outperforming baseline models in assessing segmentation continuity and completeness.
• On the fibrosis dataset, FANN achieved a terminal bronchiole (TB) detection rate of 47.15%, which, while lower than on other datasets, indicates room for improvement in complex fibrotic lung structures with honeycombing.
• The model achieved 99.38% detection rate for large branches (LB) on the BAS dataset and 100% on the COVID-19 dataset, demonstrating strong performance on larger airways.
• FANN reduced the average missing rate (AMR) to 2.36% on the COVID-19 dataset and 10.29% on fibrosis, indicating fewer leakages and better completeness than competing models.
• The ablation study confirmed that the combination of fuzzy attention and JCAM loss yielded the best performance, with BL+fuzzy+JCAM achieving a 93.81% DLR and 91.87% DBR on the COVID-19 dataset, outperforming other configurations.
• The proposed CCF-score effectively guided model selection, as the model with the highest CCF-score also showed the most consistent and clinically plausible segmentation across all datasets.