NeRD: Neural Representation of Distribution for Medical Image Segmentation

TL;DR

NeRD introduces a neural module that models feature distribution shifts in CNNs, improving medical image segmentation accuracy by reducing over-segmentation and missing regions, validated on white matter lesion and atrial segmentation tasks.

Contribution

The paper presents NeRD, a novel neural module that estimates feature distributions to enhance CNN-based medical image segmentation, addressing distribution shift issues.

Findings

Improved segmentation accuracy on white matter lesions.

Reduced over-segmentation and missed regions.

Validated effectiveness on multiple medical imaging tasks.

Abstract

We introduce Neural Representation of Distribution (NeRD) technique, a module for convolutional neural networks (CNNs) that can estimate the feature distribution by optimizing an underlying function mapping image coordinates to the feature distribution. Using NeRD, we propose an end-to-end deep learning model for medical image segmentation that can compensate the negative impact of feature distribution shifting issue caused by commonly used network operations such as padding and pooling. An implicit function is used to represent the parameter space of the feature distribution by querying the image coordinate. With NeRD, the impact of issues such as over-segmenting and missing have been reduced, and experimental results on the challenging white matter lesion segmentation and left atrial segmentation verify the effectiveness of the proposed method. The code is available via https://github.com/tinymilky/NeRD.