Stroke Lesion Segmentation using Multi-Stage Cross-Scale Attention

TL;DR

The paper tackles the problem of accurately segmenting stroke lesions in MRI, which is essential for prognosis and understanding lesion–outcome relationships. It introduces Multi-Stage Cross-Scale Attention (MSCSA), a plug-in module for the U-Net that enables multi-stage and cross-scale feature interactions through Cross-Scale Attention with MSP and Relational Positional Encoding, along with an Intra-FFN to manage computational costs. On ATLAS v2.0, MSCSA-based approaches outperform baselines, with the ensemble MSCSA achieving the highest Dice and F1 scores on both the full dataset and a small-lesion subset, demonstrating robustness across training schemes. The work suggests promising directions such as size-aware labeling strategies and extending MSCSA to other lesion domains, highlighting its potential clinical utility in MRI-based lesion characterization.

Abstract

Precise characterization of stroke lesions from MRI data has immense value in prognosticating clinical and cognitive outcomes following a stroke. Manual stroke lesion segmentation is time-consuming and requires the expertise of neurologists and neuroradiologists. Often, lesions are grossly characterized for their location and overall extent using bounding boxes without specific delineation of their boundaries. While such characterization provides some clinical value, to develop a precise mechanistic understanding of the impact of lesions on post-stroke vascular contributions to cognitive impairments and dementia (VCID), the stroke lesions need to be fully segmented with accurate boundaries. This work introduces the Multi-Stage Cross-Scale Attention (MSCSA) mechanism, applied to the U-Net family, to improve the mapping between brain structural features and lesions of varying sizes. Using the Anatomical Tracings of Lesions After Stroke (ATLAS) v2.0 dataset, MSCSA outperforms all baseline methods in both Dice and F1 scores on a subset focusing on small lesions, while maintaining competitive performance across the entire dataset. Notably, the ensemble strategy incorporating MSCSA achieves the highest scores for Dice and F1 on both the full dataset and the small lesion subset. These results demonstrate the effectiveness of MSCSA in segmenting small lesions and highlight its robustness across different training schemes for large stroke lesions. Our code is available at: https://github.com/nadluru/StrokeLesSeg.

Figures (3)

• Figure 1: (a) Multi-Stage Cross-Scale Attention (MSCSA) as an add-on module. Feature maps from each Convolution (Conv) Block in the encoder are down(up)sampled to a uniform scale, concatenated along the channel dimension, and directed into the MSCSA block. The resulting output feature maps after a Batch Normalization (BN) layer undergo a reversal process, reverting to their original shapes. These refined maps are subsequently fused with feature maps sourced from the encoder and transmitted to the decoder. (b) Cross-Scale Attention (CSA) and Intra-Feed-Forward Network (Intra-FFN). In addition to the standard attention mechanism, CSA introduces the Multi-Scale key and value Projection (MSP) and Relational Positional Encoding (RPE). The MSP manipulates the key and value tensors across three scales, culminating in the generation of a cross-scale attention map. The RPE, which includes an activation layer and a Depth-Wise Convolution (DEConv) layer, enhances local feature aggregation through positional encoding. For Intra-FFN, feature maps are segmented into distinct parts corresponding to the number of channels in each encoder stage. Following individual processing through the standard Feed-Forward Network (FFN), these segmented maps are concatenated, reinstating their multi-stage format.
• Figure 2: Visualization of lesion segmentation results. All the methods work well on larger lesion areas (see lesions in the right hemisphere of the brain), while MSCSA generalizes better for small lesions (see lesions in the left hemisphere).
• Figure 3: Dice score vs. stroke lesion volume. We can observe that MSCSA (blue) surpasses the baseline (red) for small lesions while simultaneously being competitive with the baseline on medium and large lesions.