CDNet: Causal Inference inspired Diversified Aggregation Convolution for Pathology Image Segmentation

Authors: dawei  fan
yifan  gao
jiaming  yu
changcai  yang
riqing  chen
lifang  wei
Venue: MICCAI_org_2024_Workshop_COMPAYL
Published: 2024-06-24 08:57:40

Abstract: Deep learning models have shown promising performance for Nuclei segmentation in the field of pathology image analysis. However, training a robust model from multiple domains remains a great challenge for Nuclei segmentation. Additionally, the shortcomings of background noise, highly overlapping between Nuclei, and blurred edges often lead to poor performance. To address these challenges, we propose a novel framework termed CDNet, which combines Causal Inference Module (CIM) with Diversified Aggregation Convolution (DAC) techniques. The DAC module is designed which incorporates diverse downsampling features through a simple, parameter-free attention module (SimAM), aiming to overcome the problems of edge blurring. Furthermore, we introduce CIM to leverage sample weighting by directly removing the spurious correlations between features for every input sample and concentrating more on the correlation between features and labels. Extensive experiments on the MoNuSeg and GLySAC datasets yielded promising results, with mean intersection over union (mIoU) and Dice similarity coefficient (DSC) scores increasing by 3.59% and 2.61%, and 2.71% and 2.04%, respectively, outperforming other state-of-the-art methods.

Main paper: 
 introduction: 
 1 Over the past decade, deep learning has reached promising results for Nuclei segmen-tation in the field of pathology image analysis (Ronneberger et al., 2015; Zhou et al., 2018; Oktay et al., 2018; Chen et al., 2021; Cao et al., 2022) . For instances, the U-Net pro-posed by (Ronneberger et al., 2015) and its improved version, U-Net++ by (Zhou et al., 2018) , with their unique encoder-decoder structure and skip connections, effectively integrate low-level to high-level semantic information of images, which greatly promotes the development of Nuclei segmentation technology. Moreover, the introduction of the Transformer (Vaswani et al., 2017) architecture, with its unique attention mechanism (Schlemper et al., 2019) , has optimized the processing of complex medical images, captured the intricate spatial relation-ships and feature hierarchies within images. Additionally, TransUNet (Chen et al., 2021) combines the image encoding capabilities of Convolutional Neural Network (CNN) with the deep contextual understanding of visual Transformers (Dosovitskiy et al., 2020) , bringing new perspectives and powerful potential to Nuclei segmentation technology. The application of classical medical image segmentation models to pathological sections, in particular those containing cell nuclei, is a challenging endeavour. In 2019, (Graham et al., 2019) introduced HoVer-Net, which employs horizontal and vertical distance maps to effectively separate overlapping cell nuclei. In 2022, (Wang et al., 2022) developed UCTransNet, an optimised U-Net architecture that improves segmentation by narrowing the semantic gap and exploiting multi-scale features. However, existing models encounter difficulties in the presence of background noise, overlapping nuclei, and blurred edges (Fig. 1 ), which negatively impact segmentation accuracy and clarity. These models are unable to distinguish between closely spaced nuclei and their surrounding tissue with sufficient accuracy. Furthermore, variations in lighting, resolution, and staining techniques across different hospitals or devices present challenges for the application of deep learning models in clinical practice.To tackle these challenges, we propose the use of CDNet, which combines the Causal Inference Module (CIM) with Diversified Aggregation Convolution (DAC). The DAC module employs a variety of downsampling techniques through SimAM (Yang et al., 2021) to mitigate edge blurring, thereby enhancing the precision and clarity of nuclei segmentation. The CIM addresses the issue of performance decline resulting from variations in data distribution across hospitals and devices. It does so by leveraging sample weighting to remove spurious correlations and focus on meaningful feature-label relationships, as proposed by (Zhang et al., 2021) . The following key contributions are made: (1) The DAC module was designed for feature fusion with the objective of optimising semantic information and reducing edge blurring issues. (2) The introduction of CIM with a causal inference strategy enables the dynamic adjustment of sample weights, thereby enhancing nuclei feature recognition and reducing spurious correlations. A comprehensive series of experiments has demonstrated that CDNet is a superior method to other state-of-the-art techniques, offering new insights for addressing domain shift in pathology image analysis and improving segmentation accuracy. 
method: 
 In this section, we first introduce an overview of our CDNet. Then, we detail each part of the CDNet.A graphical representation of the CDNet overview is provided in Fig. 2 . The framework commences with an input micrograph image, which is subjected to five layers of CNN and multiscale block convolutional neural network (MBConv) (Tan and Le, 2019) downsampling, thereby creating multilevel feature maps. In the subsequent stage, the feature maps generated by each downsampling stage are combined using a dynamic attention mechanism (DAC). The Transformer enhances the initial CNN layer's feature maps, thereby producing a more integrated set. CIM learns the sample weights and computes the weighted loss for the final feature maps. The aforementioned maps are then fused using skip connections and a convolutional decoder, in a manner analogous to that observed in U-Net. Ultimately, the softmax function is responsible for generating the nuclei segmentation map.To enhance the precision of Nuclei segmentation, particularly in addressing segmentation challenges arising from issues such as blurred edges and background noise, we introduce an effective Diversified Aggregation Convolution (DAC) module. The DAC is specifically designed to integrate multi-scale feature maps (Cui et al., 2016) with rich semantic information. It optimizes the feature extraction process using CNN and MBConv and focuses on regions of interest by a simple, parameter-free attention module (SimAM), achieving precise discrimination between nucleus and non-nucleus.As demonstrated in Fig. 2 (b), the input data undergoes a downsampling f i 1 and f i 2 after using CNN and MBConv, where f i 1 represents the i-th layer downsampling using CNN and f i 2 represent the corresponding layer using MBConv. Afterwards, we concatenate f i 1 and f i 2 by learned weights k i 1 and k i 2 . The obtained features are then fed into SimAM, which aims to accurately evaluate and determine the most important features of semantic information in the Nuclei segmentation. Subsequently, these features from SimAM are downsampled using two-dimensional convolution (Conv2D). Hence, the DAC can be represented as follows:f i = Conv2D(SimAM([k i 1 f i 1 , k i 2 f i 2 ])) (1)Where [•] denotes the CONCAT operation and f i denotes i-th layer feature maps after DAC. We initialize k i 1 and k i 2 as 1.0.To address data heterogeneity problem, we follow the StableNet (Zhang et al., 2021) framework and design a causal inference module (CIM). We obtain the feature maps f 5 from the 5-th layer DAC as shown in Fig. 2 . To remove spurious correlations between features, we leverage Random Fourier Features (RFF) extractor and sample weighting for segmentation task. Specifically, we use A and B represent the feature variable in the feature map from f 5 . As Frobenius norm of the partial cross-covariance matrix ∥ AB ∥ 2 F tends to zero, the two variables A and B are independent (Zhang et al., 2021) . Hence, the partial cross-covariance matrix be:EQUATIONWhere u(•) and v(•) represent the RFF mapping functions, n represents the number of input images. We use w ∈ R n + to represent the sample weights and n i=1 w i = n in the Weight Learner as shown in Fig. 3 . After weighting, the partial cross-covariance matrix for variables A and B is as follows:Σ AB;w = 1 n -1 n i=1      w i u(A i ) - 1 n n j=1 w j u(A j )   T •   w i v(B i ) - 1 n n j=1 w j v(B j )      (3) A i and B i (i ∈ [1, n]) are sampled from the feature distribution of A and B.The objective function of weight learner is: Where m represents the number of features and ∆ n = w ∈ R n + | n i=1 w i = n . Hence, weighting training samples with the optimal w * can mitigate the dependence between features and consequently remove spurious correlations between features to the greatest extent.w * = arg min w∈∆n 1≤i≤j≤m ∥Σ AB;w ∥ 2 F (4)Finally, the causal inference loss is as follows:EQUATIONWhere i represents the size of batch size. L CE (•, •)is the cross-entropy loss and ⊗ represents element-wise multiplication. f (•) represents the flattening operation. Furthermore, we leverage dice loss and focal loss (Lin et al., 2017) to optimize our framework simultaneously. The dice loss is a common loss in the segmentation tasks. And the focal loss aims to solve the problem of category imbalance by reducing the impact of easily classifiable samples and enhancing the challenging ones:L F L = -α t (1 -p t ) γ log(p t ) (6)Where α t is a balancing coefficient for Nuclei and background. p t represents the predicted probability for the correct class. γ is a tuning factor to adjust the weights of background.In this paper, we set α t as 0.8 and γ as 2.0. Hence, the total loss is:L = L CIM + λL Dice + (1 -λ)L F L (7)Where λ represents the coefficient of dice loss and (1 -λ) represents the coefficient of focal loss. We set λ as 0.5. 
dataset and experimental details: 
 The precise pixel-level labelling required for nuclei segmentation is a labour-intensive task.In order to evaluate the performance of the model, two smaller datasets were chosen: MoNuSeg (Kumar et al., 2017) and GLySAC (Doan et al., 2022) . MoNuSeg includes tissue images from various tumour patients, which have been annotated by multiple healthcare providers. This results in discrepancies due to differences in the organs, patients and staining protocols. GLySAC covers a range of nuclei, including lymphocyte and cancer epithelial nuclei, with varied sample characteristics. The data was divided into three sets: training, validation, and test. The ratio of the training set to the validation and test sets was 6:2:2. The experiments were conducted using PyTorch 2.2.1 on an NVIDIA RTX4090 GPU. The model was trained with the AdamW optimiser and a cosine annealing learning rate strategy (Loshchilov and Hutter, 2016) (initial rate 1e-3). The batch size was 8, and the training was conducted over 400 epochs. Early stopping was used to prevent overfitting. The images were resized to 224 × 224 and augmented with a variety of techniques, including horizontal flipping, rotation, Gaussian blur, colour intensity enhancement, and random cropping.In this work, we use two core metrics to evaluate model performance: the Dice Coefficient (DSC) and the Mean Intersection over Union (mIoU). DSC is suitable for binary image segmentation and measures sample similarity, while mIoU assesses the overlap between predicted and actual areas. These metrics together provide a precise evaluation of performance, verifying the model's effectiveness in complex image segmentation tasks.A comparison of the performance of seven state-of-the-art (SOTA) models was conducted.As demonstrated in Table 1 , our framework exhibits superior performance compared to Attention U-Net, achieving 3.59% higher mIoU and 2.61% higher DSC on the MoNuSeg dataset. It is noteworthy that the proposed approach outperforms transformer-based models such as TransUNet by 10.95% in mIoU and 8.16% in DSC. In comparison to UNet, our model demonstrates an improvement of 2.71% in mIoU and 2.04% in DSC on the GLySAC dataset. In comparison with TransUNet, the model demonstrates an increase of 11.46% in mIoU and 9.27% in DSC. These results demonstrate that CDNet attains the highest scores in the field of pathology image segmentation. Nuclei segmentation is a crucial aspect of pathology image analysis. To demonstrate the superiority of CDNet, we conducted a comparative analysis with some of the more advanced models. As illustrated in Fig. 4 , CDNet exhibited superior performance in terms of accuracy and efficiency, particularly in distinguishing closely neighbouring nuclei and identifying nucleus boundaries. Ablation studies were conducted to evaluate the efficacy of key components in CDNet. As demonstrated in Table 2 , the Backbone (comprising CNN downsampling and transformer) exhibited limited performance. The incorporation of DAC led to a 4.17% improvement in mIoU and a 3.02% improvement in DSC on MoNuSeg, as well as a 2.05% improvement in mIoU and a 1.58% improvement in DSC on GLySAC. This was achieved by enhancing edge recognition and feature fusion. The incorporation of CIM resulted in an improvement of 1.55% and 1.13% in mIoU and DSC, respectively, on MoNuSeg, and 0.91% and 0.54% on GLySAC, due to the elimination of spurious feature correlations. The combination of Backbone, DAC, and CIM led to an improvement in mIoU by 2.25% and DSC by 1.63% on MoNuSeg, and by 2.35% and 1.82% on GLySAC. These results demonstrate that CDNet enhances cell nucleus segmentation in pathology image analysis.In addition, we show the visualisation of the Grad-CAM (Selvaraju et al., 2017) feature activation before and after applying these two modules, as shown in Fig. 5 . DAC effec- tively reduces the model's focus on background regions compared to the case without these modules, whereas CIM results in migrating the model's focus from irrelevant features to the nucleus. 
conclusion: 
 In this paper, focusing on the challenges of domain shifts and blurred edges in the Nuclei segmentation tasks, we propose a novel CDNet that combines CIM with DAC. Compared to other methods, our CDNet can effectively improve the results of Nuclei segmentation by removing spurious correlations between features, and focusing more on the correlations between features and labels. Meanwhile, it improves the accuracy of Nuclei edges recognition by optimizing various types of semantic information for feature fusion. The extensive experiments demonstrate that our CDNet can effectively alleviate the issues in the Nuclei segmentation and achieves a robust performance. 
 
 Appendix:  



Reviews: review: title: title: Interesting method yet questionable application
main_review: main_review: Summary of paper: 

The authors develop a novel deep learning model by integrating two previously developed modules into a mixed convolutional and transformer model. These diversified aggregation convolution and causal inference module are then shown to improve dice and mean IoU for binary segmentation in two datasets. Included ablation study and some gradcam examples further demonstrate the impact of the introduced modules.

Strengths:

- The proposed method integrates two interesting modules and takes a different approach to the nuclei segmentation task.
- The manuscript is well written with included examples, nice figures and an ablation study to support their claims.

Weaknesses:

Major:

- The proposed method is applied to a binary segmentation task instead of an instance segmentation task which misses the point of nuclei segmentation completely. The authors should include an instance segmentation evaluation and comparison or address how this binary segmentation output is supposed to be used for further downstream tasks in histopathology.
- The authors compare their method against e.g. HoVer-Net which also does instance segmentation and classification which is an unfair comparison. They also do not explain where the results for these models come from and perhaps the *-symbol is supposed to address that, but I cannot find an explanation for the star.
- The authors propose a combination of new methods and include an ablation study. Yet they do this with a single split. Please include cross-validation experiments to further support the claims.
- Ground-truth annotations are mostly done manually by human pathologists. These annotations are by design slightly random, especially in blurry borders. The method is supposed to accurately match the human segmentation, yet the human segmentation might also be wrong. Can the authors include reasons as to why a more accurate nuclei boundary segmentation would be beneficial?

Minor:

- The images without overlay in Figure 5 have the wrong color which usually happens if you load images with opencv and forget to change channels from BGR to RGB.

Other comments:

- It would be nice to know how computationally expensive this approach is compared to the other included methods. How long does inference for a single tile take in comparison?

Code, model weights, and/or data availability:
- not available

Conclusion:

The authors propose an interesting segmentation method which includes two previously developed building blocks for deep learning models and evaluate the method on a histopathology task. Yet, the authors should include valid comparisons against models that solve the same task and include a thorough evaluation using cross-validation. Moreover, and more importantly, the authors need to explain what kind of problem they are trying to solve with this approach.
title: title: Application of StableNet feature decorrelation plus attention equals better performance
main_review: main_review: The application of StableNet's learned weights for decorrelation in nucleus segmentation makes a lot of sense. It's good to see the application of a method that directly addresses the problem where the model learns to correlate various types of background noise. The ablation does appear to show that both DAC and CIM are necessary to achieve such performance.

The question of whether DAC is a generally better way of integrating multiscale feature maps or not remains to be seen. Due to the lack of simpler baselines replacing DAC with other feature aggregators, it is difficult to say whether DAC is a superior feature aggregator.

Overall the paper proposed a strong method that makes sense. It demonstrated good performance.
title: title: Review for Dawei Fan
main_review: main_review: Quality:
The quality of the work is overall very good.
Pros:
The detailed comparative analysis with SOTA models underscores the robustness of the proposed method and more important the model outperforms the compared SOTA models.

Clarity: The clarity of the work is overall okay. 
Pros:
Clear and logical structure and organization of the work.
Cons:
The use of technical terminology, mainly in the methodology sections, may pose a challenge for readers not well-versed in the field.

Originality:
The work presents original contributions through the innovative combination of CIM and DAC incorporated in the CDNet framework, improving segmentation accuracy.

Significance:
The work has significant potential to improve pathology image segmentation