[Paper Review] Medical Image Segmentation via Sparse Coding Decoder

Long Zeng, Kaigui Wu|arXiv (Cornell University)|Oct 17, 2023

Advanced Image and Video Retrieval TechniquesComputer Science3 citations

TL;DR

This paper proposes CASCSCDE, a novel cascaded multi-layer convolutional sparse vector coding decoder that enhances medical image segmentation by representing features as sparse vectors to improve spatial detail recovery. Integrated with TransUNet, it achieves a 3.15% absolute improvement in Dice score and 1.16% in mIoU on the Synapse multi-organ segmentation benchmark, demonstrating superior boundary reconstruction and robustness through sparse encoding and multi-level feature fusion.

ABSTRACT

Transformers have achieved significant success in medical image segmentation, owing to its capability to capture long-range dependencies. Previous works incorporate convolutional layers into the encoder module of transformers, thereby enhancing their ability to learn local relationships among pixels. However, transformers may suffer from limited generalization capabilities and reduced robustness, attributed to the insufficient spatial recovery ability of their decoders. To address this issue, A convolution sparse vector coding based decoder is proposed , namely CAScaded multi-layer Convolutional Sparse vector Coding DEcoder (CASCSCDE), which represents features extracted by the encoder using sparse vectors. To prove the effectiveness of our CASCSCDE, The widely-used TransUNet model is chosen for the demonstration purpose, and the CASCSCDE is incorporated with TransUNet to establish the TransCASCSCDE architecture. Our experiments demonstrate that TransUNet with CASCSCDE significantly enhances performance on the Synapse benchmark, obtaining up to 3.15\% and 1.16\% improvements in DICE and mIoU scores, respectively. CASCSCDE opens new ways for constructing decoders based on convolutional sparse vector coding.

Motivation & Objective

• To address the limited spatial recovery and reduced robustness of transformer-based decoders in medical image segmentation.
• To improve generalization and noise resilience in segmentation by leveraging sparse vector representations.
• To develop a decoder that effectively fuses high-level semantic features with low-level spatial details for precise boundary localization.
• To demonstrate the effectiveness of convolutional sparse coding in a cascaded multi-layer architecture for 2D medical image segmentation.

Proposed method

• The CASCSCDE decoder uses a cascaded multi-layer convolutional sparse coding framework to represent features extracted by the encoder as sparse vectors.
• It employs iterative optimization within a learnable neural module (ML-block) that alternates between transposed convolutions and residual error correction to refine feature reconstruction.
• The ML-block performs iterative updates using learned weights W1, W2, and scaling factors c1, c2, with batch normalization and ReLU activation to enforce sparsity and stability.
• The decoder is integrated into the TransUNet architecture, replacing its original decoder to form TransCASCSCDE, enabling end-to-end training.
• Sparse coding is applied through a multi-layer structure that suppresses irrelevant information and enhances feature interpretability and robustness.
• The model is trained using a combination of cross-entropy and Dice loss, with data augmentation via random flipping and rotation on the Synapse dataset.

Experimental results

Research questions

• RQ1 Can a sparse coding-based decoder improve the spatial detail recovery capability of vision transformer encoders in medical image segmentation?
• RQ2 How does multi-layer convolutional sparse coding compare to standard deconvolutional decoders in terms of segmentation accuracy and robustness?
• RQ3 What is the optimal number of iterative refinement steps (T) in the sparse coding process for balancing sparsity and performance?
• RQ4 Does the proposed CASCSCDE decoder generalize across organs of varying sizes, especially small and complex-structured ones?
• RQ5 To what extent does sparse coding enhance noise resilience and boundary precision in medical image segmentation?

Key findings

• TransCASCSCDE achieves a 3.15% absolute improvement in mean Dice score (80.63%) compared to the baseline TransUNet on the Synapse multi-organ segmentation benchmark.
• The model improves mIoU by 1.16%, indicating enhanced overall segmentation consistency and boundary localization.
• Significant gains are observed for small organs: +4.05% for gallbladder, +2.13% for left kidney, and +1.11% for right kidney.
• Large organs also benefit: +6.69% for stomach, +3.47% for spleen, and +1.94% for pancreas, indicating strong performance on complex structures.
• An ablation study shows that T=2 provides optimal performance, with further increases in T (to 3 and 4) leading to performance degradation due to over-sparse representations.
• The sparse coding mechanism effectively suppresses noise and irrelevant features, enhancing model robustness and generalization across diverse anatomical structures.