Устойчивость обучения GAN (Копаем глубже)12.04.2019 09:47

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

Введение

На процесс обучения GAN можно смотреть как на процесс минимизации нижней грани дивергенции или как на динамическую систему состоящую из дискриминатора и генератора. Первый подход очень хорошо проработан, если кому интересно может найти очень много интересных работ на эту тему, приведу лишь самые значительные на мой взгляд: f-GAN, WGAN. С другой стороны The Numerics of GANs, Lars Mescheder et al, 2017 показали что на GAN нельзя смотреть только с позиции минимизации дивергенции (хотя это и является основной целью), но еще нужно учитывать свойства динамической системы. Наша цель объединить эти два подхода и посмотреть какими же теоретическими свойствами будут обладать траектории такой динамической системы.

Далее мы будем рассматривать достаточно общую форму функционала GAN (GAN objective).

$I(\theta,\phi)=\ \int{p_d(x)f_D(D_\theta(x))dx\ +\ \int{p(\epsilon)f_G(D_\theta(G_\phi(\epsilon))d\epsilon}}$

где $\theta$ — параметры дискриминатора, $\phi$ — параметры генератора. А также предполагаем что генератор и дискриминатор дважды дифференцируемые функции.

Процесс обучения математически можно сформулировать как:

$\underset{\phi}{min}\{\underset{\theta}{max}I(\theta,\phi)\}$

Точка покоя

Несмотря на огромное количество работ по GAN вопрос о существовании точки покоя остается открытым. Сейчас мы попробуем ответить на этот вопрос.

Предположим, (Предположение 1) что для любого генератора существует дискриминатор максимизирующий

$I(\theta, \phi)$

Т.е функция

$\theta_{opt} = \underset{\theta}{argmax}\{I(\theta,\phi) \}$

определена при любых значениях $\phi$ . Хотя возможно и многозначная — т.е существует множество различных оптимальных дискриминаторов. Мы будем называть все множество таких оптимальных дискриминаторов «репараметризацией», а множество $\{\theta_{opt}\}$ множеством репараметризации дискриминатора. Сделаем небольшое уточнение — под максимумом мы понимаем не только глобальный максимум, а любой локальный.

Математически $\theta_{opt}$ определяется системой уравнений:

$\nabla_\theta I(\theta_{opt}, \phi)=0$

$\nabla_\theta^2 I(\theta,\phi) \preceq 0$

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

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

Математически это требование можно сформулировать так: система уравнений

$\nabla_\theta^2 I(\theta_{opt}, \phi) d\theta = -\nabla_{\theta\phi} I(\theta_{opt}, \phi) d\phi \;\;\;\; (1)$

имеет решения $d\theta$ при любых $d\phi$ .
Это условие легко получается из разложения $\nabla_\theta I(\theta_{opt},\phi)=0$ в ряд Тейлора.

Сейчас мы покажем, что при выполнении предположения 2 градиент генератора $\nabla_\phi I(\theta_{opt},\phi)$ не меняется если мы двигаемся вдоль множества репараметризации дискриминатора. Двигаться вдоль множества репараметризации математически означает что $\delta\theta$ принадлежит ядру $\nabla_\theta^2 I(\theta,\phi)$ т.е:

$\nabla_\theta I(\theta_{opt}, \phi) = \nabla_\theta I(\theta_{opt}+d\theta_{opt}, \phi) = 0$

Разложим в ряд Тейлора:

$\nabla_\phi I(\theta_{opt}+d\theta_{opt}, \phi) = \nabla_\phi I(\theta_{opt}, \phi) + \nabla_{\phi\theta}I(\theta,\phi) d\theta_{opt}$

Откуда, для того чтобы градиент не менялся, необходимо чтобы $d\theta_{opt}$ принадлежал ядру $\nabla_{\phi\theta} I(\theta,\phi)$ . Но если $d\theta_{opt}$ принадлежит ядру $\nabla_\theta^2 I(\theta,\phi)$ и выполняется (1), то $d\theta_{opt}$ принадлежит ядру $\nabla_{\phi\theta} I(\theta,\phi)$ . Это легко показать умножив (1) на $d\theta_{opt}$ и учитывая что $d\phi$ может быть любым:

$d\theta_{opt}^T \nabla_\theta^2 I(\theta_{opt}, \phi) d\theta =-d\theta_{opt}^T \nabla_{\theta\phi} I(\theta_{opt}, \phi) d\phi=0$

Только что мы доказали замечательный факт, что градиент $\nabla_\phi \underset{\theta}{max}\{ I(\theta,\phi) \}$ независит от того какой выбран $\theta$ из множества репараметризации дискриминатора, а значит $\underset{\theta}{max}I(\theta,\phi)$ является дифференцируемой функцией от $\phi$ и мы можем ее минимизировать используя градиентные методы оптимизации.

Что же можно сказать о минимуме этой функции? К сожалению, сделать каких либо самых общих выводов пока не получается. Нам приходится сделать еще одно предположение (Предположение 3) о свойствах дискриминатора, а именно $\underset{\theta}{max}I(\theta,\phi) \geq 0$ . Из мат-анализа мы знаем что непрерывная функция ограниченная снизу либо достигает своего минимального значения, либо существуют монотонно убывающие последовательности точек (т.е функция достигает своего минимума на бесконечности).

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

«Оптимальная траектория»

Мы только что выяснили, что при если дискриминатор удовлетворяет достаточно «мягким» предположениям, то либо точка покоя существует, либо она распололожена на бесконечности и мы можем ассимптотически к ней приближаться.

Для этого предположим, что $\theta$ и $\phi$ текущие параметры, причем дискриминатор натренирован до сходимости. Теперь мы немножко меняем генератор ( $\phi_{k+1} = \phi_k + d\phi$ ), так что $\underset{\theta}{argmax}\{I(\theta,\phi) \}$ уменьшается — например делаем один шаг градиентного спуска. Как мы должны обновить дискриминатор? Ответ дает формула (1):

$d\theta = -\nabla_\theta^2 I(\theta, \phi)^ \dagger \nabla_{\theta\phi} I(\theta, \phi) d\phi \; \; \; \; (2)$

где $\nabla_\theta^2 I(\theta, \phi)^ \dagger$ — псевдообратная матрица.
Если быть математически строгим — то формула (2) определяет касательную гиперплоскость к поверхности в пространстве параметров на которой «живут» оптимальные дискриминаторы. Будем называть эту поверхность «оптимальной».

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

$d\theta = -\nabla_\theta^2 I(\theta, \phi)^ \dagger [\nabla_\theta I(\theta,\phi) + \nabla_{\theta\phi} I(\theta, \phi) d\phi] \; \; \; \; (3)$

Что же мы только что сделали? Во-первых заметим, что на оптимальной поверхности $\nabla_\theta I(\theta, \phi)=0$ , т.е. если если $\theta$ был натренированный до сходимости дискриминатор, то мы не поменяли траекторию. С другой стороны если мы начали где то возле «оптимальной» точки, то добавочный член будет «тянуть» нас на «оптимальную» траекторию. Т.е мы сделали «оптимальную» траекторию устойчивой к шуму.

К сожалению, математически это доказать у меня не получается. Но это нам и не понадобится. Давайте присмотримся более внимательно к формуле (3), ее можно переписать в виде:

$d\theta = -\nabla_\theta^2 I(\theta, \phi)^ \dagger \nabla_\theta I(\theta, \phi+d\phi)$

Ничего не напоминает? Выглядит почти как метод Ньютона. Это говорит о том что поочередные (в англ. язычной литературе используют alternating) обновления параметров генератора и дискриминатора, где дискриминатор обновляется методом Ньютона (на самом деле нам не обязательно делать полный Ньютоновский шаг, а можно сделать небольшой шажок в направлении указываемом методом Ньютона), а генератор вообще говоря может делать шаг в любом направлении, главное чтобы уменьшалась $\underset{\theta}{argmax}\{I(\theta,\phi) \}$ , является хорошим приближением движения вдоль оптимальной траектории.

Кстати, возможно именно в этом кроется такая успешность Adam’a в обучении GAN. Ведь хорошо известно что Adam пытается аппроксимировать Гессиан и использует эту информацию для градиентного шага (аппроксимирует метод Ньютона).

Устойчивость траекторий

Для анализа устойчивости траекторий мы будем использовать такие же инструменты как использовали парни в Which Training Methods for GANs do actually Converge?. У кого есть время рекомендую хорошо разобраться в этой работе, она того стоит.

Я считаю, что самый главный недостаток этой статьи — авторы анализируют устойчивость системы в точке покоя, которой просто не существует для большинства реальных GAN (для игрушечных 2D или 3D примеров конечно такое возможно). Я имею ввиду не вообще точки покоя не существует, а точки в которой генератор имеет такое же распределение как данные, а дискриминатор тождественный 0. Поэтому мы не будем анализировать устойчивость системы в точке покоя, мы пойдем немножко другим путем — мы попытаемся проанализировать устойчивость траекторий системы. Хотя этот анализ можно легко перенести и на точку покоя, ведь точка покоя это тоже траектория.

В качестве исследуемых траекторий выберем траектории получающиеся при обучении градиентным спуском. Почему это важно? Во-первых ранее мы выяснили что же представляет из себя «оптимальная» траектория и обнаружили что она может быть хорошо аппроксимирована обновлениями дискриминатора в направлении указываемом методом Ньютона. А метод Ньютона самый что ни на есть градиентный метод, использующий информацию от вторых производных для того чтобы скорректировать направление. Ну, а главная причина, почему мы делаем именно это очень проста — потому что можем. Любое отклонение от этого приводит к таким сумасшедшим формулам, что понять что-либо там просто невозможно.

Стандартный способ анализа устойчивости траекторий — анализ линеаризованных моделей. Аналогично авторам Which Training Methods for GANs do actually Converge? мы будем исследовать Якобиан траекторий:

$A = \begin{bmatrix} \nabla_\theta^2 I(\theta, \phi) & \nabla_{\theta\phi}^2 I(\theta, \phi)\\ -\nabla_{\theta\phi}^2 I(\theta, \phi)^T & -\nabla_\phi^2 I(\theta, \phi) \end{bmatrix}$

$J = I + \eta A$

где $I$ — единичная матрица.

Мы предположим что дискриминатор, хоть и не оптимальный, но достаточно близок к оптимальному чтобы выполнялось $\nabla_\theta^2 I(\theta,\phi) \preceq 0$ . А также что регуляризатор (я уже упоминал что в реальных GAN дискриминатор не имеет оптимальной точки и что мы должны создать эту точку с помощью регуляризатора) не завист от генератора. Забегу немного вперед — самый лучший регуляризатор который мы нашли на практике зависит, но для наших теоретических исследований это предположение не так уж сильно нас ограничивает.

Для дальнейшего нам понадобится один элементарный факт о матрицах: если $\nabla_\theta^2 I(\theta,\phi) \preceq 0$ и $\nabla_\phi^2 I(\theta,\phi) \succeq 0$ , то существует такое $$\eta > 0$» /> что спектральная норма матрицы <img src=$ будет меньше либо равна 1.
Доказательство элементарно — запишем $J^T J$ :

$J^T J=[I+\eta A]^T [I + \eta A] = I + \eta[A+A^T + \eta (A^T A)]$

но:

$A+A^T = \begin{bmatrix} 2 \nabla_\theta^2 I(\theta, \phi) & 0 \\ 0 & -2 \nabla_\phi^2 I(\theta, \phi)\end{bmatrix}$

Отсюда в силу непрерывности собственных чисел, очевидно что можно подобрать такое $$\eta > 0$» /> что собственные числа <img src=$ будут меньше либо равны 1. Отсюда следует нужный нам факт что $\left \| J \right \| \leqslant 1$ .

Мы знаем что на нашей траектории $\nabla_\theta^2 I(\theta,\phi) \preceq 0$ выполняется (поскольку траектория «оптимальная») и если бы было $\nabla_\phi^2 I(\theta,\phi) \succeq 0$ то любая бы траектория (в том числе точка покоя «equilibrium point») была устойчивой — т.е при небольшом отклонении от траектории система стремилась бы на нее вернутся. Это очевидно потому что любую дискретную траекторию можно представить как произведение якобианов, а произведение матриц с нормой меньше либо равной 1 не может иметь нормы больше 1.

Но огромное множество фактов говорят что для GAN это не всегда так. Яркие пример — мод коллапс, который к тому же происходит не всегда и не для всех моделей. Я утверждаю что сложность обучения GAN связана с двумя фактами: не существует оптимального дискриминатора (лечится регуляризатором) и якобиан генератора не является неотрицательно определенной матрицей. Сейчас мы попробуем понять как если не полностью избавиться от второй проблемы, то хотя бы минимизировать ее влияние на устойчивость.

Для этого выпишем якобиан генератора. Он будет состоять из трех слагаемых:

$J_1 = \int{p(\epsilon)f_G''(D_\theta(G_\phi(\epsilon))) [\nabla_\phi G_\phi(\epsilon) \nabla_x D_\theta(G_\phi(\epsilon))] [\nabla_\phi G_\phi(\epsilon) \nabla_x D_\theta(G_\phi(\epsilon))]^T d\epsilon}$

$J_2 = \int{p(\epsilon)f_G'(D_\theta(G_\phi(\epsilon))) \nabla_\phi^2 G_\phi(\epsilon) \nabla_x D_\theta(G_\phi(\epsilon)) d\epsilon}$

$J_3 = \int{p(\epsilon)f_G'(D_\theta(G_\phi(\epsilon))) \nabla_\phi G_\phi(\epsilon)^T \nabla_x^2 D_\theta(G_\phi(\epsilon)) \nabla_\phi G_\phi(\epsilon) d\epsilon}$

Поскольку выражения для якобиана очень сложные и не получается увидеть какой то структуры все наши действия будут направлены на то чтобы эти матрицы были как можно ближе к 0.

$J_1$ видно что это слагаемое очень легко обнулить — достаточно либо выбрать функцию $f_G$ такую чтобы $f_G''(x)=0$ либо чтобы «рабочая точка» была в области где вторая производная близка к 0. Кстати, выбор функции соответствует WGAN. Связано ли с этим что WGAN известна своей стабильностью?

На два других слагаемых можно влиять только архитектурой нейронных сетей. Например, они обнуляются в случае линейных генератора и дискриминатора (вторые производные нулевые). Кстати, на ум приходят ReLU активации — к сожалению этот анализ не применим для таких сетей, потому что нарушается предположение о дифференцируемости. Но из практики известно что действительно GAN с нейросетями с кусочно линейными активациями обучаются гораздо стабильнее, чем например с гиперболическими тангенсами. Связано ли это с якобианом или есть другая более важная причина остается гадать, но в будущем я планирую немного исследовать этот вопрос.

Еще можно сделать вывод что GAN «предпочитает» широкие сети. И действительно мы знаем что GAN с глубокими сетями более часто коллапсится.

Заключение

В этой статье я попробовал дать ответы на долгоиграющие вопросы в теории GAN. Мы с вами увидели что для того чтобы существовала точка покоя, достаточно чтобы дискриминатор обладал достаточно простыми свойствами. Так же мы увидели что вероятностная (минимизация дивергенции) трактовка GAN имеет все права на жизнь и что достаточно простым алгоритмом мы можем очень хорошо аппроксимировать «оптимальные» траектории. А также что алгоритм Adam уже это делает и в теории должен стабилизировать обучение GAN, что кстати было отмечено во многих работах.

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

Что осталось вне статьи? Практика. В следующей статье я опишу регуляризатор, который я использую для того чтобы обеспечить точку покоя дискриминатору и еще кучу трюков улучшающих стабильность процесса обучения GAN (кстати некоторые из них мотивированы этим анализом).

И как пример глубокая сеть (50 слоев генератор и 50 слоев дискриминатор) натренированная на LSUN bedroom датасете (всего 100 тыс. итераций и мы не использовали усреднение генератора)