Как мы выиграли соревнование CLEF 2024 по генерации медицинских снимков

Всем привет! Меня зовут Михаил Чайчук, я учусь в магистратуре Вышки на ФКН, где также являюсь исследователем в НУЛ моделей и методов вычислительной прагматики. А недавно я пришел работать в AIRI на должность инженера-исследователя в команду Прикладное NLP, которой руководит Елена Тутубалина. Вместе с ней мы приняли участие в соревновании ImageCLEFmed MEDVQA-GI 2024 по генерации медицинских картинок, которое проводилось в рамках конференции CLEF 2024. 

Наш результат оказался лучшим среди 27 зарегистрированных команд! Мы уже опубликовали статью в сборнике соревнования, здесь же я расскажу, как нам удалось добиться победы.

ДИСКЛЕЙМЕР

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

Приятного аппетитачтения!

Постановка задачи

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

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

Одним из решений данной проблемы как раз и является разработка и использование генеративных нейронных сетей для создания обучающих данных. Чтобы мотивировать исследователей поработать в этом направлении, в рамках конференции CLEF 2024 было организовано соревнование ImageCLEFmed MEDVQA?GI.

CLEF (Conference and Labs of the Evaluation Forum) — это европейская конференция по искусственному интеллекту, которая существует уже 25 лет. Ее основная цель — стимулировать исследования в различных актуальных областях ИИ. Для этого каждый год перед конференцией организуется несколько соревнований, в рамках которых исследователи решают разнообразные задачи, связанные с разработкой и оценкой качества моделей. Участники предоставляют свои решения и статьи, описывающие процесс их получения, которые затем презентуются на конференци, а само участие открытое и бесплатное для всех. MEDVQA?GI — одно из таких соревнований.

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

Участникам был предоставлен уникальный набор изображений колоноскопии, который ранее не использовался для решения похожих задач. Основной метрикой качества стала метрика FID (Fr?chet inception distance). В остальном участникам была предоставлена полная свобода в выборе инструментов и подходов для решения задачи.

Данные

Для разработки модели нам предоставили два набора данных: обучающий и тестовый. Обучающая выборка содержала 20 241 пару «текст?картинка», где тексты были примерами запросов, которые наша модель должна была научиться обрабатывать. В тестовом датасете содержались только тексты, которые нужно было использовать для генерации набора изображений финальной моделью.

Большой проблемой этого датасета была нехватка разнообразия в данных: картинок было всего 2000, а из 20 241 текста только 483 были уникальными. То есть каждой картинке соответствовало в среднем около 10 текстов. Из 5000 текстов в тестовой выборке только 260 были уникальными, причем все они встречались также и в обучающем датасете.

Очевидно, что это достаточно маленький и не очень разнообразный датасет, пригодность которого для обучения качественной модели выглядела сомнительной. Как мы попытались справиться с этим напишу чуть дальше.

Наше решение

Мы начали наше исследование с предположения, что предобученные большие генеративные модели смогут наилучшим образом справиться с созданием нужных изображений. Эти модели могут создавать отличные изображения благодаря обучению на больших и разнообразных наборах данных, что помогает им улавливать сложные структуры.

Но в своем изначальном виде эти модели не могут создать специализированные медицинские изображения, так как они ориентированы на общие задачи и недостаточно адаптированы к узким областям. Вот примеры изображений, которые генерировались некоторыми популярными моделями по запросам из нашего набора данных:

Поэтому мы решили пофайнтьюнить какую?нибудь большую модель с помощью LoRa, которая сейчас является наиболее популярным методом для этого.

Ребята из команды Kandinsky предложили использовать их модель, поскольку она показывает высокое качество генерации и входит в топ лучших в мире, при этом ее веса есть в открытом доступе на HuggingFace, что значительно облегчает процесс работы с моделью.

Kandinsky 2.2

В первую очередь мы опробовали модель Kandinsky 2.2, так как она показывает хорошее качество, требуя при этом меньших затрат на обучение, чем более новая модель Kandinsky 3.0.

Архитектура Kandinsky 2.2 состоит из следующих частей:

1. Image Prior модель — DiffusionMapping модель, которая позволяет генерировать визуальный эмбеддинг CLIP по текстовому промту или текстовому эмбеддингу CLIP, одновременно оставаясь в парадигме латентного визуального пространства. В качестве модели CLIP используется CLIP?ViT?G.

2. Image Decoder — диффузионная модель U?Net для непосредственной генерации изображений.

3. Sber?MoVQGAN — модифицированная разработчиками версия VQGAN, продемонстрировавшая хорошее качество в экспериментах.

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

Подробнее про Kandinsky 2.2 можно почитать в статье разработчиков Kandinsky 2.2 — новый шаг в направлении фотореализма.

Основной сложностью при дообучении модели Kandinsky 2.2 стала необходимость параллельного файнтьюнинга обеих ее частей для достижения наилучших результатов и корректного подсчета метрик. На рисунке ниже представлен пример генераций, полученных при дообучении только одной из двух частей модели.

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

На первом этапе была обучена модель Decoder с рангом LoRa 32 в течение 40 эпох, при этом метрики не подсчитывались. На втором этапе, используя уже дообученную модель Decoder, проводились эксперименты с моделью Prior. Основным параметром, подбиравшимся в процессе экспериментов, был ранг LoRa. Мы перебирали его, начиная с 4, затем от 8 до 64 с шагом 8. Дополнительно был рассмотрен ранг 128. На графике ниже отражены результаты указанных экспериментов. При оценке качества моделей мы опирались в основном на метрику FID.

В целом наблюдается постоянное улучшение метрик с увеличением ранга. Однако изменение метрики FID при переходе от ранга 64 к рангу 128 оказалось довольно незначительным. Это, учитывая большую разницу между данными значениями, было воспринято нами, как сигнал о том, что дальнейшее увеличение ранга не даст существенных результатов.

Следующим шагом стало дообучение модели Decoder с использованием Prior?модели, с добавлением LoRa с рангом 128. Ранг LoRa для модели Decoder мы подбирали аналогичным образом. Результаты представлены на графике:

Влияние значения ранга на итоговое качество изображений в данном случае оказалось менее выраженным. На основании полученных данных можно заключить, что наибольшее влияние на качество генерируемых изображений оказывает именно Prior?модель.

Kandinsky 3.0

После достаточно успешного начала с Kandinsky 2.2 логичным продолжением экспериментов стало дообучение более новой версии данной модели — Kandinsky 3.0, которая, по словам разработчиков, позволет добиваться лучших результатов при файнтюнинге с использованием LoRa.

В Kandinsky 3.0 разработчики отказались от двухстадийной генерации, которая использовалась в Kandinsky 2.2. Вместо этого они решили использовать более классический подход с передачей в модель непосредственно закодированных текстов. Этому способствовало появление новых больших языковых моделей, которые понимают текст гораздо лучше, чем текстовый энкодер CLIP. По итогу новая архитектура состоит из трех частей:

• FLAN?UL2 — большая языковая модель, основанная на архитектуре T5. В Kandinsky 3.0 используется только Encoder данной модели.

• U?Net с модифицированной архитектурой, состоящей в основном из блоков BigGAN?deep, что позволяет увеличить глубину архитектуры в два раза, при этом сохраняя неизменным общее количество параметров.

• Sber?MoVQGAN — такой же декодер, как в версии 2.2.

Такие изменения в архитектуре позволили значительно упростить процесс обучения и дообучения, так как теперь обучать необходимо только непосредственно U?Net, а все остальные модели используются в замороженном виде. Подробности можно прочитать в статье Kandinsky 3.0 — новая модель генерации изображений по тексту.

Как и в случае с Kandinsky 2.2, в первую очередь мы провели эксперименты с различными значениями ранга для LoRa при обучении. На рисунке ниже представлены примеры полученных изображений. Как можно увидеть, им явно не хватает фотореалистичности. Изображения обладают необычной текстурой, напоминающей рисунки на бумаге или созданные в 3D?редакторе. Это негативно отразилось на значениях метрики FID, которые варьировались в диапазоне 90–110.

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

