ИИ для медицины: новый метод обучения нейросетей

Медицинские учреждения не всегда имеют возможность собирать данные в идеальном порядке. Клиника может располагать несколькими тщательно размеченными изображениями с одного сканера, но при этом иметь тысячи неразмеченных снимков из других центров, каждый со своими особенностями настроек, профилем пациентов и артефактами визуализации. Такая путаница делает задачу медицинской сегментации изображений еще сложнее. Модели, обученные на идеальных предположениях, могут сбоить при развертывании в других условиях, особенно при работе с небольшими, тусклыми или низкоконтрастными объектами.

Ассистент профессор Чжао На из SUTD и ее коллеги поставили своей целью справиться с этой неразберихой, а не игнорировать ее. Вместо стандартного подхода, когда размеченные и неразмеченные данные предполагаются принадлежащими к схожим распределениям, они работают в более реалистичном сценарии, известном как кросс-доменная полуавтоматическая генерализация домена (CD-SSDG). В этом сценарии немногочисленные размеченные изображения поступают из одного домена, в то время как обширный пул неразмеченных данных охватывает несколько различных доменов, что точно соответствует ситуации, с которой сталкиваются многие больницы.

Результаты команды подробно описаны в статье «Dual-supervised Asymmetric Co-training for Semi-supervised Medical Domain Generalization», опубликованной в IEEE Transactions on Multimedia. В настоящее время методы полуавтоматического обучения в основном полагаются на псевдо-метки. Модель, обученная на меньшем наборе размеченных данных, угадывает метки для неразмеченных изображений, а затем учится на этих догадках. Когда неразмеченные изображения сильно отличаются от размеченных, эти догадки становятся неверными, и ошибки накапливаются.

Решение исследователей — это двухсупервизорная асимметричная ко-тренинговая структура (DAC), в которой две подмодели обучаются параллельно. Они по-прежнему обмениваются псевдо-метками, но с важным дополнением: супервизией на уровне признаков. Вместо того чтобы полагаться только на догадки пиксель за пикселем, каждая модель также побуждает другую к согласованию в более богатом пространстве признаков, стимулируя согласие относительно базовой структуры, даже когда стиль и контраст различаются. Подмоделям также назначаются различные вспомогательные задачи самообучения: одна учится локализовать смешанный фрагмент в изображении CutMix; другая учится распознавать вращение фрагмента.

Эта асимметрия поддерживает разнообразие их внутренних представлений, снижая риск того, что обе модели впадут в одинаковые ошибки, и обостряя их способность отделять передний план от фона. «Как клиницисты и инженеры, мы редко имеем возможность выбирать идеальные наборы данных», — говорит ассистент профессор Чжао. «DAC — это наш способ добавить систему безопасности. Когда псевдо-метки хрупки, руководство на уровне признаков по-прежнему привязывает модель к стабильным, инвариантным к домену признакам. Асимметричные задачи затем побуждают двух обучающихся смотреть на данные под разными углами».

Тестирование на трех эталонных сценариях сегментации — глазное дно (диск и чашка зрительного нерва), изображения полипов толстой кишки и МРТ спинного мозга — показало, что DAC стабильно лучше обобщается на незнакомые домены, чем сильные базовые методы, включая методы, специально разработанные для генерализации доменов. Наиболее заметные улучшения наблюдались на небольших или низкоконтрастных структурах, таких как чашка зрительного нерва, где команда зафиксировала двузначное улучшение показателя Dice по сравнению с передовыми подходами при низком соотношении размеченных данных. Важно отметить, что вспомогательные задачи и супервизия признаков используются только во время обучения, поэтому стоимость вывода DAC соответствует стоимости обычных моделей.

«Что нас удивило, так это стабильность», — добавил ассистент профессор Чжао. «Даже когда мы уменьшали долю размеченных данных, в некоторых случаях до одной десятой, кривая не рушилась. Это дает уверенность больницам, которые могут размечать только небольшую часть данных каждый год, но при этом хотят, чтобы модели хорошо работали в различных условиях». Подход команды также прагматичен. Супервизия на уровне признаков действует как мягкое ограничение, не зависящее от точных пиксельных меток, которые notoriously зашумлены при сдвиге домена. Асимметричные задачи — локализация смешанных фрагментов и предсказание вращения случайных фрагментов — просты в реализации (по одной линейной голове каждая) и computationally легки, но они достаточно разнообразят двух обучающихся, чтобы со временем улучшить качество псевдо-меток.

Команда также наметила, где DAC можно усовершенствовать. Случаи сбоев включают изображения глазного дна, где несколько кровеносных сосудов пересекают диск, и сцены, где цель почти сливается с фоном. Будущие работы включают аугментацию, учитывающую сосуды, для изображений глазного дна и адаптивные многомерные представления, объединяющие многомасштабные и частотные признаки для уточнения границ в условиях низкого контраста.

«Эти компоненты не ограничиваются тремя протестированными наборами данных», — отметил ассистент профессор Чжао. «Визуализация опухолей сталкивается с теми же двойными проблемами — дорогостоящие аннотации и вариативность от центра к центру. DAC здесь применим немедленно, особенно там, где точные границы имеют клиническое значение». Хотя DAC является рецептом времени обучения, а не совершенно новой сетью, его влияние практично — более эффективное использование неразмеченных данных из разных центров без предположения, что мир независим и одинаково распределен. Метод также хорошо сочетается с существующими базовыми сетями (ResNet-DeepLabv3+ в текущем исследовании) и стандартными оптимизаторами, сокращая путь к внедрению. Исследователи сообщают о стабильных улучшениях на бенчмарках Fundus, Polyp и SCGM, более быстром обучении по сравнению с ведущим базовым методом ко-тренинга и отсутствии дополнительных затрат при развертывании.

«Прежде всего, важна генерализация», — сказал ассистент профессор Чжао. «Больницы хотят моделей, которые ведут себя предсказуемо, когда сканер другой, пациент другой, освещение другое. Супервизируя не только видимые нами метки, но и признаки, сохраняющиеся между доменами, мы приближаемся к этой цели».