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. Они необходимы для сокрытия от алгоритма лишней информации и учета контекста — чтобы ускорить обработку изображения после каждого подсчета, предлагается вместо удаления значения пикселей применять маску к изображению, что является более быстрой операцией.

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

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

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

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

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

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

Внутреннее устройство RowLSTM и 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]7\times 7[/math]
2. Блоки уменьшения размеренности с RowLSTM блоком, в котором [math]K_{is}[/math] имеет размер [math]3\times 1[/math], [math]K_{ss}[/math] — [math]3\times 2[/math]. Для Diagonal BiLSTM [math]K_{is}[/math] имеет размер [math]1\times 1[/math], [math]K_{ss}[/math] — [math]1\times 2[/math]. Количество блоков варьируется.
3. ReLU активация
4. Сверточный слой размером [math]1\times 1[/math]
5. Softmax слой

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

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

Сравнение подходов

Критерий\название	PixelCNN	PixelRNN(Row LSTM)	PixelRNN(Diagonal BiLSTM)
Время обучения	Быстрый	Средний	Медленный
Качество генерируемых изображений	Наихудшее	Средне-низкое	Средне-высокое

Примеры реализации

• PixelCNN на Pytorch
• PixelRNN на Pytorch
• PixelRNN на Keras

Источники информации

• Auto-Regressive Generative Models
• Advanced Generation Methods
• Pixel Recurrent Neural Networks
• PixelCNN