MicroAUNet: Boundary-Enhanced Multi-scale Fusion with Knowledge Distillation for Colonoscopy Polyp Image Segmentation

AI summary

Overview

• Research area: Medical image segmentation, specifically colorectal polyp segmentation in colonoscopy images; lightweight network design and knowledge distillation for computer vision.
• Technical level: Advanced (assumes familiarity with encoder-decoder segmentation networks, attention mechanisms, depthwise separable convolutions, and knowledge distillation losses).
• Scope: The paper proposes and evaluates MicroAUNet, a compact attention-based segmentation network trained with progressive two-stage knowledge distillation, benchmarked on Kvasir-SEG and CVC-ClinicDB.

What This Paper Is About

Deep learning models for polyp segmentation face a trade-off: accurate models often produce ambiguous polyp margins, while lightweight models that can run in real time tend to lose fine boundary detail. This paper aims to build a very small network (0.0249M parameters) that recovers boundary precision through depthwise separable dilated convolutions, a parameter-shared channel-spatial attention design, and a two-stage distillation scheme that transfers both feature-level and decision-boundary knowledge from a larger teacher. The stated goal is real-time clinical polyp segmentation with accuracy comparable to or better than heavier models.

Key Contributions

1. MicroAUNet architecture: A lightweight attention-based segmentation network combining depthwise separable dilated convolutions with a single-path, parameter-shared channel-spatial attention mechanism to improve boundary extraction and adaptability to polyps of diverse morphology.
2. Progressive two-stage knowledge distillation: A scheme that transfers semantic and boundary knowledge from a high-capacity teacher model (a MALUNet teacher) to the compact MicroAUNet student, decoupling feature imitation from semantic distribution matching to avoid gradient conflict.
3. Preference alignment stage: A contrastive loss built from the teacher's prediction confidence, using positive samples above τ_h = 0.8 and negative samples below τ_l = 0.2, to help the student inherit the teacher's decision boundaries in ambiguous regions.
4. Comprehensive validation: Experiments on Kvasir-SEG and CVC-ClinicDB comparing against UNet, SANet, UNeXt, and MALUNet, plus an ablation study isolating each component.

Main Findings

• Kvasir-SEG performance: MicroAUNet reaches 0.904 mDice, 0.861 mIoU, 0.951 Acc, 0.964 Spe, and 0.902 Sen, the highest mDice and mIoU among the compared models on this dataset. For comparison, UNet scores 0.818 mDice / 0.746 mIoU, SANet 0.886 / 0.847, UNeXt 0.879 / 0.837, and MALUNet 0.892 / 0.857.
• CVC-ClinicDB performance: MicroAUNet reaches 0.902 mDice, 0.869 mIoU, 0.957 Acc, 0.962 Spe, and 0.908 Sen. MALUNet is slightly higher on this dataset at 0.906 mDice / 0.875 mIoU, and the paper attributes the differing relative performance between datasets to differences in data distribution and architecture-specific feature representations.
• Extreme parameter reduction: MicroAUNet has 0.0249M parameters versus 7.77M for UNet, 23.90M for SANet, 0.30M for UNeXt, and 0.175M for MALUNet. Its GFLOPs are 0.148, slightly higher than MALUNet's 0.083 but well below UNet's 13.78 and SANet's 5.99.
• Lightweight design ablation (MicroAUNet-1): Removing depthwise separable dilated convolutions reduced Kvasir mDice by 1.34% (0.904 to 0.892) and mIoU by 0.47% (0.861 to 0.857). On CVC, MicroAUNet-1 (0.906 mDice) slightly exceeded the full model (0.902), which the authors attribute to clearer lesion boundaries in that dataset.
• Imitation learning ablation (MicroAUNet-2): Removing the imitation learning stage caused a 3.21% mDice drop (0.904 to 0.875) on Kvasir and a 2.71% mIoU decrease (0.869 to 0.846) on CVC, supporting the claim that teacher-student distillation is critical for distinguishing subtle lesions.
• Preference alignment ablation (MicroAUNet-3): Removing the preference contrastive learning mechanism reduced mDice by 3.43% and mIoU by 2.91% on Kvasir, indicating that contrastive learning strengthens intra-class feature consistency.
• Cross-dataset stability: The full model shows an mDice standard deviation of 0.0014 (0.904 vs. 0.902), lower than MicroAUNet-2 (0.0015) and MicroAUNet-3 (0.0040). The paper also states the complete model outperforms MicroAUNet-3 by 3.09% at its maximum Kvasir mDice and CVC mIoU.
• Qualitative results: U-Net and SANet miss polyp edges or produce fragmented masks; UNeXt and MALUNet are more complete but blur fine boundaries in low-contrast areas; MicroAUNet maintains boundary continuity under occlusion and irregular shapes and removes background artefacts.
• Inference speed: The paper states MicroAUNet achieves faster inference than the compared state-of-the-art models, but no per-frame latency or throughput measurements are reported. The authors list per-frame latency and throughput evaluation as future work.

Methodology in Plain English

The starting point is a U-shaped encoder-decoder network in the style of MALUNet. The authors change two things inside it. First, convolutions are split into a depthwise dilated stage for spatial sampling and a 1×1 pointwise stage for channel mixing, which reduces the parameter cost of a convolution from K×K×C×C to K×K×C while still capturing multi-scale context. Second, instead of multiple attention branches with separate parameters, they use a single-path spatial attention module—a depthwise convolution, GELU, then a 1×1 convolution and sigmoid to produce an attention mask, applied with a residual connection—plus a parameter-shared channel-spatial attention bridge reused across all decoder stages, combining channel and spatial attention with learnable weights α and β (both initialized to 0.5).

Training happens in two sequential stages rather than one combined objective, because the authors argue that low-level feature imitation and high-level semantic matching produce conflicting gradients. In stage one (imitation learning), the student matches teacher features at five encoder stages using an L2 mimicry loss and matches the teacher's output distribution using KL divergence. A cosine curriculum schedules the KL weight ω_KL, with T set to half of the total epochs, so early training focuses on features and later training on semantics. In stage two (preference alignment), the teacher's predictions are used to define high-confidence positive regions (probability at or above 0.8) and low-confidence negative regions (at or below 0.2); a contrastive loss pulls positive pairs together and pushes positive-negative pairs apart with margin m. The final objective adds a regularization term weighted by ρ = 0.3, and parameters are updated with EMA.

All models were trained for 300 epochs with batch size 8, the AdamW optimizer, an initial learning rate of 0.001, a cosine annealing scheduler, and a fixed random seed of 42, on hardware consisting of an Intel Xeon Gold 6330 CPU at 2.1 GHz with 128 GB RAM and an NVIDIA GeForce RTX 3090 GPU with 24 GB memory. Evaluation uses mean Dice, mean IoU, Accuracy, Specificity, and Sensitivity.

Why This Matters

Impact on research. The paper argues that the accuracy-efficiency trade-off in polyp segmentation is not necessarily zero-sum: a 0.0249M-parameter model can match or exceed much larger models on some benchmarks when paired with careful distillation. It also proposes decoupling feature imitation from semantic distillation into sequential stages, offering a concrete design choice for others working on compact medical segmentation models.

Real-world applications:

• Computer-aided detection during colonoscopy, where low miss rates matter—the paper cites a traditional colonoscopy miss rate of approximately 20-30% and notes physician fatigue and difficulty detecting small or flat polyps.
• Real-time decision support on clinical workstations or endoscopy suites where a 7.77M or 23.90M parameter model may be impractical to run at video rate.
• Deployment on edge or resource-constrained hardware in clinics, given the sub-0.03M parameter footprint.
• Training-data-efficient development of similar compact models for other medical imaging segmentation tasks, using a larger teacher network as a knowledge source.

Industry relevance. Medical device and endoscopy software vendors developing embedded or latency-sensitive AI tools have direct interest in a model family that claims competitive accuracy at a fraction of the parameter count. The code is publicly available at https://github.com/JeremyXSC/MicroAUNet, which lowers the barrier to reproduction and product prototyping.

Future Directions

• Cross-domain generalization: The authors state that generalization to unseen clinical environments—different endoscopic devices or lighting conditions—has not been fully validated, and call for evaluation on multi-centre or real-world datasets.
• Continuous video segmentation: Extending MicroAUNet from still frames to colonoscopy video by incorporating temporal information across adjacent frames, with temporal consistency constraints and frame-to-frame feature or mask propagation to reduce mask fluctuations, plus lightweight tracking to handle camera motion and occlusion.
• Direct latency benchmarking: Measuring per-frame latency and throughput on continuous colonoscopy streams to substantiate real-time applicability, which the current evaluation does not report.
• Reducing teacher dependence: The two-stage distillation framework relies on teacher quality, and the authors note that biased or noisy teacher representations may limit student performance. They also plan to explore domain adaptation and contrastive learning for cross-domain generalization.

Target Audience

Researchers and graduate students in medical image analysis and computer vision working on lightweight segmentation architectures, attention mechanisms, or knowledge distillation. It is also relevant to applied machine learning engineers building real-time clinical decision-support tools, and to clinical AI researchers interested in the feasibility of extremely compact models for colonoscopy. Readers need prior familiarity with encoder-decoder segmentation networks and distillation losses to follow the method section; the results tables and ablation discussion are accessible to a broader technical audience.

Authors’ abstract

Early and accurate segmentation of colorectal polyps is critical for reducing colorectal cancer mortality, which has been extensively explored by academia and industry. However, current deep learning-based polyp segmentation models either compromise clinical decision-making by providing ambiguous polyp margins in segmentation outputs or rely on heavy architectures with high computational complexity, resulting in insufficient inference speeds for real-time colorectal endoscopic applications. To address this problem, we propose MicroAUNet, a lightweight attention-based segmentation network, which synergistically combines depthwise-separable dilated convolutions with a single-path, parameter-shared channel-spatial attention block to effectively strengthen multi-scale boundary features. On the basis of it, a progressive two-stage knowledge-distillation scheme is introduced to transfer semantic and boundary cues from a high-capacity teacher. Extensive experiments on benchmarks also demonstrate the state-of-the-art accuracy under extremely low model complexity, indicating that MicroAUNet is suitable for real-time clinical polyp segmentation. The code is publicly available at https://github.com/JeremyXSC/MicroAUNet.