Това е първата публикация, написана по време на нашите изследователски сесии в NICO.LAB.
Можем ли да генерираме CT срезове на мозък? Или интерполиране между два различни мозъка? Или... ако представим образа на мозъка като набор от числа, можем ли да променим размера на мозъка, въртенето, анатомията само чрез промяна на тези числа?
Зададох тези въпроси, когато за първи път прочетох документа за Disentangled Variation Autoencoders (β-VAE).
Според статията можете да кодирате изображение в малък цифров вектор по начин, по който променливата на всеки вектор ще бъде отговорна за една независима и интерпретируема визуална характеристика. За лицата това може да бъде цвят на кожата, възраст, пол или наситеност на изображението. Какви ще бъдат тези функции за мозъчни изображения? Нека разберем.
Вариационен автоенкодер
Преди да се потопим в Disentangled Variation Autoencoders (β-VAE), нека да разгледаме обичайните автоенкодери. Механизмът зад автоенкодерите е прост — вие обучавате мрежа чрез компресиране на високомерни данни (част от енкодера) в нискоизмерно пространство за представяне и след това го реконструирате обратно в оригинални високоизмерни данни (декодерна част). Целта е да се реконструира нискоразмерният вектор спрямо оригиналните данни възможно най-близо. Загубата, използвана за обучение на такава мрежа, се нарича загуба при реконструкция - тя показва колко различни са реконструираните данни от оригиналните данни.
Нискоразмерните вектори (или латентни представяния) на автокодери обаче нямат толкова много смисъл. Ако генерирате произволен вектор на латентно представяне и го реконструирате в изображение, ще получите обратно произволно изображение.
Въпреки това, ако използвате вариационни автоенкодери (VAE), за да реконструирате вектор на произволно латентно представяне, реконструираните изображения ще имат смисъл. VAE могат да се използват за генериране на изображения!
Идеята зад VAE е, че енкодерът произвежда латентното представяне не като единичен вектор, а като вероятностно разпределение върху вектори. След това можем да вземем проби от латентно представяне от това разпределение и да реконструираме изображението.
Въпреки това, „операцията за вземане на проби“ по време на обучение не е диференцируема и не можем да разпространяваме обратно градиентите. За да се преодолее това ограничение, се прилага трикът за повторна параметризация: енкодерът произвежда два вектора: вектор на средната стойност μ и вектор на логаритъм на вариация Σ. След това някакъв произволен вектор ε се взема от единичното разпределение на Гаус. Скрито представяне z в този случай:
Използва се за възстановяване на изображението. След това се използва обратно разпространение за актуализиране на тегловните вектори μ и Σ.
За да се гарантира, че разпределението на вероятностите, произведено от енкодера, е близко до единичното гаусово разпределение, KL разминаването между тези две разпределения се добавя към загубата при реконструкция.
Разплетен вариационен автоенкодер
Ако генерираме вектор на произволно латентно представяне и променим само един компонент и го реконструираме в изображението, ще видим някои промени. Въпреки това, в повечето случаи тези промени няма да бъдат интерпретируеми. Ако реконструираме изображения на ръкописни цифри, например, можем да видим, че и размерът, и ротацията, и формата се променят едновременно. Това означава, че латентното представяне е заплетено и неговите компоненти не са независими.
При VAE загубата се състои от две части: загуба при реконструкция, която гарантира, че реконструираното изображение е близко до оригиналното изображение и KL-дивергенция, която гарантира, че вероятностното разпределение, което енкодерът произвежда, е близко до единичното разпределение на Гаус.
За да се преодолее проблемът със заплетеното представяне и да се получат по-интерпретируеми реконструкции, беше предложен Disentangled VAE или β-VAE.
В β-VAE дивергенцията на KL е претеглена с параметър β, за да се наложи независимостта на латентните компоненти на представяне:
Колкото по-голям е β, толковатолковапо-разплетени трябва да бъдат визуалните характеристики.
Подробности за изпълнението
За внедряване на VAE използвахме енкодер с 10 конволюционни слоя и декодер със 7 конволюционни слоя, всеки слой последван от транспонирана конволюция. Използвахме също нормализиране на партиди и функция за активиране на Leaky ReLU. Изпълнението е описано подробно в този среден пост.
За да приложим β-VAE, току-що добавихме параметър β към частта за дивергенция на KL във функцията за загуба.
За да обучим β-VAE, използвахме CTA (компютърна томографска ангиография) мозъчни обеми на 350 пациенти. Бяха взети срезове от всеки CTA, което доведе до 9910 изображения. Ние намалихме семплирането на всяко изображение от 512x512 до 256x256, за да опростим данните.
Разделихме данните на набор за обучение (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 оригинална хартия:
Kingma, Diederik P., et al. „Автоматично кодиране на вариационни гнезда.“ arXiv preprint arXiv:1312.6114 (2013). https://arxiv.org/pdf/1312.6114.pdf
Бета-VAE хартия:
Хигинс, Ирина и др. „β-VAE: Изучаване на основни визуални концепции с ограничена вариационна рамка.“ (2016). https://openreview.net/references/pdf?id=Sy2fzU9gl
VAE обясни публикация в блога:
https://keitakurita.wordpress.com/2017/12/19/an-intuitive-explanation-of-variational-autoencoders/
Синтез на коремни изображения с урок по VAE:
