Multi-stage Attention ResU-Net

• Multi-stage Attention ResU-Net is a family of convolutional architectures that integrates hierarchical attention and residual connections within a multi-stage U-Net for precise segmentation.
• It leverages cascaded U-Net stages with additive gating for medical images and linear attention mechanisms for remote sensing, enhancing both global and local feature extraction.
• Empirical results demonstrate superior segmentation accuracy and efficiency, with improvements validated by metrics such as Dice score and mIoU over traditional U-Net models.

Multi-stage Attention ResU-Net (MAResU-Net) refers to a family of convolutional architectures that integrate attention mechanisms and @@@@1@@@@ in a multi-stage U-Net topology for image segmentation tasks. Two instantiations with substantial methodological differences have been proposed: a variant for semantic segmentation of fine-resolution remote sensing images utilizing a Linear Attention Mechanism (LAM) (Li et al., 2020), and a double U-Net cascade with attention-gated residual blocks for medical image segmentation (Khan et al., 2023). Both approaches leverage hierarchical attention, residual connectivity, and multi-stage refinement to improve semantic segmentation, but diverge in attention mechanism, backbone, and training configurations.

1. Architectural Principles

MAResU-Net designs are characterized by the systematic stacking of attention-augmented residual U-Nets or attention-residual blocks at multiple semantic levels:

• Cascaded Stages: In medical segmentation (Khan et al., 2023), two full U-Net structures (Stage 1 and Stage 2) operate in series. The first U-Net predicts an initial mask, which is used to mask the input image spatially. The masked image feeds the second U-Net, allowing progressive focusing and refinement.
• Hierarchical Multi-stage Attention: In remote sensing segmentation (Li et al., 2020), a single encoder-decoder U-Net is augmented with attention-residual blocks inserted at each skip connection between encoder and decoder. Attention is thus exerted at all five spatial scales of the representation, from coarse/global to fine/local.
• Residual Pathways: Both designs use identity-style residual connections within convolutional blocks (and/or blocks spanning the attention modules), ensuring efficient gradient propagation and stable deep model optimization.

This architectural motif enables the capture of both global context and fine-grained boundary information, while attenuating irrelevant background signals.

2. Attention Mechanisms

Distinct attention formulations are proposed.

2.1. Additive Gating (Medical Imaging)

Attention gates are positioned in all skip connections. For encoder feature $x^{\ell}$ and decoded gating signal $g^{\ell}$ at level $\ell$, the gate computes:

$$\alpha^{\ell} = \sigma\left(W_x^{\ell} * x^{\ell} + W_g^{\ell} * g^{\ell} + b^{\ell}\right)$$

where $W_x^{\ell}, W_g^{\ell}$ are $1 \times 1$ convolutional weights and $\sigma$ is the sigmoid. The gated output is $\hat{x}^{\ell} = \alpha^{\ell} \odot x^{\ell}$, suppressing spatially irrelevant activations and enhancing decoder selectivity (Khan et al., 2023).

2.2. Linear Attention Mechanism (LAM, Remote Sensing)

To overcome the quadratic complexity of standard dot-product attention, LAM linearly approximates the attention kernel. For input feature map $X \in \mathbb{R}^{N \times C}$, with projections $Q, K, V$, LAM approximates

$$\exp(q_i^\top k_j) \approx 1 + \frac{q_i}{\|q_i\|_2}^\top \frac{k_j}{\|k_j\|_2}$$

Enabling global attention with $O(N)$ time and memory complexity, the output for each position is:

$$D_i = \frac{U + \bar{q}_i^\top R}{N + \bar{q}_i^\top M}$$

where $U = \sum_{j=1}^N v_j$, $M = \sum_{j=1}^N \frac{k_j}{\|k_j\|_2}$, $R = \sum_{j=1}^N \frac{k_j}{\|k_j\|_2} v_j^T$. LAM is inserted in a block that also includes channel-wise self-attention and a convolutional fusion (Li et al., 2020).

3. Network Topologies and Dataflows

3.1. Medical Segmentation (AttResDU-Net)

The architecture comprises:

• Stage 1: VGG-19 encoder (5 levels, feature sizes from $H \times W \times 64$ to $H/32 \times W/32 \times 512$), ASPP bottleneck, decoder with attention-gated skip-conn, all residual blocks.
• Stage 2: Encoder-decoder of identical structure, not initialized with pre-trained weights. Input is original image masked by $S_1$, the output of Stage 1.
• Final Fusion: Concatenate Stage 1 and 2 masks; process via $3\times3$ conv (64 filters), then $1\times1$ conv (sigmoid) to obtain the final prediction.
• ASPP: Parallel dilated $3 \times 3$ convolutions with rates 6, 12, 18 applied at the network bottleneck.
• Skip Connections: All five encoder-decoder levels use attention gates and residual blocks to process skip features.

