Это первый пост, написанный во время наших исследовательских сессий в NICO.LAB.
Можем ли мы создать КТ-срезы мозга? Или интерполировать между двумя разными мозгами? Или… если мы представляем изображение мозга в виде набора чисел, можем ли мы изменить размер мозга, вращение, анатомию, просто изменив эти числа?
Я задался этими вопросами, когда впервые прочитал бумагу о распутанных вариационных автоэнкодерах (β-VAE).
Согласно документу, вы можете закодировать изображение в небольшой числовой вектор таким образом, что переменная каждого вектора будет отвечать за одну независимую и интерпретируемую визуальную характеристику. Для лиц это может быть цвет кожи, возраст, пол или насыщенность изображения. Какими будут эти функции для изображений мозга? Давайте узнаем.
Вариационный автоэнкодер
Прежде чем углубиться в распутанные вариационные автоэнкодеры (β-VAE), давайте взглянем на обычные автоэнкодеры. Механизм, лежащий в основе автоэнкодеров, прост — вы обучаете сеть, сжимая данные высокой размерности (часть кодировщика) в пространство представления низкой размерности, а затем восстанавливая их обратно в исходные данные высокой размерности (часть декодера). Цель состоит в том, чтобы как можно ближе восстановить малоразмерный вектор к исходным данным. Потери, используемые для обучения такой сети, называются потерями при реконструкции — они показывают, насколько восстановленные данные отличаются от исходных данных.
Однако малоразмерные векторы (или скрытые представления) автоэнкодеров не имеют особого смысла. Если вы сгенерируете случайный вектор скрытого представления и реконструируете его в изображение, вы получите обратно случайное изображение.
Однако, если вы используете вариационные автоэнкодеры (VAE) для восстановления случайного скрытого вектора представления, восстановленные изображения будут иметь смысл. VAE можно использовать для создания изображений!
Идея VAE заключается в том, что кодировщик создает скрытое представление не в виде отдельного вектора, а в виде распределения вероятностей по векторам. Затем мы можем выбрать скрытое представление из этого распределения и восстановить изображение.
Однако «операция выборки» во время обучения не дифференцируема, и мы не можем распространять градиенты обратно. Чтобы обойти это ограничение, применяется трюк перепараметризации: кодировщик создает два вектора: вектор среднего μ и вектор логарифмической дисперсии Σ. Затем из единичного распределения Гаусса выбирается некоторый случайный вектор ε. Скрытое представление z в этом случае:
Он используется для восстановления изображения. Затем используется обратное распространение для обновления весовых векторов μ и Σ.
Чтобы гарантировать, что распределение вероятностей, создаваемое кодером, близко к единичному распределению Гаусса, расхождение KL между этими двумя распределениями добавляется к потерям при реконструкции.
Распутанный вариационный автоэнкодер
Если мы сгенерируем случайный вектор скрытого представления и изменим только один компонент и реконструируем его в изображение, мы увидим некоторые изменения. Однако в большинстве случаев эти изменения не поддаются интерпретации. Например, если мы реконструируем изображения рукописных цифр, мы увидим, что размер, поворот и форма меняются одновременно. Это означает, что скрытое представление запутано и его компоненты не являются независимыми.
В VAE потери состоят из двух частей: потери при реконструкции, которые гарантируют, что восстановленное изображение близко к исходному изображению, и KL-дивергенция, которая гарантирует, что распределение вероятностей, создаваемое кодером, близко к единичному распределению Гаусса.
Чтобы преодолеть проблему с запутанным представлением и получить более интерпретируемые реконструкции, был предложен Disentangled VAE или β-VAE.
В β-VAE расхождение KL взвешивается с помощью параметра β, чтобы обеспечить независимость компонентов скрытого представления:
Чем больше β,темболее распутанными должны быть визуальные элементы.
Детали реализации
Для реализации VAE мы использовали кодировщик с 10 сверточными слоями и декодер с 7 сверточными слоями, за каждым слоем следует транспонированная свертка. Мы также использовали пакетную нормализацию и функцию активации Leaky ReLU. Реализация подробно описана в этой публикации.
Чтобы реализовать β-VAE, мы просто добавили параметр β в часть расхождения KL в функции потерь.
Для обучения β-ВАЭ мы использовали КТА (компьютерно-томографическую ангиографию) томов головного мозга 350 пациентов. Из каждого CTA были взяты срезы, в результате чего было получено 9910 изображений. Мы уменьшили разрешение каждого изображения с 512 x 512 до 256 x 256, чтобы упростить данные.
Мы разделили данные на обучающий набор (80%) и проверочный набор (20%) и обучили 14 эпох (что было достаточно для сходимости).
Параметр размера скрытого представления
В первом эксперименте мы можем увидеть влияние скрытого размера представления на потери при проверке. Размеры скрытых векторов в этом эксперименте были [10, 20, 50, 100, 500]. Значение β = 1 (традиционный VAE).
Наихудшая производительность достигается при скрытом размере представления 500 (слишком много единиц для небольшого количества обучающих данных, которые у нас есть) и 10 (представление недостаточно мощное для представления наших данных).
Скрытые векторы в 20, 50 и 100 единиц показали одинаковую эффективность.
На графиках ниже первое число показывает размер скрытого вектора, а второе число показывает значение β. Не удивляйтесь забавным названиям экспериментов — они были сгенерированы автоматически (с использованием весов и смещений).
Ниже мы можем видеть срезы КТ головного мозга, созданные с использованием β-VAE со скрытыми размерами представления 10, 50 и 100. При слишком малом количестве скрытых единиц представления (10) сгенерированные изображения слишком круглые и простые. При слишком большом количестве единиц (500) сгенерированные изображения разрываются. Изображения, сгенерированные из 50 единиц, выглядят лучше: менее круглые и не разорванные.
Параметр β
В следующем эксперименте мы можем увидеть влияние β на потерю проверки. Скрытый размер вектора равен 50, а параметр β равен [1, 20, 50, 200]. Чем больше β, тем сильнее регуляризация и тем выше потери, как и ожидалось.
Хотя качество реконструкции ухудшается с увеличением β (это коррелирует с потерей проверки), распутывание и интерпретируемость скрытых переменных представления должны улучшаться (см. следующий эксперимент).
Визуализации распутывания
Параметр β в этом эксперименте был [1, 50, 200], размер вектора скрытого представления = 50.
Изображения были уменьшены, чтобы соответствовать сетке.
Средний столбец на рисунках ниже представляет собой реконструкцию исходного изображения среза мозга. Если мы уменьшим/увеличим первую единицу скрытого представления на небольшое значение, мы увидим некоторое изменение в реконструкции (первая строка, справа/слева от среднего столбца). Другие строки представляют остальные четыре единицы скрытого представления. Конечно, мы могли бы показать все 50 единиц, но сетка была бы слишком большой и ее было бы трудно интерпретировать.
Визуализации интерполяции
В этом эксперименте мы увидим, как интерполяция между срезами изображений мозга разных людей может быть выполнена посредством интерполяции между скрытыми представлениями этих срезов.
Для изображений ниже использовалась модель с β, равным 50, и размером скрытого представления, равным 50.
Будущая работа
Нам пришлось уменьшить разрешение изображений с 512x512 до 256x256 пикселей, чтобы упростить данные. Однако, если мы используем больше обучающих данных, имеет смысл сохранить исходный размер изображения. Возможно, в этом случае реконструированные изображения будут менее размытыми.
Также было бы интересно использовать более объективную метрику для оценки важности конкретной латентной единицы представления вместо визуального осмотра вручную. Для этой цели можно использовать метрику распутывания, предложенную в статье β-VAE.
Вывод
Сгенерированные изображения, созданные VAE и β-VAE, слишком размыты, чтобы их можно было использовать (например, для увеличения данных). Тогда все было бесполезно? Не совсем.
Тот факт, что вы можете закодировать срез мозга всего 10 (или 50) числами, впечатляет. Тот факт, что вы можете модифицировать эти числа, чтобы изменить отдельные свойства мозга и интерпретировать эти изменения, также завораживает.
Более того, было так весело проводить эти эксперименты!
Оригинальная бумага VAE:
Кингма, Дидерик П. и др. Автокодирование вариационных байесов. препринт arXiv arXiv:1312.6114 (2013 г.). https://arxiv.org/pdf/1312.6114.pdf
Бумага Бета-ВАЭ:
Хиггинс, Ирина и др. β-VAE: изучение основных визуальных концепций с ограниченной вариационной структурой. (2016). https://openreview.net/references/pdf?id=Sy2fzU9gl
VAE объяснил сообщение в блоге:
https://keitakurita.wordpress.com/2017/12/19/an-intuitive-explanation-of-variational-autoencoders/
Синтез изображений брюшной полости с помощью учебника VAE:
