nnY-Net: Swin-NeXt с кросс-вниманием для сегментации 3D медицинских изображений

Основная информация

• ID статьи: 2501.01406
• Название: nnY-Net: Swin-NeXt with Cross-Attention for 3D Medical Images Segmentation
• Авторы: Haixu Liu¹, Zerui Tao¹, Wenzhen Dong², Qiuzhuang Sun¹
• Учреждения: ¹Университет Сиднея, ²Китайский университет Гонконга
• Классификация: cs.CV (компьютерное зрение)
• Ссылка на статью: https://arxiv.org/abs/2501.01406

Аннотация

В данной работе предложена новая архитектура модели nnY-Net для сегментации 3D медицинских изображений. Модель получила своё название благодаря Y-образной структуре, образованной добавлением модуля кросс-внимания в нижней части U-Net. Авторы интегрировали преимущества двух современных моделей SOTA — MedNeXt и SwinUNETR, используя Swin Transformer в качестве кодировщика и ConvNeXt в качестве декодировщика, создав инновационную архитектуру Swin-NeXt. Модель использует карты признаков самого низкого уровня кодировщика в качестве Key и Value, а клинические характеристики пациента (патологическую и терапевтическую информацию) в качестве Query для вычисления весов кросс-внимания. Кроме того, на основе фреймворков dynUnet и nnU-Net упрощены методы предварительной и постобработки для 3D сегментации изображений, а также разработана функция потерь DiceFocalCELoss для повышения эффективности обучения при классификации несбалансированных вокселей.

Исследовательский контекст и мотивация

Определение проблемы

Сегментация 3D CT-изображений опухолей печени является важной задачей в сегментации медицинских изображений. Точная сегментация CT-изображений помогает врачам оценить объём опухоли и разработать адекватный план лечения.

Ограничения существующих методов

1. Ограничения традиционных методов: До 2016 года в основном использовались неконтролируемые алгоритмы, такие как рост регионов, с ограниченной точностью
2. Дефекты единственной архитектуры: Существующие методы основаны либо на чистой свёртке (например, nnU-Net), либо на чистом Transformer (например, SwinUNETR), не полностью объединяя преимущества обоих подходов
3. Недостаточное мультимодальное слияние: Существующие методы в основном обрабатывают информацию изображения, неэффективно используя патологическую и терапевтическую информацию пациента
4. Проблема дисбаланса классов: Фоновая метка занимает более 90% пикселей, что затрудняет обучение

Исследовательская мотивация

Данная работа направлена на разработку мультимодальной модели сегментации, объединяющей преимущества Transformer и свёрточных нейронных сетей, способной одновременно обрабатывать 3D-изображения и клиническую информацию пациента для повышения точности и практичности сегментации медицинских изображений.

Основные вклады

1. Предложена архитектура Swin-NeXt: Инновационное объединение кодировщика SwinUNETR с декодировщиком MedNeXt, полностью использующее преимущества Transformer в извлечении признаков и свёртки в пиксельном декодировании
2. Разработан механизм кросс-внимания для слияния: Предложены три метода слияния признаков, выявлено, что механизм кросс-внимания показывает лучшие результаты при мультимодальном слиянии и стабильно повышает производительность модели
3. Построена функция потерь DiceFocalCELoss: Объединение DiceLoss, FocalLoss и кросс-энтропийной потери эффективно решает проблему дисбаланса классов при классификации вокселей
4. Упрощение процесса предварительной обработки: На основе фреймворков dynUnet и nnU-Net упрощены и оптимизированы методы предварительной и постобработки для 3D сегментации изображений

Подробное описание методов

Определение задачи

Входные данные:

• 3D CT-сканирование χ ∈ R^(H×W×D×C)
• Клиническая информация пациента (патологическая, терапевтическая информация и т.д.)

Выходные данные: Многоклассовая маска сегментации, включающая печень, опухоль, кровеносные сосуды, аорту и другие структуры

Ограничения: Ограничения памяти при обработке высокоразрешённых 3D-изображений, задача сегментации с экстремальным дисбалансом классов

Архитектура модели

1. Кодировщик Swin Transformer

Использует 3D Swin Transformer в качестве кодировщика с конкретной реализацией:

• Разбиение на патчи: Разделение входного изображения на 3D окна размером M×M×M
• Механизм оконного внимания:

  z^l = W-MSA(LN(z^(l-1))) + z^(l-1)
  z^l = MLP(LN(z^l)) + z^l
  z^(l+1) = SW-MSA(LN(z^l)) + z^l
  z^(l+1) = MLP(LN(z^(l+1))) + z^(l+1)
  

• Вычисление внимания:

  Attention(Q,K,V) = Softmax(QK^T/√d + B)V
  

2. Декодировщик ConvNeXt

Использует структуру декодировщика MedNeXt:

• Транспонированная свёртка для повышения дискретизации:

  Y_{i,j,k} = ∑∑∑ K_{p,q,r} · X_{(expanded)}_{i+p,j+q,k+r}
  

• Функция активации GELU:

  GELU(x) = x/2[1 + erf(x/√2)]
  

3. Модуль кросс-внимания для слияния

Интеграция мультимодальной информации в нижнем слое кодировщика (Bottleneck):

• Key и Value: Карты признаков самого низкого уровня кодировщика
• Query: Вектор клинических характеристик пациента
• Стратегия слияния: Отображение клинических характеристик в подходящее измерение через полносвязный слой, затем вычисление кросс-внимания

Технические инновации

1. Проектирование гибридной архитектуры: Первое объединение кодировщика Swin Transformer с декодировщиком ConvNeXt, использующее преимущества каждого
2. Мультимодальное кросс-внимание: Инновационное использование механизма кросс-внимания для слияния признаков изображения и клинической информации
3. Комбинированная функция потерь: Разработка DiceFocalCELoss для решения проблемы дисбаланса классов в сегментации медицинских изображений

Экспериментальная установка

Набор данных

• Данные изображений: 110 CT-сканирований печени от 98 пациентов
• Клинические данные: 56 переменных патологической и терапевтической информации пациента
• Очистка данных: Удалены 4 проблемных набора данных (HCC 017, 008, 025, 009)
• Обработка пропущенных значений: Использование моделей машинного обучения для заполнения пропусков

Метрики оценки

1. Коэффициент Dice: Dice = 2×|X∩Y|/(|X|+|Y|)
2. Средний IoU: MIoU = (1/N)∑|X_i∩Y_i|/|X_i∪Y_i|
3. Расстояние Хаусдорфа (HD95): 95-й процентиль расстояния Хаусдорфа
4. Точность, полнота, прецизионность: Стандартные метрики классификации

Методы сравнения

• U-Net
• UNETR
• SwinUNETR
• MedNeXt
• Предложенные Swin-NeXt и его варианты

Детали реализации

• Фреймворк: Разработано на основе MONAI, совместимо с Jupyter Notebook
• Предварительная обработка: Обрезка изображений, нормализация, передискретизация, увеличение данных
• Стратегия обучения: Стратегия скользящего окна, блочная обработка для предотвращения переполнения памяти

Результаты экспериментов

Основные результаты

Сравнение производительности при сегментации опухолей печени:

Абляционные исследования

Сравнение методов слияния признаков:

Ключевые выводы:

1. Механизм кросс-внимания обеспечивает стабильное повышение производительности на всех моделях
2. Методы Add и Concat не сходятся при использовании с моделями с транспонированной свёрткой (MedNeXt, Swin-NeXt)
3. Механизм кросс-внимания имеет лучшую совместимость со структурой транспонированной свёртки

Анализ конкретных случаев

Визуализация результатов сегментации CT-срезов пациента HCC066 показывает, что хотя Swin-NeXt не достигает оптимальных числовых показателей, форма и размер сегментации наиболее близки к истинной метке, обладая лучшей клинической практичностью.

Связанные работы

