# KLO-Net: A Dynamic K-NN Attention U-Net with CSP Encoder for Efficient Prostate Gland Segmentation from MRI

arXiv: https://arxiv.org/abs/2512.13902 · HTML: https://mlsanity.com/paper/2512.13902/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 31 references that are in the map.
- First posted on arXiv 15 Dec 2025.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [YOLOv4: Optimal Speed and Accuracy of Object Detection](https://mlsanity.com/paper/2004.10934/) (2020)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [CSPNet: A New Backbone that can Enhance Learning Capability of CNN](https://mlsanity.com/paper/1911.11929/) (2019)
- [Understanding Robustness of Transformers for Image Classification](https://mlsanity.com/paper/2103.14586/) (2021)
- [CE-Net: Context Encoder Network for 2D Medical Image Segmentation](https://mlsanity.com/paper/1903.02740/) (2019)
- [Boundary loss for highly unbalanced segmentation](https://mlsanity.com/paper/1812.07032/) (2018)
- [A Novel Focal Tversky loss function with improved Attention U-Net for lesion segmentation](https://mlsanity.com/paper/1810.07842/) (2018)
- [U-Net Transformer: Self and Cross Attention for Medical Image Segmentation](https://mlsanity.com/paper/2103.06104/) (2021)
- [KVT: k-NN Attention for Boosting Vision Transformers](https://mlsanity.com/paper/2106.00515/) (2021)
- [CAT-Net: A Cross-Slice Attention Transformer Model for Prostate Zonal Segmentation in MRI](https://mlsanity.com/paper/2203.15163/) (2022)
- [CSAM: A 2.5D Cross-Slice Attention Module for Anisotropic Volumetric Medical Image Segmentation](https://mlsanity.com/paper/2311.04942/) (2023)

## Abstract (arXiv)

Real-time deployment of prostate MRI segmentation on clinical workstations is often bottlenecked by computational load and memory footprint. Deep learning-based prostate gland segmentation approaches remain challenging due to anatomical variability. To bridge this efficiency gap while still maintaining reliable segmentation accuracy, we propose KLO-Net, a dynamic K-Nearest Neighbor attention U-Net with Cross Stage Partial, i.e., CSP, encoder for efficient prostate gland segmentation from MRI scan. Unlike the regular K-NN attention mechanism, the proposed dynamic K-NN attention mechanism allows the model to adaptively determine the number of attention connections for each spatial location within a slice. In addition, CSP blocks address the computational load to reduce memory consumption. To evaluate the model's performance, comprehensive experiments and ablation studies are conducted on two public datasets, i.e., PROMISE12 and PROSTATEx, to validate the proposed architecture. The detailed comparative analysis demonstrates the model's advantage in computational efficiency and segmentation quality.