31 / 33 · Авторегрессионные модели
Авторегрессионные модели: PixelRNN и PixelCNN
Плотность изображения как произведение одномерных условных распределений, каждое из которых — softmax по значениям пикселя: точное правдоподобие, биты на размерность, параллельное обучение и последовательное сэмплирование, рекуррентные (PixelRNN) и маскированные свёрточные (PixelCNN) сети, слепое пятно, Gated PixelCNN и PixelCNN++.
На этой странице 12
Семейства этого курса по-разному обходятся с невычислимым правдоподобием: VAE его оценивает снизу, GAN избегает, поток вычисляет через обратимое отображение, диффузия снова оценивает снизу. Авторегрессионные модели убирают трудность у самого источника. Они записывают плотность изображения как произведение одномерных условных распределений, каждое из которых — softmax по нескольким сотням значений и потому нормировано по построению. Правдоподобие точное, обучение сводится к кросс-энтропии, а платить приходится при сэмплировании: пиксель за пикселем.
Обозначения. Изображение — вектор из значений в фиксированном порядке, обычно построчно от левого верхнего угла, а у цветных изображений ещё и красный, зелёный, синий внутри каждого пикселя. — значения до позиции . Значения — целые числа из , если не сказано иное.
Цепное правило как модель#
Любое совместное распределение можно записать как
Это цепное правило теории вероятностей, а не допущение: оно верно для любого распределения и любого порядка координат. Ограничения модели появляются только из того, как параметризованы условные распределения. Метод максимального правдоподобия, который заметка о KL выводит из минимизации , превращается в сумму задач классификации: в каждой позиции предсказать следующее значение по предыдущим и заплатить кросс-энтропию истинного значения.
Вся трудность — в компактном представлении условных распределений. Возьмём бинарные изображения из пикселей. Таблица совместного распределения содержит свободных чисел, и столько же в сумме содержат таблицы всех условных распределений: условному распределению пикселя нужно по числу на каждую из предысторий. Полностью видимая сигмоидная сеть доверия (fully visible sigmoid belief network) заменяет каждую таблицу логистической регрессией,
с параметрами для пикселя и параметрами всего. Для изображений MNIST размера , , это 307 720 параметров вместо . Более поздние модели разделяют параметры сильнее: NADE (Larochelle and Murray, 2011) вычисляет все условные распределения одним скрытым слоем, веса которого переиспользуются между позициями, а MADE (Germain et al., 2015) получает их все за один проход автоэнкодера с масками на весах, так что выход зависит только от входов до ; те же маски строят авторегрессионные потоки. PixelRNN и PixelCNN идут дальше и применяют одну рекуррентную или свёрточную сеть во всех позициях, так что число параметров перестаёт зависеть от размера изображения.
Сэмплирование следует произведению: вытянуть , затем при данном , и так далее. Каждый шаг — метод обращения функции распределения для дискретной величины: взять и вернуть наименьшее значение, накопленная вероятность которого достигает . Авторегрессионная модель — по-прежнему шум плюс обучаемое преобразование, где шумом служат равномерных чисел.
Пиксели как дискретные значения#
PixelRNN (van den Oord, Kalchbrenner and Kavukcuoglu, 2016) считает каждое значение цвета категорией и предсказывает его softmax на 256 классов. Дискретному распределению не нужны допущения о форме: оно может быть многомодальным или скошенным, и оно назначает вероятности, а не плотности: , и деквантование (dequantization) не нужно, в отличие от непрерывных плотностей потоков.
Правдоподобие моделей изображений сообщают в битах на размерность (bits/dim):
с усреднением по тестовым изображениям. Модель, распределяющая вероятность равномерно по 256 значениям, получает ровно 8 бит на размерность. У этого числа прямой смысл: с арифметическим кодированием модель с битами на размерность сжимает изображения без потерь примерно до бит. Для CIFAR-10 , и лучший PixelRNN достигает 3.00 бит на размерность: бит, около 1.1 КиБ на изображение, или 37.5% от 24 576 бит исходного изображения.
Модель, предсказывающая меньше уровней, решает более грубую задачу. При 10 уровнях серого равномерное угадывание стоит ната, или бита на пиксель, и сравнивать такие числа с битами на размерность по 256 значениям нельзя.
Параллельное обучение, последовательное сэмплирование#
При обучении изображение известно целиком. Если сеть устроена так, что её выход в позиции зависит только от входов , один прямой проход по настоящему изображению даёт сразу все условных распределений, а функция потерь — их сумма. Подача на вход истинных предыдущих значений, а не собственных сэмплов модели, называется teacher forcing, и благодаря ей обучение так же параллельно, как у любой свёрточной сети.
Сэмплирование параллельным быть не может: вычисляется по , который нужно сначала вытянуть. Модели, которая предсказывает по одному значению, на сэмпл нужно проходов сети: 3072 для изображения CIFAR-10 и для серого изображения пикселей. Модели, которая за проход предсказывает целый пиксель, как PixelCNN++ ниже, нужен один проход на пиксель, 1024 для CIFAR-10, а три значения пикселя она вытягивает из предсказанных параметров; позиции по-прежнему идут последовательно. Простая реализация к тому же на каждом проходе пересчитывает всё изображение, хотя использует один выход. Та же асимметрия — дешёвая плотность и дорогое сэмплирование — есть у маскированного авторегрессионного потока.
PixelRNN и PixelCNN#
Две сети из van den Oord, Kalchbrenner and Kavukcuoglu (2016) различаются тем, как выход в позиции собирает информацию из всего , ничего не видя из и последующих значений.
Рекуррентные слои#
Row LSTM обрабатывает изображение строка за строкой сверху вниз и вычисляет состояние сразу для всей строки. Переход вход–состояние — свёртка вдоль строки (), а переход состояние–состояние — одномерная свёртка по предыдущей строке состояний. Каждый шаг смотрит немного влево и вправо в строке выше, поэтому контекст пикселя — треугольник над ним, и часть доступного прошлого в него не попадает.
Diagonal BiLSTM охватывает всё прошлое. Он скашивает входную карту, сдвигая каждую строку на одну позицию вправо относительно предыдущей, так что диагонали изображения становятся столбцами карты . Свёртка состояние–состояние имеет размер вдоль этих столбцов, и два LSTM сканируют в противоположных направлениях; выход прохода справа налево сдвигается на строку вниз перед сложением, чтобы не просочились будущие пиксели. В моделях до 12 слоёв с остаточными связями между ними. Рекуррентности идут последовательно по строкам или диагоналям даже при обучении, поэтому обе модели обучаются медленно.
Маскированные свёртки#
PixelCNN заменяет рекуррентности свёртками, ядра которых умножаются на маску. Для ядра с центром
так что позиция видит строки над собой и пиксели слева от себя. Первый слой использует маску A, которая исключает центр: на входе первого слоя в центре стоит значение, которое нужно предсказать. Последующие слои используют маску B, которая центр включает: после первого слоя признак в позиции вычислен только по , и передавать его дальше безопасно. PixelCNN из статьи состоит из пятнадцати слоёв: первый с маской A, затем остаточные блоки со свёртками . Обучение полностью параллельно, сэмплирование — нет.
Маски свёртки 5×5 в PixelCNN
Цветовые каналы#
В цветном изображении каждый пиксель содержит три значения в порядке R, G, B, и факторизация идёт по всем (число пикселей) значениям:
где теперь весь пиксель. Карты признаков каждого слоя делятся на три группы, по одной на цвет. Маска A не даёт красной группе видеть ни одного значения текущего пикселя, зелёной позволяет видеть красный, а синей — красный и зелёный; маска B вдобавок соединяет каждую группу с самой собой.
Простейшая альтернатива моделирует каналы независимо, например тремя моделями полутоновых изображений, по одной на канал. Тогда каждый канал сам по себе — правдоподобное изображение, но ничто не связывает три канала между собой: в сэмпле красное, зелёное и синее значения пикселя, а также границы, которые они рисуют, не обязаны согласовываться. Более мягкий вариант ставит в условие каждого канала все предыдущие пиксели всех каналов, но отбрасывает зависимость между тремя значениями текущего пикселя. Цепному правилу эта зависимость нужна, поэтому такой вариант — допущение, а не факторизация. Модель, сохраняющая эту зависимость, либо соблюдает порядок каналов внутри масок, как описано выше, либо, как PixelCNN++ ниже, предсказывает три значения пикселя совместно с простой зависимостью между ними.
Слепое пятно#
Стопка маскированных свёрток не даёт всего прошлого. Маска дотягивается на строку вверх и на столбец в каждую сторону или на столбец влево в текущей строке. После слоёв строка на выше цели покрыта от столбцов слева лишь до столбцов справа, каким бы ни было . Пиксели правее в ближайших строках сверху остаются недосягаемыми при любой глубине. В окне из строк и столбцов это слепое пятно содержит из предшествующих пикселей, 6 из 40 при ; его доля растёт к четверти — эту цифру приводят van den Oord et al. (2016b). Раздел с кодом ниже находит те же 6 пикселей, дифференцируя сеть со случайной инициализацией.
Слепое пятно PixelCNN
Gated PixelCNN#
van den Oord et al. (2016b) закрывают разрыв между PixelCNN и PixelRNN двумя изменениями. Первое устраняет слепое пятно двумя стопками свёрток. Вертикальная стопка не маскируется внутри своего окна и видит все строки выше текущей, так что её рецептивное поле растёт прямоугольником. Горизонтальная стопка видит текущую строку до текущего пикселя. После каждого слоя вертикальная стопка передаёт выход горизонтальной, но не наоборот: информация из горизонтальной стопки занесла бы в вертикальную пиксели правее или ниже цели.
Второе изменение заменяет ReLU между слоями мультипликативными блоками, аналогом гейтов LSTM:
где — свёртка, — номер слоя, — поэлементное произведение. Условие входит дополнительным слагаемым под обеими нелинейностями: для вектора , например one-hot метки класса, , а для пространственной карты произведения становятся свёртками . Условная модель — .
На CIFAR-10 Gated PixelCNN достигает 3.03 бита на размерность против 3.14 у PixelCNN и 3.00 у PixelRNN. На ImageNet он немного лучше PixelRNN: 3.83 против 3.86 при и 3.57 против 3.63 при , а качества Row LSTM достигает менее чем за половину времени обучения (60 часов на 32 GPU). Модель для ImageNet — 20 слоёв по 384 канала с фильтрами .
PixelCNN++#
Softmax на 256 классов ничего не знает о порядке своих классов: 127 и 128 для него так же не связаны, как 0 и 255, пока сеть не выучит иное; значение, не встреченное при обучении, получает почти нулевую вероятность; а выходной слой содержит чисел на пиксель. PixelCNN++ (Salimans et al., 2017) предполагает непрерывную интенсивность со смесью логистических распределений и моделирует наблюдаемое значение как , округлённое до целого:
где — логистическая сигмоида, функция распределения логистического закона. Для слагаемое заменяется на , а для слагаемое — на , так что вероятности по дают в сумме единицу, а крайние значения собирают массу всего, что лежит за ними. С одной компонентой, и , значения 126, 127, 128 и 129 получают вероятности 0.112, 0.124, 0.121 и 0.105: соседи близки автоматически. При и крайнее значение 255 получает 0.438, а 254 — 0.124, что соответствует избытку насыщенных пикселей в реальных фотографиях.
Остальные изменения:
- Целые пиксели. Сеть обусловливается на целые пиксели сверху и слева и предсказывает три канала пикселя вместе. Индикатор компоненты смеси общий для каналов, а внутри компоненты средние линейно зависят от уже вытянутых каналов: и , где выдаёт сеть. Выход — чисел на пиксель ( весов, среднее и масштаб для каждого канала и каждой компоненты и три коэффициента на компоненту), 50 при , чего, по данным статьи, достаточно.
- Понижение разрешения свёртками с шагом 2 вместо дилатированных свёрток, чтобы видеть большие области на нескольких разрешениях, и обходные связи (short-cut connections), которые проносят признаки полного разрешения мимо слоёв пониженного; без них модель не обучается.
- Dropout на остаточном пути. Без него функция потерь на обучении падает ниже 2.0 бита на субпиксель, а на тесте поднимается выше 6.0.
PixelCNN++ достигает 2.92 бита на размерность на CIFAR-10. Та же модель, обученная на изображениях с добавленным равномерным шумом, то есть с непрерывной плотностью вместо округления, даёт оценку 3.11.
Сравнение правдоподобий#
Отрицательное логарифмическое правдоподобие на тестовой части CIFAR-10 в битах на размерность (van den Oord et al., 2016a, 2016b; Salimans et al., 2017):
| Модель | Бит на размерность |
|---|---|
| Равномерное распределение | 8.00 |
| Многомерное нормальное | 4.70 |
| NICE (поток) | 4.48 |
| Real NVP (поток) | 3.49 |
| PixelCNN | 3.14 |
| Row LSTM | 3.07 |
| Gated PixelCNN | 3.03 |
| Diagonal BiLSTM (PixelRNN) | 3.00 |
| PixelCNN++ | 2.92 |
Несколько лет после 2016 года лучшие правдоподобия на этих бенчмарках держали авторегрессионные модели. Правдоподобие измеряет, насколько хорошо модель сжимает тестовую выборку, а это не то же самое, насколько хорошо выглядят её сэмплы: это отдельно измеряют заметки об оценке, и модель может выигрывать по одному и проигрывать по другому.
Сильные стороны и ограничения#
- Точное правдоподобие и простое обучение. Функция потерь — сумма кросс-энтропий без оценок снизу, без противника и без якобиана, и её можно отслеживать на отложенных данных.
- Медленное сэмплирование. Одно вычисление сети на позицию: на изображение у модели, работающей по субпикселям, и одно на пиксель ( для цветного изображения), когда сеть, как PixelCNN++, предсказывает пиксели целиком, а значения каналов затем вытягиваются из её выхода. В обоих случаях стоимость растёт с разрешением. Диффузии тоже нужно много проходов, но их число задаётся расписанием шума, а не числом пикселей.
- Нет латентного кода. Ничто не описывает изображение несколькими числами, а порядок произволен: модель должна выучить глобальную структуру из предсказания слева направо и сверху вниз. PixelVAE (Gulrajani et al., 2017) соединяет две идеи: это VAE, декодер которого — PixelCNN. Ему нужно гораздо меньше авторегрессионных слоёв, чем одному PixelCNN, а его коды компактнее, чем у обычного VAE, и всё же передают большую часть структуры изображения.
- Сильные априорные распределения. На небольшой сетке дискретных кодов авторегрессионная модель — отличное априорное распределение. Такова её роль в VQ-VAE, где PixelCNN обучается на индексах кодов, и с этого начинается следующая заметка.
Ноутбук#
Ноутбук о PixelCNN из репозитория HSE-LAMBDA Deep Generative Models обучает небольшой PixelCNN на постерах фильмов (номера ячеек считают все ячейки, начиная с 1). Сохранённых выводов в нём нет, поэтому он описывает процедуру, а не результаты.
- Ячейки 1–3: подготовка. Постеры размера уменьшаются в 3 раза до ; батчи по 32.
- Ячейки 4–5, задание 1:
bucketizeпереводит 0–255 в 10 уровней целочисленным делением на 26 с ограничением сверху числом 9. Корзины не совсем равные, как утверждает docstring: в уровнях 0–8 по 26 значений, а в уровне 9 — 22 значения 234–255. В тексте задания форма указана как batch × channels × width × height; у тензоров она batch × channels × height × width. - Ячейки 6–7, задание 2: изменение размера, перевод в серое,
ToTensor, умножение на 255 иbucketize. Произведение здесь возвращает точные целые, так что ни один пиксель не попадает не в ту корзину. - Ячейки 9–13:
CustomImageDatasetсоставляет список файлов изображений в../data/movie-genre-from-its-poster/posters/— так называется папка датасета Kaggle Movie Genre from its Poster с постерами из IMDB. Он один раз открывает каждый файл методомverify()из PIL и пропускает те, которые не удаётся прочитать, переводит каждое изображение в RGB перед преобразованиями и возвращает с каждым изображением фиктивную метку 0, чтобы работал стандартный загрузчик. Последние 100 файлов в отсортированном порядке — валидация. Валидационный загрузчик отбрасывает последний неполный батч, так что используются 96 из них. - Ячейки 15–19, задание 3:
make_maskиMaskedCNN, который на каждом прямом проходе умножаетweight.dataна маску на месте. Это работает: замаскированные веса всё равно получают градиенты, и оптимизатор сдвигает их между шагами, но следующий прямой проход снова их обнуляет — тот же эффект, что обнуление после каждого обновления в статье о PixelRNN. Умножение в прямом проходе без.data, как в разделе с кодом, оставляет их градиенты нулевыми. - Ячейки 20–21, задание 4: модель из одного слоя с маской A, одного с маской B и выхода на 10 логитов на пиксель. Два маскированных слоя видят немного: каждое предсказание дотягивается не выше чем на строк вверх, 6 из 89 строк для ядер . Модель получает индексы уровней 0–9 как float без масштабирования.
- Ячейки 23–24: Adam, кросс-энтропия по 10 уровням и 50 эпох. После каждой эпохи два валидационных постера достраиваются со строки 10 и столбца 10: область ниже и правее обнуляется и заполняется заново пиксель за пикселем, проходов сети на область, а пиксели левее столбца 10 в нижних строках остаются настоящими. Функция потерь — в натах на пиксель, равномерное угадывание даёт ; она измеряет задачу с 10 уровнями и несравнима с битами на размерность по 256 значениям.
- Ячейка 25: восемь достраиваний со строки 30 и столбца 30, проходов сети на батч.
- Ячейка 26: вопрос о цветных изображениях, ответ — в разделе Цветовые каналы выше.
Код#
PixelCNN для серых изображений с levels значениями на пиксель на PyTorch и проверка его рецептивного поля:
import torch
import torch.nn as nn
import torch.nn.functional as F
def make_mask(include_center, height, width):
mask = torch.zeros(height, width)
mask[: height // 2] = 1 # строки выше центра
mask[height // 2, : width // 2] = 1 # слева от центра в его строке
if include_center:
mask[height // 2, width // 2] = 1
return mask
class MaskedConv2d(nn.Conv2d):
def __init__(self, include_center, *args, **kwargs):
super().__init__(*args, **kwargs)
self.register_buffer('mask', make_mask(include_center, *self.kernel_size))
def forward(self, x):
return F.conv2d(x, self.weight * self.mask, self.bias, padding=self.padding)
class PixelCNN(nn.Module):
def __init__(self, levels=10, hidden=64, layers=4, k=3):
super().__init__()
self.levels = levels
blocks = [MaskedConv2d(False, 1, hidden, k, padding=k // 2), nn.ReLU()]
for _ in range(layers - 1):
blocks += [MaskedConv2d(True, hidden, hidden, k, padding=k // 2), nn.ReLU()]
self.net = nn.Sequential(*blocks, nn.Conv2d(hidden, levels, 1))
def forward(self, x): # x: [B, 1, H, W], целые уровни
return self.net(x.float() / (self.levels - 1)) # логиты: [B, levels, H, W]
@torch.no_grad()
def sample(self, x, start=(0, 0)): # заполнить x в порядке развёртки начиная со start
x = x.clone()
H, W = x.shape[-2:]
for i in range(start[0], H):
for j in range(start[1] if i == start[0] else 0, W):
probs = self(x)[:, :, i, j].softmax(dim=1)
x[:, 0, i, j] = torch.multinomial(probs, 1).squeeze(1)
return x
# шаг обучения: loss = F.cross_entropy(model(x), x.long().squeeze(1)), в натах на пиксель
torch.manual_seed(0)
model = PixelCNN()
x = torch.randint(0, 10, (1, 1, 11, 11)).float().requires_grad_(True)
model(x)[0, :, 5, 5].sum().backward() # какие входы влияют на выход в (5, 5)?
print((x.grad[0, 0, 1:6, 1:10] != 0).int())
Напечатанное окно покрывает четыре строки над позицией , её собственную строку и по четыре столбца в каждую сторону:
tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 0],
[1, 1, 1, 1, 1, 1, 1, 0, 0],
[1, 1, 1, 1, 1, 1, 0, 0, 0],
[1, 1, 1, 1, 0, 0, 0, 0, 0]], dtype=torch.int32)
Ни у одного входа в позиции или после неё в порядке развёртки нет ненулевого градиента, а нули в строках выше — это 6 пикселей слепого пятна с рисунка.
Дальше#
PixelCNN предсказывает пиксель по фиксированной окрестности. Авторегрессионные трансформеры заменяют свёртки self-attention, которое на каждом слое видит всё прошлое, и для больших изображений переходят от пикселей к дискретным кодам VQ-автоэнкодера: сетка из 256–1024 токенов описывает изображение , и её можно продолжать по текстовому запросу.
Литература#
- A. van den Oord, N. Kalchbrenner, K. Kavukcuoglu. Pixel Recurrent Neural Networks. ICML, 2016. (2016a)
- A. van den Oord, N. Kalchbrenner, O. Vinyals, L. Espeholt, A. Graves, K. Kavukcuoglu. Conditional Image Generation with PixelCNN Decoders. NeurIPS, 2016. (2016b)
- T. Salimans, A. Karpathy, X. Chen, D. P. Kingma. PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications. ICLR, 2017.
- H. Larochelle, I. Murray. The Neural Autoregressive Distribution Estimator. AISTATS, 2011.
- M. Germain, K. Gregor, I. Murray, H. Larochelle. MADE: Masked Autoencoder for Distribution Estimation. ICML, 2015.
- I. Gulrajani, K. Kumar, F. Ahmed, A. A. Taiga, F. Visin, D. Vazquez, A. Courville. PixelVAE: A Latent Variable Model for Natural Images. ICLR, 2017.