3.2. Remote Sensing Segmentation

• Encoder: ResNet-34 backbone; feature stages C1–C5 (channels: 64, 64, 128, 256, 512).
• Decoder: Five-stage upsampler (D5→D1), halving channels at each step.
• Residual Attention Blocks: At each encoder–decoder connection, mixing encoder and decoder features via channel-wise self-attention, LAM, $1\times1$ projections, and residual addition.
• Output: Per-pixel semantic label predictions.

4. Training Procedures and Practical Configurations

4.1. Medical Segmentation (Khan et al., 2023)

• Loss: Dice loss, $$L_{dice} = 1 - \frac{2 \sum y_i \hat{y}_i + \epsilon}{\sum y_i + \sum \hat{y}_i + \epsilon}$$.
• Optimizer: Nadam; learning rate $10^{-4}$ (CVC-ClinicDB, ISIC-2018), $10^{-5}$ (Data Science Bowl).
• Pre-processing: Sample-wise centering, normalization, Shades-of-Gray color constancy.
• Augmentation: Random rotation, flips, HSV jitter, brightness/contrast, histogram equalization (implemented via Albumentations).
• Batch Size: Fitted to Tesla V100 (typically 8–16).
• Schedule: Reduce-on-plateau, early stopping after ~10 epochs without improvement; total up to 40 epochs.

4.2. Remote Sensing Segmentation (Li et al., 2020)

• Loss: Pixel-wise soft cross-entropy.
• Optimizer: AdamW, fixed learning rate $3 \times 10^{-4}$; batch size 1–4 tiles per GPU.
• Augmentation: Random flip and crop to $512 \times 512$.
• Input: Near-infrared, red, green bands.

5. Empirical Results and Ablative Findings

Medical Image Segmentation (Khan et al., 2023)

Dataset	Dice (%)	IoU (%)	Recall (%)	Precision (%)
CVC-ClinicDB	94.35	89.32	87.50	97.37
ISIC-2018	91.68	84.68	87.55	94.19
Data Science Bowl	92.45	85.96	65.55	96.29

• Ablation: Full attention at all skip connections raises Dice by 0.74 percentage points over Half-Attention (Stage 1 only). Color constancy gives +0.91, residuals +0.04 but with faster convergence (~20 epochs vs. ~30).

Remote Sensing Segmentation (Li et al., 2020)

Network/Metric	Params (M)	FLOPs (G)	OA (%)	mIoU (%)
MAResU-Net (ResNet-34)	26.3	26.3	90.86	83.30
CE-Net	29.0	29.0	90.40	81.49
EaNet	44.3	44.3	89.99	80.22

• Statistical Tests: Paired Kappa $z$-tests confirm the improvements (exceed $z=1.96$, 95% confidence).
• Attention Placement: Adding more attention blocks at lower decoder levels (C1→D1) yields higher mIoU; total mIoU gain up to +2.1 over U-Net/ResNet-34 baseline.

6. Implementation Complexity and Performance

• Medical MAResU-Net: 36.5M parameters, 92.1 GFLOPs for $256 \times 256$ input, model size 146 MB.
• Remote Sensing MAResU-Net: 26.3M parameters, 26.3 GFLOPs per tile, epoch time 8.8s (ResNet-34 backbone).
• Efficiency: LAM allows spatial global attention at each level with linear (vs. quadratic) complexity. Empirically, computation and memory are considerably reduced compared to traditional attention modules.

7. Methodological Significance and Limitations

• Multi-stage Cascade: Stagewise refinement enables robust extraction of both gross object morphology and fine boundary details, by spatially focusing downstream computation via Stage 1 masking.
• Attention Gates and LAM: Both additive gating (Khan et al., 2023) and LAM (Li et al., 2020) facilitate selective spatial information flow, filtering encoder activations in accordance with decoder context. LAM extends this to large images with linear runtime.
• Residual Pathways: Residual blocks mitigate vanishing gradients and accelerate convergence, especially critical in deep cascaded or multi-stage topologies.
• Generalization and Open Problems: Validation in (Li et al., 2020) was restricted to one imaging source; visible generalization to multisensor, DSM data, or time series is an open question. The use of first-order kernel approximation in LAM may miss nonlinear interactions. Potential extensions include multi-head attention, positional encoding, and learned normalization schemes.

• AttResDU-Net: Medical Image Segmentation Using Attention-based Residual Double U-Net (Khan et al., 2023)
• Multi-stage Attention ResU-Net for Semantic Segmentation of Fine-Resolution Remote Sensing Images (Li et al., 2020)