PixelRNN и PixelCNN

Рисунок 1. Пример использования PixelRNN/PixelCNN сетей

PixelRNN/PixelCNN - алгоритмы машинного обучения, входящие в семейство авторегрессивных моделей. Используются для генерации и дополнения изображений. Алгоритмы были представлены в 2016 году компанией DeepMind и являются предшественниками алгоритма WaveNet, который используется в голосовом помощнике Google.

Основным преимуществом PixelRNN/PixelCNN является уменьшение времени обучения, по сравнению с наивными способами попиксельной генерации изображений.

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

Пусть дано черно-белое изображение [math]X[/math] размером [math]N\times N[/math]. Построчно преобразуем картинку в вектор [math]V_X = \{x_1, x_2, \dots, x_{N^2} \}[/math], соединяя конец текущей строки с началом следующей. В таком представлении изображения можно предположить, что значение любого пикселя [math]x_i\in V_X[/math] может зависеть от значений предыдущих пикселей [math]x_j, j = 1,2,\dots i-1[/math].

Тогда значение пикселя [math]x_i\in V_X[/math] можно выразить через условную вероятность [math]p(x_i|x_1, x_2, \dots x_{i-1})[/math], и, используя цепное правило для вероятностей, оценка совместного распределения всех пикселей будет записываться в следующем виде: [math]p(X)=\prod_{i=1}^{N^2}p(x_i|x_1, x_2, \dots x_{i-1})[/math].

Задача алгоритма - восстановить данное распределение. Учитывая тот факт, что любой пиксель принимает значение [math]0\lt =x_i\lt =255[/math], необходимо восстановить лишь дискретное распределение.

Идея

Т.к. утверждается, что значение текущего пикселя зависит от значений предыдущего, то уместно использовать RNN, а точнее LSTM. В ранних работах уже использовался данный подход, и вычисление скрытого состояния происходило следующим образом: [math]h_{i,j}=f(h_{i-1,j}, h_{i,j-1}, x_{i,j})[/math], т.е. для того, чтобы вычислить текущее скрытое состояние, нужно было подсчитать все предыдущие, что занимает достаточно много времени.

Авторы алгоритма модернизировали LSTM в RowLSTM и Diagonal BiLSTM таким образом, чтобы стало возможным распараллеливание вычислений, что в итоге положительно сказывается на времени обучения модели.

RowLSTM

Рисунок 2. Визуализация работы модификаций LSTM. Снизу кружками обозначены пиксели, сверху - состояния на каждом пикселе. Синим обозначено то, что влияет на текущее скрытое состояние. Пустые кружки не принимают участие в вычислениях для данного скрытого состояния

В данной модификации LSTM предлагается рассчитывать скрытое состояние следующим образом: [math]h_{i,j}=f(h_{i-1,j-1}, h_{i-1,j}, h_{i-1,j+1}, x_{i,j})[/math].

Как видно из формулы и Рисунка 2, значение текущего скрытого состояния не зависит от предыдущего слева, а зависит от предыдущих сверху, которые можно параллельно рассчитать.

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

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

Diagonal BiLSTM

Рисунок 3. Операция сдвига в Diagonal BiLSTM. Параллелизация происходит по диагоналям.

В данной версии скрытое состояние считается таким же образом, как и в наивном подходе: [math]h_{i,j}=f(h_{i-1,j}, h_{i,j-1}, x_{i,j})[/math], но при этом есть хитрость в самом вычислении. Построчно сдвинем строки вправо на один пиксель относительно предыдущей, а затем вычислим скрытые состояния в каждом столбце, как показано на рисунке 3.

Данная версия позволяет учитывать контекст более качественно, но при этом занимает больше времени, чем RowLSTM.

PixelCNN

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

Архитектура

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

Маскированные сверточные слои

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

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

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

Уменьшение размерности

Рисунок 4. Блоки уменьшения размерности. Слева - блок для PixelCNN, справа - PixelRNN.

На вход в любой их указанных выше алгоритмов (PixelCNN, Row LSTM, Diagonal BiLSTM) подается большое количество объектов. Поэтому внутри каждого из них сначала происходит уменьшение их количества в 2 раза, а затем обратное увеличение в 2 раза. Структура алгоритма с учетом уменьшения размерности показана на рисунке 4.

Внутреннее устройства LSTM

Модификация Row LSTM, Diagonal BiLSTM одинакова с архитектурной точки зрения. Отличия Diagonal BiLSTM только в том, что добавляется операция сдвига в начале и возврат к исходной структуре изображения в конце.

Структура LSTM блока: 1. MaskB слой input-to-state [math]K_{is}[/math] учитывает контекст из входа. 2. Сверточный слой state-to-state [math]K_{ss}[/math] учитывает контекст из предыдущих скрытых слоев.

Используя эти два сверточных слоя формально вычисление LSTM блока можно записать следующим образом:

[math] [o_i, f_i, i_i, g_i] = \sigma (K_{ss}\circledast h_{i-1} + K_{is}\circledast x_{i}) \\ c_i=f_i\odot c_{i-1} + i_i\odot g_i\\ h_i = o_i\odot tanh(c_i) [/math]

где [math]\sigma[/math] - функция активации, [math]\circledast[/math] - операция свертки, [math]\odot[/math] - поэлементное умножение.

Архитектура PixelRNN

1. MaskA размером [math]7x7[/math]
2. Блоки уменьшения размеренности с Row LSTM блоком, в котором [math]K_{is}[/math] имеет размер [math]3x1[/math], [math]K_{ss}[/math] - [math]3x2[/math]. Для Diagonal BiLSTM [math]K_{is}[/math] имеет размер [math]1x1[/math], [math]K_{ss}[/math] - [math]1x2[/math]. Количество блоков варьируется.
3. ReLU активация
4. Сверточный слой размером [math]1x1[/math]
5. Softmax слой

Архитектура PixelCNN

1. MaskA размером [math]7x7[/math]
2. Блоки уменьшения размеренности для PixelCNN.
3. ReLU активация
4. Сверточный слой размером [math]1x1[/math]
5. Softmax слой