MSDM

В ходе изысканий мы обнаружили статью Diffusion Probabilistic Models beat GANs on Medical Images. В ней авторы тоже генерируют медицинские картинки и для этого с нуля обучают небольшую диффузионную модель, которую они назвали Medfusion. Их модель показала хорошее качество, при том, что она значительно меньше предобученных моделей (всего в районе 400 миллионов параметров).

Мы решили попробовать повторить их успех. Однако их модель была не приспособлена под text?to?image генерацию, а генерила картинки только по меткам классов. Поэтому мы адаптировали ее под нашу задачу, предложив свою модификацию их архитектуры, которую назвали MSDM (Medical Synthesis with Diffusion Models).

В плане архитектуры эта модель в целом похожа на Kandinsky и также идейно основана на Stable Diffusion. Наша модификация заключалась в основном в добавлении Cross?Attention блоков в слои для того, чтобы учитывать контекст из промпта.

Эта модель сразу продемонстрировала заметно лучшее, по сравнению с Kandinsky, значение FID (~35), однако именно при ее обучении сильнее всего проявились недостатки датасета: получившимся изображениям заметно не хватало разнообразия. Более того, модель имела сильную склонность просто генерировать картинки из трейна.

Чтобы решить эту проблему мы попытались аугментировать данные. Это было непросто сделать, поскольку в данном датасете большое значение имеют цвет и расположение объектов на картинках. Следовательно, применять наиболее популярные аугментации, такие, как поворот или изменение цвета, было бы некорректно. Поэтому мы подошли к делу с другой стороны и постарались максимально разнообразить тексты, чтобы модель меньше запоминала конкретные промпты. Для этого мы с помощью GPT Turbo 3.5 сгенерировали перефразы для каждого промпта из датасета. Таким образом нам удалось увеличить число уникальных текстов с ~500 до ~11 000.

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

В дополнение к улучшенным текстам мы увеличили значения dropout в модели, а также добавили context?dropout и self?conditioning. Все вместе это помогло улучшить ситуацию, и генерируемые моделью изображения сразу стали заметно более разнообразными.

Результаты

В качестве итогового решения мы отправили три набора изображений от следующих моделей: модель Kandinsky 2.2 + LoRa c рангом 128 и две версии нашей модели — одна без добавления перефразированных текстов, другая с ними.

Организаторы протестировали их с использованием нескольких скрытых от нас датасетов и прислали результаты, которые можно посмотреть в таблице:

Как можно увидеть, значения метрик на тестах значительно отличаются от тех, которые мы наблюдали в ходе экспериментов. Также интересно, что модель Kandinsky 2.2 на тестовой выборке показала лучшие результаты среди трех моделей, в то время как во время экспериментов ее результаты были хуже.

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

Для себя мы сделали следующие выводы.

Во?первых, оба опробованных нами подхода вполне имеют место быть и хорошо решают нашу задачу. Модель с 460 миллионами параметров (MSDM) продемонстрировала не сильно худшие значения метрик по сравнению с моделью в 4.6 миллиарда параметров (Kandinsky 2.2). Однако даже несмотря на то, что данная модель в целом меньше, на ее обучение с нуля всё равно уходит больше времени и ресурсов, чем на дообучение больших моделей с LoRa, так что мы всё же считаем дообучение более эффективным и перспективным методом решения подобных задач. Во?вторых, аугментации данных действительно помогают значительно улучшить результаты — можно и нужно искать разные способы разнообразить свои данные. Наш подход с перефразами текстов показал себя весьма эффективно.

В дальнейшем мы планируем протестировать наш подход на других датасетах и моделях, чтобы детальнее изучить влияние параметров LoRa на итоговый результат.

Мы благодарим команду разработчиков Kandinsky за их труд по развитию open source, в особенности, @VArkhipkin за его советы по обучению и настройке гиперпараметров 3й версии модели.