Is Exchangeability better than I.I.D to handle Data Distribution Shifts while Pooling Data for Data-scarce Medical image segmentation?

TL;DR

This paper explores whether assuming exchangeability rather than i.i.d. conditions better handles distribution shifts in data pooling for medical image segmentation, proposing a method that improves feature representation and achieves state-of-the-art results.

Contribution

It introduces a novel approach based on causal insights to control feature discrepancies across network layers, enhancing segmentation performance under data addition scenarios.

Findings

Achieves state-of-the-art segmentation on five datasets.

Demonstrates improved feature representations across models.

Provides qualitative evidence of more accurate segmentation maps.

Abstract

Data scarcity is a major challenge in medical imaging, particularly for deep learning models. While data pooling (combining datasets from multiple sources) and data addition (adding more data from a new dataset) have been shown to enhance model performance, they are not without complications. Specifically, increasing the size of the training dataset through pooling or addition can induce distributional shifts, negatively affecting downstream model performance, a phenomenon known as the "Data Addition Dilemma". While the traditional i.i.d. assumption may not hold in multi-source contexts, assuming exchangeability across datasets provides a more practical framework for data pooling. In this work, we investigate medical image segmentation under these conditions, drawing insights from causal frameworks to propose a method for controlling foreground-background feature discrepancies across all layers of deep networks. This approach improves feature representations, which are crucial in data-addition scenarios. Our method achieves state-of-the-art segmentation performance on histopathology and ultrasound images across five datasets, including a novel ultrasound dataset that we have curated and contributed. Qualitative results demonstrate more refined and accurate segmentation maps compared to prominent baselines across three model architectures.