Развитие 3D сегментации медицинских изображений

1. Традиционные методы: Неконтролируемые методы, такие как рост регионов
2. Эпоха глубокого обучения: 3D U-Net заложил основу
3. Автоматизированные фреймворки: nnU-Net интегрирует автоматическую предварительную обработку и выбор параметров
4. Применение Transformer: UNETR вводит ViT, SwinUNETR использует Swin Transformer
5. Современная свёртка: MedNeXt разработан на основе ConvNeXt

Позиционирование вклада работы

Данная работа является первым систематическим объединением кодировщика Transformer и современного декодировщика свёртки для 3D сегментации медицинских изображений с введением механизма мультимодального кросс-внимания.

Заключение и обсуждение

Основные выводы

1. Эффективность архитектурного слияния: Объединение кодировщика Swin Transformer и декодировщика ConvNeXt является эффективным
2. Превосходство кросс-внимания: При мультимодальном слиянии механизм кросс-внимания значительно превосходит простые методы сложения и конкатенации
3. Эффективность комбинированной функции потерь: DiceFocalCELoss улучшает сходимость обучения при решении проблемы дисбаланса классов

Ограничения

1. Консервативные параметры: Для справедливого сравнения сокращено количество параметров, не использовано оптимальное соотношение блоков извлечения признаков 1:1:3:1
2. Ограничения вычислительных ресурсов: Из-за очередей на сервере не удалось завершить все запланированные эксперименты
3. Масштаб набора данных: Относительно небольшой набор данных может ограничить обобщающую способность модели

Направления будущих исследований

1. Исследование более оптимальных пропорций архитектуры сети
2. Проверка эффективности метода на более крупных наборах данных
3. Изучение применимости к другим задачам сегментации медицинских изображений

Глубокая оценка

Преимущества

1. Высокая инновационность: Первое систематическое объединение преимуществ Swin Transformer и ConvNeXt
2. Мультимодальное слияние: Эффективное использование клинической информации для вспомогательной сегментации изображений
3. Полные эксперименты: Включены подробные абляционные исследования и сравнительный анализ
4. Высокая практическая ценность: Разработано в соответствии с практическими потребностями сегментации медицинских изображений

Недостатки

1. Ограниченное повышение производительности: Числовое улучшение по сравнению с базовыми методами недостаточно значительно
2. Несогласованность экспериментальных условий: Существуют различия во времени обучения и условиях для разных моделей
3. Недостаточный теоретический анализ: Отсутствует глубокий анализ причин, по которым кросс-внимание более эффективно
4. Неизвестная обобщающая способность: Проверено только на сегментации опухолей печени, производительность на других задачах неизвестна

Влияние

1. Методологический вклад: Предоставляет новые идеи для проектирования архитектур 3D сегментации медицинских изображений
2. Мультимодальное слияние: Предлагает эффективное решение для слияния мультимодальной информации в медицинском анализе изображений
3. Практическая ценность: Методы, учитывающие клиническую информацию, лучше соответствуют реальным медицинским потребностям

Применимые сценарии

1. 3D сегментация медицинских изображений: Особенно в сценариях, требующих объединения клинической информации
2. Мультимодальный медицинский анализ: Задачи, объединяющие изображения и структурированные данные
3. Сегментация с дисбалансом классов: Задачи сегментации, где фон доминирует

Список литературы

Статья цитирует важные работы в этой области, включая:

• 3D U-Net: Основополагающая работа по сегментации 3D медицинских изображений
• nnU-Net: Автоматизированный фреймворк для сегментации медицинских изображений
• SwinUNETR: Применение Swin Transformer в сегментации медицинских изображений
• MedNeXt: Метод сегментации медицинских изображений на основе ConvNeXt

Общая оценка: Это работа с определённой инновационностью в области 3D сегментации медицинских изображений, особенно в аспектах мультимодального слияния и проектирования архитектур. Хотя повышение производительности относительно ограничено, практичность метода и его вклад в область заслуживают признания.