2405.01503v1

# PAM-UNet: Shifting Attention on Region of Interest in Medical Images > *This work is supported by the NIH funding: R01-CA246704, R01-CA240639, U01-DK127384-02S1, and U01-CA268808.1A. Das, D. Jha, V. Gorade, K. Biswas, H. Pan, Z. Zhang and U. Bagci are with the Machine & Hybrid Intelligence Lab, Department of Radiology, Northwestern University, USA.2D.P. Ladner is with Comprehensive Transplant Center, Feinberg School of Medicine, Northwestern University, USA & Division of Transplantation, Department of Surgery, Northwestern Medicine, USA. Abstract Computer-aided segmentation methods can assist medical personnel in improving diagnostic outcomes. While recent advancements like UNet and its variants have shown promise, they face a critical challenge: balancing accuracy with computational efficiency. Shallow encoder architectures in UNets often struggle to capture crucial spatial features, leading in inaccurate and sparse segmentation. To address this limitation, we propose a novel P rogressive A ttention based M obile UNet (PAM-UNet) architecture. The inverted residual (IR) blocks in PAM-UNet help maintain a lightweight framework, while layerwise Progressive Luong Attention ( $\mathcal{PLA}$ ) promotes precise segmentation by directing attention toward regions of interest during synthesis. Our approach prioritizes both accuracy and speed, achieving a commendable balance with a mean IoU of 74.65 and a dice score of 82.87, while requiring only 1.32 floating-point operations per second (FLOPS) on the Liver Tumor Segmentation Benchmark (LiTS) 2017 dataset. These results highlight the importance of developing efficient segmentation models to accelerate the adoption of AI in clinical practice. I Introduction The growing demand for precise segmentation in medical images has led to the development of complex models like HRNet [18], DeepLabv3+ [17], AttUNet [15], and FANet [13] which rarely fit on edge devices. While adopting shallower backbones in UNet architectures, such as Efficient-UNet [1] and Mobile-UNet [10], reduces computational requirements, it frequently comes at the cost of segmentation quality. These models often struggle to differentiate indistinguishable peripheral and occluded patterns, and handle the large inter-patient variation of anomalies, due to the inability of lightweight encoding modules to generate significantly rich feature maps. Fully convolutional CNNs (FCNNs) [4] and UNet [12] have paved the way in biomedical image segmentation. UNet inspired networks, like DoubleUNet [8], ResNet-UNet [6], Attention-UNet [15] excel in organ segmentation and lesion detection. However, extensive down-sampling using deeper backbones results in spatial information loss. Mobile-UNet aims to mitigate overfitting and optimize computation by employing a pre-trained MobileNetV2 [7] as the encoder, but the limited capacity for capturing fine-grained features results in spatial context loss. This motivated us to propose PAM-UNet, a new segmentation architecture based on depth-wise separable convolution and a novel Progressive Luong Attention ( $\mathcal{PLA}$ ) mechanism. $\mathcal{PLA}$ selectively incorporates long-range dependencies from shallow features to the final segmentation mask, guided by residual cues from the encoder. This targeted approach, unlike conventional attention mechanisms, focuses only on relevant regions. This synergistic design, combined with Mobile convolutions (MB Conv) and layerwise Luong Loss, significantly reduces FLOPs while enhances feature representation. Our key contributions are summarized as follows: - Novel architecture: We introduce a novel architecture, PAM-UNet, that combines mobile convolution blocks along with $\mathcal{PLA}$ mechanism for precise medical image segmentation with heightened efficiency. - Advanced attention mechanism: Our study show superior performance of $\mathcal{PLA}$ attention over self-and cross-attention mechanisms. - Detailed interpretation: We provide a detailed interpretation to show how PAM-UNet is different from and superior to existing baselines using center kernel alignment (CKA) [16]. - Extensive evaluation: We have throughly evaluated PAM-UNet on two common benchmarks: LiTS-2017 [2] and Kvasir-SEG [9] dataset. The results show that PAM-UNet achieves comparable precision with $30x$ less computational overhead. <details> <summary>2405.01503v1/extracted/5573681/embs-cka-comp.png Details</summary>  ### Visual Description # Technical Document Extraction: Heatmap Analysis of UNet Variants ## Image Description The image presents four comparative heatmaps labeled **(a) PAM-UNet**, **(b) UNet**, **(c) AttUNet**, and **(d) MobileUNet**. Each heatmap visualizes a matrix of values using a color gradient from **purple (low values)** to **yellow (high values)**. No explicit legend or colorbar is provided to quantify the values. --- ### **Key Labels and Axis Titles** 1. **Y-Axis**: Labeled "PAM-UNet" across all heatmaps. 2. **X-Axis**: Unlabeled in all heatmaps. 3. **Heatmap Labels**: - (a) PAM-UNet - (b) UNet - (c) AttUNet - (d) MobileUNet --- ### **Heatmap Analysis** #### **(a) PAM-UNet** - **Pattern**: Strong diagonal dominance with high values (yellow) along the diagonal from top-left to bottom-right. - **Trend**: Values decrease symmetrically away from the diagonal, forming a gradient. - **Spatial Grounding**: Highest values concentrated along the diagonal axis. #### **(b) UNet** - **Pattern**: Uniform distribution with no clear diagonal trend. - **Trend**: Values are relatively consistent across the matrix, with occasional localized peaks (yellow regions). - **Spatial Grounding**: No dominant axis; values are scattered. #### **(c) AttUNet** - **Pattern**: Similar to UNet but with slightly more pronounced vertical/horizontal striations. - **Trend**: Values cluster in horizontal bands, suggesting feature-group dependencies. - **Spatial Grounding**: Vertical/horizontal groupings dominate. #### **(d) MobileUNet** - **Pattern**: Diagonal trend similar to PAM-UNet but with reduced intensity. - **Trend**: Diagonal values are high (yellow), but off-diagonal regions show more variability than PAM-UNet. - **Spatial Grounding**: Diagonal remains dominant but less pronounced than in PAM-UNet. --- ### **Observations** 1. **Diagonal Dominance**: PAM-UNet and MobileUNet exhibit diagonal patterns, likely indicating self-similarity or identity preservation in feature mappings. 2. **Uniformity**: UNet and AttUNet lack strong diagonal trends, suggesting more distributed feature interactions. 3. **Color Consistency**: All heatmaps use the same purple-to-yellow gradient, but no numerical scale is provided to quantify differences. --- ### **Limitations** - No explicit legend or colorbar to interpret absolute value ranges. - Unlabeled x-axis prevents identification of specific categories or sub-categories. - No textual data or tables present in the image. --- ### **Conclusion** The heatmaps compare feature interaction patterns across four UNet variants. PAM-UNet and MobileUNet show stronger diagonal correlations, while UNet and AttUNet exhibit more uniform or grouped distributions. Further analysis with quantitative metrics (e.g., correlation matrices) is required for precise comparisons. </details> Figure 1: CKA plot illustrating that PAM-UNet learns from the very beginning. (a) PAM-UNet have highly similar structure throughout the model, (b) and (d) A large number of lower layers in the UNet and MobileUNet are similar to smaller set of the lower PAM-UNet layers and (c) PAM-Net takes a few layers lesser than AttUNet to learn the similar representations during training. II Preliminaries Vanilla vs Depth-wise separable convolution [5]: Traditional convolution performs a linear combination of input values and kernels across all dimensions and channels. In contrast, depth-wise separable convolution decomposes the convolution operations into depth-wise and point-wise convolutions separately and performs: <details> <summary>2405.01503v1/extracted/5573681/EMBC-Model.jpg Details</summary>  ### Visual Description # Technical Document Extraction: Neural Network Architecture and Attention Mechanism ## Section a: Encoder-Decoder Architecture with Attention ### Diagram Overview The diagram illustrates a U-Net-like architecture with attention mechanisms for medical image segmentation. Key components include convolutional layers, attention gates, residual connections, and an IR (Intensity Ratio) block. ### Component Breakdown 1. **Input** - MRI scan: `H x W x 3` (Height x Width x 3 channels) - Processed through Conv + IR block 2. **Encoder Path** - **Layer 1**: `H/2 x W/2 x 16` → Conv + IR → `H/2 x W/2 x 16` - **Layer 2**: `H/4 x W/4 x 24` → Conv + IR → `H/4 x W/4 x 24` - **Layer 3**: `H/8 x W/8 x 32` → Conv + IR → `H/8 x W/8 x 32` - **Layer 4**: `H/16 x W/16 x 96` → Conv + IR → `H/16 x W/16 x 96` - **Bottleneck**: `H/32 x W/32 x 1028` → Conv + IR → `H/32 x W/32 x 1028` 3. **Decoder Path** - **Upsampling**: - `H/16 x W/16 x 96` → Conv Transpose → `H/8 x W/8 x 32` - `H/8 x W/8 x 32` → Conv Transpose → `H/4 x W/4 x 24` - `H/4 x W/4 x 24` → Conv Transpose → `H/2 x W/2 x 16` - **Attention Gates**: - Symbol: `⊙` (Concatenation) - Gating signal flow: `↑` (Upward arrows indicate signal propagation) 4. **Output** - Segmentation mask: `H x W x 3` - Final dimensions match input size for pixel-wise comparison ### Legend - `🔵 IR Block`: Intensity Ratio processing - `🟣 Residual`: Residual feature integration - `🟢 Attention Gate`: Feature modulation - `⚫ Concatenation`: Feature fusion - `↑ Conv Transpose`: Upsampling operation - `↑ Gating Signal`: Control signal for attention - `🟥 Residual Feature`: Skip connection features - `🟩 Attended Feature`: Attention-modulated features ## Section b: Cross-Attention Mechanism ### Component Breakdown 1. **Global Context Vectors** - Processed through Sigmoid activation - Dimensions: Not explicitly stated 2. **Attention Computation** - Query (Q) matrix: Residual from encoder - Key (K^T): Transposed key matrix - Value (V): Feature from lower decoder layer - Cross-attention operation: `Q ⊗ K<sup>T</sup> ⊗ V` 3. **Output** - Attention gated feature map: `H x W x 3` - Visualized as heatmap with red/yellow intensity gradients ## Section c: Feature Map Visualization ### Layer Analysis | Layer | Before Processing | After Processing | Dimensions | |-------|-------------------|------------------|------------------| | L1 | Low contrast | High contrast | `H/16 x W/16 x 96`| | L2 | Blurred edges | Sharp edges | `H/8 x W/8 x 32` | | L3 | Faint structures | Clear structures | `H/4 x W/4 x 24` | | L4 | Noisy background | Clean segmentation| `H/2 x W/2 x 16` | ### Ground Truth vs Prediction - **GT (Ground Truth)**: White silhouette on black background - **Pred (Prediction)**: - L1: Noisy outline - L2: Partial segmentation - L3: Improved accuracy - L4: Near-perfect match ### Color Coding - Blue: Low activation - Red: High activation - Yellow: Intermediate activation ## Key Trends 1. **Encoder**: Progressive downsampling with increasing channel depth 2. **Decoder**: Upsampling with residual feature integration 3. **Attention**: Enhances relevant features while suppressing noise 4. **Cross-Attention**: Aligns global context with local features ## Spatial Grounding - Legend position: Bottom-right quadrant - Color verification: - IR Block (`🔵`): Blue arrows in diagram - Residual (`🟣`): Purple circles - Attention Gate (`🟢`): Green squares ## Trend Verification - **Encoder Path**: - Spatial resolution decreases by factor of 2 per layer - Channel depth increases by 8-16x per layer - **Decoder Path**: - Spatial resolution increases by factor of 2 per layer - Channel depth decreases by 2-4x per layer - **Attention Gates**: - Modulate feature maps to emphasize diagnostically relevant regions ## Missing Information - No explicit numerical data table present - No textual description of training parameters - No quantitative performance metrics (e.g., Dice score, IoU) ## Language Note All text in the diagram is in English. No foreign language content detected. </details> Figure 2: Here we present (a) Holistic architecture of proposed PAM-UNet, (b) Diagram of working of one unit of attention gate ( $\mathcal{PLA}$ ) and (c) Features extracted from intermediate layers of PAM-UNet showing the effectiveness of proposed $\mathcal{PLA}$ . Predicted mask before and after applying $\mathcal{PLA}$ is significantly different and proposed method produces precise segmentation. $$ \displaystyle DSConv=\mathcal{DP}(\mathbf{I},\mathbf{K}_{\mathrm{d}},\mathbf{K% }_{\mathrm{p}})= \displaystyle\underbrace{\sum_{m,n}\mathbf{I}(i+m,j+n,k)\ast\mathbf{K}_{% \mathrm{d}}(m,n,k)}_{\text{Depth-wise ($\mathcal{D}$)}}\,\underbrace{\sum_{k}% \ast\mathbf{K}_{\mathrm{p}}(k,l)}_{\text{Point-wise ($\mathcal{P}$)}} \tag{1} $$ where $\mathcal{DP}(·)$ performs depth-wise and point-wise convolution, $\mathbf{I}$ is the input tensor, $\mathbf{K}_{\mathrm{d}}$ and $\mathbf{K}_{\mathrm{p}}$ are depth-wise and point-wise kernels, respectively. $\sum_{m,n}$ works over all spatial dimensions $m$ and $n$ . Depth-wise separable convolution $DSConv$ reduces computational cost by applying separate kernels to individual input channels, making it more efficient than vanilla convolution. Self vs Cross [14] vs Luong attention [11]: Self-attention captures intra-sequence relationships via the weighted sum of input features by learnable parameters and activation ( $\sigma$ ). Cross-attention extends this to inter-sequence relationships. In contrast, Luong attention computes the context vector $c_{t}$ for each decoder time step $t$ as a weighted sum of encoder hidden states $h_{i}$ : $$ c_{t}=\sum_{i=1}^{T}\alpha_{t,i}\cdot h_{i} \tag{2} $$ The attention weights $\alpha_{t,i}$ are determined by a $\sigma$ function applied to a score function $score(s_{t},h_{i})$ , measuring the similarity between the decoder hidden state $s_{t}$ and encoder hidden states: $$ \alpha_{t,i}=\frac{\exp(score(s_{t},h_{i}))}{\sum_{j=1}^{T}\exp(score(s_{t},h_% {j}))} \tag{3} $$ III Proposed Method The proposed PAM-UNet adopts a U-shaped architecture with dedicated down-sampling and up-sampling arms. Each arm employs mobile convolution blocks equipped with layer-wise additive luong attention to capture relevant features at different scales. A detailed description is mentioned below: III-A Decoding the Lightweight U-shaped backbone The proposed PAM-UNet model, illustrated in Fig. 2 a, comprises an encoding arm and a decoding arm with skip connections between corresponding layers. The encoding arm incorporates MB Conv, IR bottleneck, and a convolutional layer for dimensionality reduction. The decoding arm mirrors the encoding arm’s structure, utilizing layerwise Luong attention on skip connections. Unlike UNet’s bottleneck block, MBConv employs an IR bottleneck and performs ConvTranspose instead of up-sampling during synthesis. The attention mechanism focuses on relevant information in the encoder’s residual features. The final layer by through a ConvTranspose generates the segmentation mask. III-B Progressive Luong Attention $\mathcal{PLA}$ is designed to capture hierarchical context dependencies. While the standard Luong attention attends to the information from a single layer, $\mathcal{PLA}$ refines this process by considering a progressive aggregation of information from multiple layers. As shown in Fig. 2 b, feature map from lower decoding layer is passed through a IR Block and ConvTranspose splits into Key and Value and the residual coming from the current encoding block is used as Query. Given query ( $Q$ ) and key ( $K$ ) vectors, the attention scores ( $A$ ) are calculated as: $$ A=\sigma\left(\frac{Q\cdot K^{T}}{\sqrt{\text{key\_dimension}}}\right) \tag{4} $$ The attention weights ( $W$ ) are obtained by applying $\sigma$ to the attention scores, and attended values ( $V$ ) are computed by: $$ W={\sigma}(A)\quad\text{and}\quad V_{\text{attended}}=W\cdot V \tag{5} $$ TABLE I: Results comparison of PAM-UNet with other popular benchmarking algorithms. Here, UNet variants with an enhanced attention mechanism are shown in orange and variants with improved backbone architecture is shown in cyan. | U-Net (ResNet50) [12] | 78.02 | 69.00 | 87.20 | 83.57 | 76.8 | 84.26 | 35.00 | | --- | --- | --- | --- | --- | --- | --- | --- | | DeepLabv3+ [3] | 82.34 | 73.59 | 91.11 | 81.72 | 78.36 | 85.24 | 27.90 | | Mobile-UNet [10] | 71.04 | 62.93 | 80.23 | 76.45 | 71.65 | 78.43 | 0.85 | | FCN8 [4] | 76.86 | 67.65 | 82.5 | 68.98 | 64.52 | 73.50 | 24.90 | | HRNet [18] | 76.21 | 72.85 | 92.34 | 80.46 | 77.87 | 87.43 | 11.70 | | ResUNet [6] | 74.09 | 64.45 | 78.98 | 65.21 | 59.41 | 70.45 | 14.82 | | AttUNet [15] | 82.38 | 74.32 | 76.95 | 83.24 | 79.91 | 87.16 | 35.35 | | PAM-UNet (Ours) | 82.87 | 74.65 | 92.14 | 84.8 | 78.40 | 86.63 | 1.32 | III-C Loss function During back-propagation, gradients from irrelevant regions are effectively down-weighted, prioritizing updates of relevant spatial regions. Assuming $\hat{Y}$ represents the predicted segmentation map and $Y$ represents the ground truth segmentation map, the segmentation loss ( $\mathcal{L}_{\text{seg}}$ ) is: $$ \mathcal{L}_{\text{seg}}=-\frac{1}{N}\sum_{i=1}^{N}\left(Y_{i}\cdot\log(\hat{Y% }_{i})+(1-Y_{i})\cdot\log(1-\hat{Y}_{i})\right) \tag{6} $$ The additional objective of attention mechanism is further enhanced by Attention Regularization to mitigate over-attentiveness that puts attention on both background and anomaly together. This is defined as: $$ \mathcal{R}(\mathbf{A})=\frac{1}{N}\sum_{i=1}^{N}(\mathbf{A}_{i}-\bar{A})^{2} \tag{7} $$ $$ \mathcal{L}_{\text{total}}=\mathcal{L}_{\text{seg}}+\lambda\cdot\mathcal{L}_{% \text{$\mathcal{R}$}} \tag{8} $$ This formulation aims to train the model by encouraging the attention mechanism to distribute its focus more evenly. IV Experiments Evaluation datasets: We used two publicly available datasets, namely Liver Tumor Segmentation (LiTS)-2017 [2] and Kvasir-SEG [9], each representing distinct imaging modalities (CT and Colonoscopy respectively) for comprehensive evaluation of performance and compatibility of our model. The LiTS dataset contains 201 CT images of the abdomen, but only the training dataset (130 scans) is made publicly available, so we have split the train set of LiTS into train and test. Kvasir-SEG comprises 1000 colonoscopy images. We used a split of 80% for training and the remaining for testing and validation. Metrics: Aligning with segmentation challenge evaluation standards, we adopt the Dice Similarity Coefficient (Dice), mean Intersection Over Union (mIoU) and Recall to measure the quality of predicted segmentation masks and to evaluate computation efficiency we have picked FLOPS. To further compare the representations learned by PAM-UNet and other models across different layers we use CKA plots. CKA goes beyond comparing individual activations and instead compares how data points relate to each other within the representation space. Implementation Details: Our models were implemented using PyTorch and trained on a NVIDIA RTX-A4000 GPU with 16GB of memory. The input image size was set to $128× 128$ , and training utilized a batch size of 8 with a learning rate of 0.01. We employed the SGD optimizer with momentum 0.9 and weight decay of 0.0001 for optimization. Data augmentations such as flipping and rotating were applied during training to enhance model generalization. Architecture Configuration: PAM-UNet employs MBConv blocks and $\mathcal{PLA}$ mechanism within a mirrored U-shaped architecture. It incorporates DSConv for flexible expansion factors and strides (1 to 2), along with optional deconvolution in IR blocks. To mitigate overfitting, attention regularization is applied. Key hyperparameters include an expansion factor of 6, stride variation from 1 to 2, and an attention regularization weight of 0.01. Techniques for Comparison: We compare PAM-UNet against UNet and six other variants: MobileUNet [10], DeepLabV3+ [3], FCN-8 [4], HRNet [18], ResUNet [6], and AttUNet [15]. The first four variants utilize diverse backbone implementations, while the remaining two incorporate attention gating methodologies. This separation aims to highlight the efficacy of a mobile convoluted backbone and the proposed $\mathcal{PLA}$ gating in PAM-UNet with all fairness. <details> <summary>2405.01503v1/x1.png Details</summary>  ### Visual Description # Technical Document: Medical Image Segmentation Model Comparison ## Image Structure The image is a comparative visualization of medical image segmentation results across three datasets: **LITS** (Liver Tumor Segmentation), **LITS** (repeated), and **Kvasir-SEG** (polyp segmentation). Each row represents a dataset, and each column represents a model's output or ground truth. --- ## Labels and Axis Titles - **Rows (Datasets):** 1. **LITS** (Liver Tumor Segmentation) 2. **LITS** (repeated, likely a different slice or case) 3. **Kvasir-SEG** (polyp segmentation in endoscopic images) - **Columns (Models/Datasets):** 1. **Image**: Original input medical images. 2. **GT (Ground Truth)**: Reference segmentation masks (gold standard). 3. **PAM-UNet**: Segmentation output from the PAM-UNet model. 4. **UNet**: Segmentation output from the UNet model. 5. **DeepLabv3+**: Segmentation output from the DeepLabv3+ model. 6. **HRNet**: Segmentation output from the HRNet model. 7. **FCN8**: Segmentation output from the FCN8 model. 8. **ResUNet**: Segmentation output from the ResUNet model. 9. **MobileUNet**: Segmentation output from the MobileUNet model. 10. **AttUNet**: Segmentation output from the AttUNet model. --- ## Key Observations 1. **Ground Truth (GT)**: - White regions represent the target anatomical structures (e.g., liver tumors, polyps). - Used as the reference for evaluating model performance. 2. **Model Performance**: - **PAM-UNet** (red label): Highlights segmentation results with varying accuracy across datasets. - **UNet**: Shows moderate segmentation quality, with some over/under-segmentation. - **DeepLabv3+**: Captures larger tumor regions but may miss smaller details. - **HRNet**: Demonstrates finer segmentation but with potential noise in Kvasir-SEG. - **FCN8**: Struggles with small structures (e.g., polyps in Kvasir-SEG). - **ResUNet**: Balanced performance, closely matching GT in LITS. - **MobileUNet**: Lightweight model with acceptable results but lower precision. - **AttUNet**: Attention-based model shows improved focus on small structures (e.g., polyps). 3. **Dataset-Specific Notes**: - **LITS**: Focuses on liver tumor segmentation. Models like ResUNet and PAM-UNet perform well. - **Kvasir-SEG**: Polyp segmentation is challenging due to small size and variability. AttUNet and HRNet show better adaptation. --- ## Color Legend - **Green**: Ground Truth (GT). - **Red**: PAM-UNet segmentation results. - **Black/White**: Segmentation outputs from other models (binary masks). --- ## Transcribed Text Embedded in Diagrams - Column headers: `Image`, `GT`, `PAM-UNet`, `UNet`, `DeepLabv3+`, `HRNet`, `FCN8`, `ResUNet`, `MobileUNet`, `AttUNet`. - Row headers: `LITS`, `LITS`, `Kvasir-SEG`. --- ## Notes - All segmentation outputs are binary masks (black background, white segmented regions). - Model performance varies by dataset complexity (e.g., Kvasir-SEG's small polyps vs. LITS' larger tumors). - No numerical metrics (e.g., Dice score) are provided; qualitative comparison is implied. </details> Figure 3: Segmentation masks generated by PAM-UNet (Ours) and all the baselines highlight the capability justifying quantitative results. PAM-UNet segments both liver and tumour precisely detecting boundaries better than MobileUNet and AttUNet. DeepLabv3+, FCN8 and ResUNet due to overfitting takes peripheral regions in segmentation mask also while segmenting liver. V Results V-A Quantitative results From Table I, we note that proposed PAM-UNet outperforms both backbone and attention mechanism based baselines. Overall, PAM-UNet secures the top spot in Dice score across both datasets, surpassing complex models like U-Net (ResNet50), DeepLabv3+, and FCN8. It even challenges the impressive performance of attention-based variants like AttUNet, achieving comparable mIoU and Dice scores while maintaining significantly lower computational cost of 1.32 FLOPS (vs. 35.35 FLOPS). This efficiency advantage becomes even more evident when compared to MobileUNet, where PAM-UNet achieves substantial improvements in Dice score of 11.83% on LiTS and 8.35% on Kvasir-SEG and mIoU of 11.72% on LiTS and 6.75% on Kvasir-SEG, while requiring only marginally more FLOPs (1.32 vs. 0.85). TABLE II: Effect of mobile blocks and different attentions. | w/ Mob (M) Encoder only w/ M Encoder-decoder (MED) MED + Self attention | 71.04 75.89 78.32 | 62.93 71.28 73.56 | 0.85 0.72 0.98 | | --- | --- | --- | --- | | MED + Cross attention | 81.26 | 74.44 | 1.24 | | MED + Additive attention | 79.40 | 74.34 | 1.17 | | MED+ $\mathcal{PLA}$ (Ours) | 82.87 | 74.65 | 1.32 | V-B Qualitative results As shown in Fig. 3 PAM-UNet consistently identifies regions of interest with more refined boundaries, outperforming baselines. PAM-UNet demonstrates precise predictions of lesions, even in cases involving varying small locations, sizes, and modality. Furthermore, it efficiently suppresses irrelevant information, including background elements. V-C Ablation study Table II summarizes the ablation study, revealing the $\mathcal{PLA}$ as the most effective mechanism, resulting in significant enhancements in mIoU and Dice scores. Solely relying on self-attention may not fully exploit contextual information from other parts of the input, resulting in suboptimal segmentation performance. Although cross attention improves segmentation performance compared to self-attention alone, the increased computational cost may limit its practicality for resource-constrained scenarios. Whereas, additive attention introduces a learnable transformation matrix to compute attention scores, offering flexibility in capturing complex relationships between input tokens. However, the performance improvement is modest compared to cross attention. Our proposed method, $\mathcal{PLA}$ effectively captures long-range dependencies while incorporating positional information, allowing the model to better understand spatial relationships within the input. Although it incurs a slightly higher computational cost compared to other attention mechanisms, the substantial performance gains justify its adoption. By striking a balance between accuracy and computational efficiency, $\mathcal{PLA}$ emerges as the preferred choice for semantic segmentation tasks. V-D Discussion Experimental results and ablation studies shows that the significant improvement in Dice and mIoU scores stems primarily from the incorporation of $\mathcal{PLA}$ . Fig. 2 (c) demonstrates this effect by illustrating how PLA alters feature map representations at each decoding layer. Furthermore, Fig. 1 (a) highlights how the symmetric structure of PAM-UNet ensures equal learning opportunities across all layers. As shown in Fig. 1 (b, c and d), PAM-UNet exhibits greater learning efficiency from initial layers compared to UNet, AttUNet, and MobileUNet. This implies that PAM-UNet achieves comparable representation learning with a very shallow light-weight encoder. <details> <summary>2405.01503v1/extracted/5573681/embs-corner-case.png Details</summary>  ### Visual Description # Technical Document Extraction: Image Analysis ## Image Panel - **Label**: "Image" (top center, black text on gray background) - **Content**: Endoscopic view of the gastrointestinal tract (likely colon/rectum) - **Textual Elements**: - **Timestamp**: "09/06/2015 19:12:24" (bottom left corner, white text on black background) - **Duration**: "00:03:38" (bottom left corner, white text on black background) - **Scale Bar**: "Scale: 1 cm" (bottom left corner, white text on black background) - **Date**: "09/06/2015" (bottom left corner, white text on black background) - **Visual Features**: - Mucosal surface with visible folds and vascular patterns - Reflective sheen from endoscope light - No visible lesions or abnormalities in the provided view ## GT Mask Panel - **Label**: "GT Mask" (top center, black text on gray background) - **Content**: Binary mask representation of anatomical structures - **Textual Elements**: - No additional text present - **Visual Features**: - White polygonal shapes on black background - Shapes correspond to anatomical landmarks (e.g., folds, valves) - No scale or coordinate system indicated ## Prediction Panel - **Label**: "Prediction" (top center, red text on gray background) - **Content**: Model-generated segmentation output - **Textual Elements**: - No additional text present - **Visual Features**: - White polygonal shapes on black background - Shapes show partial overlap with GT Mask - No scale or coordinate system indicated ## Cross-Reference Analysis - **Legend Alignment**: No legend present in any panel - **Spatial Grounding**: - GT Mask and Prediction panels use identical coordinate system (implied by shape alignment) - No axis markers or numerical values visible - **Trend Verification**: Not applicable (no continuous data series) - **Component Isolation**: - Header: Panel titles ("Image", "GT Mask", "Prediction") - Main Content: Image and mask representations - Footer: Timestamp/duration metadata in Image panel ## Conclusion The image contains three panels: 1. A clinical endoscopic image with metadata (timestamp, duration, scale) 2. A ground truth segmentation mask 3. A model prediction output No numerical data, charts, or complex diagrams are present. All textual information is limited to panel labels and metadata in the Image panel. </details> Figure 4: Corner case showing where PAM-UNet failed to detect all 10 polyps, which is a rare case in the Kvasir-SEG dataset (No training image contains more than 5 polyps). V-E Limitations The PAM-UNet model, while showing promising performance, has limitations. As shown in Fig. 4, PAM-UNet failed to segment precisely. The investigation of the failure cases suggests that no image in training data consists of more than five polyps. Leveraging the training data with a variety of training samples can mitigate the issue. We also recognize the recent rise of Transformer based segmentation architectures. However, it should be noted that these architectures become very successful when combined with UNet-style architectures, indicating the necessity of U-shaped architectures and/or encoder-decoder styles in successfully locating the boundary of objects. Since our focus in this study is to provide an effective segmentation engine with lightweight architecture, and since Transformer based architectures are currently costly in computational burden, we focus on the architecture stemming from lightweight and U-shaped architectures. In future work, we will extend our strategy by encoding linear Transformers in the same problem setting. Also we will evaluate proposed method on 3D segmentation task over variety of modalities. VI Conclusion In summary, we present a novel medical image segmentation architecture, PAM-UNet, that utilizes mobile convolutions and $\mathcal{PLA}$ to achieve state-of-the-art biomedical image segmentation performance against the baselines while maintaining a low computational cost of only 1.32 FLOPS. Our experiments demonstrate superior performance on LiTS and Kvasir-SEG benchmarks, with ablation studies confirming the crucial role of progressive attention in capturing long-range dependencies. Qualitative analysis further showcases PAM-UNet’s ability to learn relevant representations from initial layers. While acknowledging limitations in specific scenarios, we see PAM-UNet as an efficient solution for accurate and efficient medical image segmentation, paving the way for future advancements in the field. References - [1] Bhakti Baheti, Shubham Innani, Suhas Gajre, and Sanjay Talbar. Eff-unet: A novel architecture for semantic segmentation in unstructured environment. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 358–359, 2020. - [2] Patrick Bilic, Patrick Christ, Hongwei Bran Li, Eugene Vorontsov, Avi Ben-Cohen, Georgios Kaissis, Adi Szeskin, Colin Jacobs, Gabriel Efrain Humpire Mamani, Gabriel Chartrand, et al. The liver tumor segmentation benchmark (lits). Medical Image Analysis, 84:102680, 2023. - [3] Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam. Encoder-decoder with atrous separable convolution for semantic image segmentation. In Proceedings of the European Conference on Computer Vision (ECCV), September 2018. - [4] Junlong Cheng, Shengwei Tian, Long Yu, Hongchun Lu, and Xiaoyi Lv. Fully convolutional attention network for biomedical image segmentation. Artificial Intelligence in Medicine, 107:101899, 2020. - [5] François Chollet. Xception: Deep learning with depthwise separable convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1251–1258, 2017. - [6] Foivos I Diakogiannis, François Waldner, Peter Caccetta, and Chen Wu. ResUNet-a: A deep learning framework for semantic segmentation of remotely sensed data. ISPRS Journal of Photogrammetry and Remote Sensing, 162:94–114, 2020. - [7] Ke Dong, Chengjie Zhou, Yihan Ruan, and Yuzhi Li. Mobilenetv2 model for image classification. In Proceedings of the International Conference on Information Technology and Computer Application (ITCA), pages 476–480, 2020. - [8] Debesh Jha, Michael A Riegler, Dag Johansen, Pål Halvorsen, and Håvard D Johansen. DoubleU-Net: A Deep Convolutional Neural Network for Medical Image Segmentation. In Proceedings of the 33rd International symposium on computer-based medical systems (CBMS), pages 558–564, 2020. - [9] Debesh Jha, Pia H Smedsrud, Michael A Riegler, Pål Halvorsen, Thomas de Lange, Dag Johansen, and Håvard D Johansen. Kvasir-SEG: A segmented polyp dataset. In Proceedings of the International Conference on Multimedia Modeling (MMM), pages 451–462, 2020. - [10] Junfeng Jing, Zhen Wang, Matthias Rätsch, and Huanhuan Zhang. Mobile-unet: An efficient convolutional neural network for fabric defect detection. Textile Research Journal, 92(1-2):30–42, 2022. - [11] Minh-Thang Luong, Hieu Pham, and Christopher D Manning. Effective approaches to attention-based neural machine translation. arXiv preprint arXiv:1508.04025, 2015. - [12] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234–241, 2015. - [13] Nikhil Kumar Tomar, Debesh Jha, Michael A Riegler, Håvard D Johansen, Dag Johansen, Jens Rittscher, Pål Halvorsen, and Sharib Ali. Fanet: A feedback attention network for improved biomedical image segmentation. IEEE Transactions on Neural Networks and Learning Systems, 2022. - [14] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. - [15] Sihan Wang, Lei Li, and Xiahai Zhuang. Attu-net: attention u-net for brain tumor segmentation. In International MICCAI Brainlesion Workshop, pages 302–311, 2021. - [16] Tinghua Wang, Yunzhi Qiu, and Jialin Hua. Centered kernel alignment inspired fuzzy support vector machine. Fuzzy Sets and Systems, 394:110–123, 2020. - [17] Lingjuan Yu, Zhaoxin Zeng, Ao Liu, Xiaochun Xie, Haipeng Wang, Feng Xu, and Wen Hong. A lightweight complex-valued deeplabv3+ for semantic segmentation of polsar image. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 15:930–943, 2022. - [18] Rui Zhang, Jie Chen, Liang Feng, Shuang Li, Wei Yang, and Ding Guo. A refined pyramid scene parsing network for polarimetric sar image semantic segmentation in agricultural areas. IEEE Geoscience and Remote Sensing Letters, 19:1–5, 2021.