Диффузионные модели появились как новый метод генеративного моделирования, который демонстрирует исключительную производительность при создании синтетических данных. Примеры их впечатляющих результатов в настоящее время можно увидеть во многих моделях преобразования текста в изображение, таких как DALL-E 2, Imagen и Stable Diffusion. В этом сообщении блога мы быстро объясним основы моделей диффузии и покажем их применение для создания трехмерных медицинских изображений, таких как компьютерная томография (КТ) или магнитно-резонансная томография (МРТ).

Прежде чем мы перейдем к диффузионным моделям, давайте сначала рассмотрим другие генеративные модели, показавшие отличные результаты за последние годы.

Вариационные автоэнкодеры (VAE)

Впервые опубликованные в 2013 году Kingma и Welling, вариационные автоэнкодеры (ВАЭ) [1] показали многообещающие результаты в создании синтетических изображений путем моделирования распределения скрытых переменных. Основная интуиция, стоящая за этими моделями, заключается в том, что кодировщик используется во время обучения для извлечения скрытых функций из входных изображений. Затем эти функции передаются в декодер для восстановления изображений. Заставляя распределение скрытых переменных быть близким к априорному распределению Гаусса, и в то же время заставляя реконструированное изображение быть близким к входному изображению, модели получают хорошее приближение к основному распределению данных. Затем изображения могут быть сгенерированы путем выборки нового набора скрытых признаков из гауссовского априора, и декодер реконструирует эти признаки в реалистично выглядящие изображения.

Генеративно-состязательные сети (GAN)

Еще одна модель, получившая широкую популярность за последние пару лет благодаря впечатляющим генеративным характеристикам, — это генеративно-состязательные сети (GAN) [2]. Они состоят из двух основных компонентов: сети генератора и сети дискриминатора. В то время как генератору поручено генерировать синтетические данные, похожие на обучающие данные, дискриминатор обучается определять, является ли данный образец реальным или поддельным. Во время обучения обе сети обучаются одновременно. Учитывая обратную связь от дискриминатора, генератор соответствующим образом корректирует свои параметры до тех пор, пока сгенерированные синтетические данные не станут неотличимы от реальных данных. После обучения GAN можно использовать для создания синтетических изображений, вводя случайный шум в генератор и позволяя ему создавать новое изображение.

Рассмотрев, как работают предыдущие современные методы генеративного моделирования, давайте кратко рассмотрим функциональные возможности диффузионных моделей и то, как мы можем применить их для создания трехмерных данных.

Диффузионные модели

В отличие от двух других моделей, диффузионные модели [3] позволяют генерировать синтетические изображения, моделируя генеративный процесс как цепь Маркова по скрытым переменным. На высоком уровне модели диффузии работают, постепенно искажая изображения во время обучения с гауссовским шумом и обучая нейронную сеть восстанавливать данные, обращая этот шумовой процесс. Когда этот шумовой процесс сделан достаточно точным, что означает, что мы добавляем к изображению только небольшое количество шума на каждом шаге, задача восстановления изображения из зашумленного ввода становится выполнимой. Путем многократного добавления этого небольшого количества шума в течение достаточно большого количества шагов (это гиперпараметр, который должен быть установлен при обучении сети), исходное изображение превращается в зашумленную версию, которая распределяется в соответствии со стандартной гауссианой. Затем во время логического вывода мы можем произвести выборку из стандартной гауссовой диаграммы и проследить процесс диффузии в обратном направлении, чтобы превратить выбранный шум в реалистично выглядящее синтетическое изображение.

Моделирование трехмерных медицинских данных с помощью моделей распространения

Диффузионные модели имеют один заметный недостаток по сравнению с двумя другими методами: синтез новых изображений занимает довольно много времени. Для каждого шага обратного процесса обученная нейронная сеть должна восстанавливать добавленный крошечный бит небольшого шума. Чтобы нейронная сеть восстановила шум, допускается только небольшой вклад шума на каждом шаге, и поэтому для преобразования изображения в гауссовский шум требуется большое количество шагов. Это приводит к длительному времени вывода по сравнению с GAN и VAE. Это особенно важно при моделировании трехмерных данных, поскольку нейронная сеть, используемая в диффузионной модели, довольно велика.

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

Один из способов преодолеть эти проблемы — сначала обучить вариационный автокодер извлечению скрытых особенностей изображений в сжатом пространстве. Затем можно использовать диффузионную модель для моделирования скрытых признаков вместо всего изображения. Это позволяет ускорить вывод и снизить требования к вычислительным ресурсам, поскольку размер модели может быть уменьшен при работе в сжатом скрытом пространстве. Ромбах и др. использовали эту технику в своей работе над моделями скрытой диффузии [4] для синтеза двумерных данных. Это также метод, лежащий в основе широко популярной модели Стабильная диффузия.

В нашей работе (Medical Diffusion: https://arxiv.org/abs/2211.03364) мы придерживаемся аналогичного подхода и сначала сжимаем трехмерные объемы изображений КТ и МРТ в скрытое представление. Для создания этого скрытого представления мы используем 3D-версию VQ-GAN [5] (Vector Quantized Generative Adversarial Network) — вариант VAE, который позволяет лучше реконструировать путем квантования скрытого представления с помощью обучаемой кодовой книги и дискриминатора. потеря. Процесс распространения, который был адаптирован для поддержки трехмерных изображений, затем используется для моделирования этого скрытого пространства и, следовательно, позволяет использовать эффективный способ синтеза трехмерных медицинских изображений: изображений с помощью обученного декодера VQ-GAN.

Мы обнаружили, что этот подход может привести к реалистично выглядящим изображениям (по оценке нескольких рентгенологов), которые превосходят предыдущие базовые уровни GAN, современные достижения в создании трехмерных изображений. Кроме того, мы показываем, что модель можно использовать в условиях, когда учреждение А хочет поделиться своими данными с другим учреждением Б, не предоставляя доступа к реальным данным: используя синтетические данные учреждения А в условиях обучения с самоконтролем, мы обнаруживаем, что учреждению B нужно только аннотировать несколько изображений, чтобы достичь почти полной производительности сегментации.

Наш код общедоступен на GitHub: https://github.com/FirasGit/medicaldiffusion и может быть адаптирован для запуска модели Medical Diffusion на ваших собственных наборах данных.

Для получения дополнительной информации и понимания нашей модели медицинской диффузии, пожалуйста, прочитайте наш препринт на https://arxiv.org/abs/2211.03364.

Ссылки

  1. Кингма, Д., и Веллинг, М.. (2013). Автокодирование вариационного Байеса.
  2. Гудфеллоу И., Пуже-Абади Дж., Мирза М., Сюй Б., Вард-Фарли Д., Озаир С., Курвиль А. и Бенжио Ю. (2014). Генеративно-состязательные сети.
  3. Хо, Дж., Джайн, А., и Аббил, П.. (2020). Шумоподавление диффузионных вероятностных моделей.
  4. Ромбах Р., Блаттманн А., Лоренц Д., Эссер П. и Оммер Б.. (2021). Синтез изображений высокого разрешения с использованием моделей скрытой диффузии.
  5. Эссер П., Ромбах Р. и Оммер Б.. (2020). Укрощение трансформаторов для синтеза изображений с высоким разрешением.