Разделение сэмплирования и бюджета обучения в сегментации состава тела на КТ с несбалансированными классами

Аннотация

Дисбаланс классов является фундаментальной проблемой в сегментации медицинских изображений, где частые классы, как правило, доминируют в обучении за счет редких классов. Методы, основанные на функциях потерь, смягчают дисбаланс путем перевзвешивания попиксельных потерь в пределах батча, в то время как стратегии выборки контролируют, какие изображения попадают в батч. Однако ни один из этих подходов явно не контролирует, какие классы присутствуют в батче, оставляя экспозицию редких классов лишь частично сбалансированной. В данной работе мы применяем эпизодическую выборку из обучения с несколькими примерами для обеспечения сбалансированного по классам построения батчей в условиях полностью контролируемого обучения. Мы отделяем эпизодическую выборку от ее традиционного контекста метрического обучения и оцениваем ее в сегментации компонентов тела на КТ-изображениях. Мы сравниваем эпизодическую выборку со случайной и взвешенной выборкой на девяти типах мышечной и жировой тканей, полученных из 210 снимков общедоступного набора данных SAROS. Обучение проводится в режимах полного и малого объема данных, а также с дополнительными сравнениями при одинаковом бюджете итераций обучения. При обучении на полных данных все три стратегии показали сопоставимые результаты (средний Dice 0,882 для эпизодической, 0,878 для случайной и взвешенной). При обучении на малых данных эпизодическая выборка превзошла случайную и взвешенную (0,787 против 0,758 и 0,762), что обусловлено 12-кратной разницей в количестве итераций обучения. При одинаковом бюджете итераций случайная и взвешенная выборки переобучались раньше, тогда как эпизодическая улучшалась примерно в три раза дольше до выхода на плато. Наши результаты выявляют бюджет итераций обучения как недостаточно учитываемый confounding-фактор в стратегиях выборки, что мотивирует внедрение протоколов оценки, учитывающих количество итераций для малых наборов данных. Кроме того, остаточное преимущество эпизодической выборки согласуется с эффектом неявной регуляризации за счет сбалансированных по классам батчей, предлагая низкозатратную, модельно-независимую стратегию для сегментации медицинских изображений с дисбалансом классов. Код доступен по адресу https://github.com/iasonsky/episodic-sampling.

English

Class imbalance is a fundamental challenge in medical image segmentation, where frequent classes typically dominate training at the expense of rare classes. Loss-based approaches mitigate imbalance by reweighting the per-pixel loss within the batch, while sampling strategies control which images enter the batch. Yet neither explicitly controls which classes appear within the batch, leaving rare-class exposure only partially rebalanced. In this work, we adopt episodic sampling from few-shot learning to promote class-balanced batch construction in a fully supervised setting. We decouple episodic sampling from its conventional metric-learning context and evaluate it in body composition segmentation in CT. We compare episodic sampling against random and weighted sampling on nine muscle and adipose tissues, derived from 210 scans of the public SAROS dataset. Training is performed under full- and low-data regimes, with additional comparisons under matched training iteration budgets. Under full-data training, all three strategies performed comparably (mean Dice 0.882 for episodic, 0.878 for random and weighted). Under low-data training, episodic sampling outperformed random and weighted (0.787 vs. 0.758 and 0.762), driven by a 12-fold difference in training iterations. Under matched training budgets, random and weighted overfit earlier, while episodic improved for approximately three times more iterations before plateauing. Our findings identify the training iteration budget as under-recognized confound in sampling strategies, motivating iteration-aware evaluation protocols for small datasets. Furthermore, the residual advantage of episodic sampling is consistent with an implicit regularization effect of class-balanced batches, offering a low-cost, model-agnostic strategy for class-imbalanced medical image segmentation. Code is available at https://github.com/iasonsky/episodic-sampling.