8.4. Нормализующие потоки

Введение

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

• генерировать объекты , где – параметры модели;
• оценивать вероятность встретить данный объект среди набора наблюдаемых данных ;
• выучивать скрытые представления для объекта .

Известными примерами генеративных моделей являются:

• Авторегрессионные модели:

• Вариационные автокодировщики:

Но оба эти метода не позволяют одновременно:

• получать скрытые представления для объектов
• точно вычислять функцию правдоподобия

Нормализующие потоки способны решить обе эти задачи.

Мотивация

Пусть , где неизвестно, а . Мы хотим найти отображение , для которого и .

Отображение преобразует базовую функцию плотности к более сложной . С его помощью мы можем генерировать сложный объект путем сэмплинга простого объекта (скрытой переменной) из распределения и применения «генератора» . Обратное отображение «нормализует» сложное распределение , приводя его к простому .

Найдя такое отображение , мы сможем генерировать новые объекты , а оценить плотность поможет формула преобразования плотности случайной величины. Давайте её вспомним.

Формула замены переменной

Пусть , , при этом отображение дифференцируемо, обратимо и . Тогда:

где – якобиан отображения .

Определение

Итак, нормализующий поток – это обратимое дифференцируемое отображение , которое переводит исходные представления объектов в скрытые: и .

При этом функция правдоподобия вычисляется по формуле:

Умея вычислять функцию правдоподобия, мы можем обучать нашу модель методом максимума правдоподобия (ММП):

где – выборка наблюдаемых данных из распределения .

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

Тогда якобиан вычисляется по формуле:

Но вычисление якобиана является очень затратной операцией. Для того, чтобы мы могли обучать модели эффективно на высокоразмерных данных (аудио, изображения), необходимо использовать такое отображение , подсчет якобиана которого был бы эффективен!

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

Примером такого отображения является планарный поток (Planar Flow), где отображение принадлежит следующему семейству функций:

где – обучаемые параметры, а – гладкая нелинейная функция, например, .

Якобиан такого отображения можно будет посчитать за . Обозначим

Тогда

Действие планарного нормализующего потока на нормальное и равномерное распределение (ссылка на статью).

Развитие идеи

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

Следующие модели активно использовали этот трюк.

NICE: Non-linear Independent Component Estimation и RealNVP

Авторы модели NICE предложили использовать в качестве следующее семейство преобразований:

где , а – произвольная нейросеть с входами и выходами. Такое преобразование называют аддитивным связыванием (additive coupling).

Обратное преобразование вычисляется с такой же легкостью, а якобиан равен . То есть, , что является довольно сильным ограничением модели.

Далее, из-за того, что

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

Позже авторы NICE позже предложили использовать между слоями нормализующих потоков зафиксированные перестановки признаков/каналов , что стало основой работы RealNVP. Использование перестановок позволяет добиться того того, чтобы все выходные каналы оказались затронуты преобразованием ; при этом градиент перестановки вычисляется легко.

где – поэлементное умножение, а – нейросеть, которая может быть произвольной, но, как правило, выбирается такой же архитектуры, как и . Такое преобразование называют аффинным связыванием (affine coupling).

Получившееся отображение тоже легко обращается, а его якобиан равен:

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

Чтобы улучшить сходимость глубоких () нормализующих потоков, авторы предложили использовать Batch Normalization. Данное преобразование тоже является обратимым, а его якобиан вычисляется крайне просто.

В результате, выразительная способность модели сильно повысилась, и она стала способна выучивать сложные распределения:

Masked Autoregressive Flows

Ссылка на статью

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

где и – нейросети произвольной архитектуры.

Как видно из формулы, напрямую зависит от . Таким образом, элементы генерируются авторегрессивно, что и дало название архитектуре.

Якобиан такого преобразования вычисляется по следующей формуле:

Таким образом, шаг генерации выглядит следующим образом:

1. $ z \sim \mathcal{N}(0, I) $
2. $ x_1 = z_1 \exp(\alpha_1) + \mu_1 $
3. $ x_2 = z_2 \exp(f_{\alpha_2}(x_1)) + f_{\mu_2}(x_1) $
   ...

Однако вычисление скрытых переменных не является авторегрессивным:

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

Inverse Autoregressive Flows

Ссылка на статью

Чтобы быстро генерировать объекты из сложного распределения, мы можем избавиться от авторегрессивности на шаге генерации, поставив в авторегрессивную зависимость не наблюдаемые, а латентные переменные:

Можем заметить, что проблема долгого вычисления авторегрессивных выражений никуда не уходит. Мы лишь изменяем построение модели таким образом, чтобы генерировать объекты быстрее:

Но вычисление , а значит и правдоподобия, становится долгим, и обучение занимает больше времени.

Звук

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

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

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

Нормализующий поток учится генерировать сигнал в виде waveform-а на основе спектрограммы путем обратного преобразования.

Чтобы генерировать более длинные фрагменты звука, модель генерирует короткие звуковые кадры (фреймы) за раз, которые затем объединяются для формирования полного waveform-а.

Теперь мы готовы узнать про применение нормализующих потоков в генерации аудио!

Probability Density Distillation и Parallel WaveNet

Ссылка на статью

Архитектура Inverse Autoregressive Flow (IAF) была изначально предложена для задачи генерации аудио. Она позволяет генерировать объекты крайне эффективно, но обучение методом максимального правдоподобия занимает много времени из-за авторегрессивности вычислений. Метод Probability Density Estimation позволяет решить эту проблему с помощью использования второй предобученной авторегрессивной модели в качестве учителя. IAF обучается в качестве модели-студента, минимизируя KL дивергенцию где и – распределения студента и учителя соответственно. Ключевым достижением данного подхода является то, что вычисление функции потерь требует вычисления кросс-энтропии между учителем и студентом, а не правдоподобия, что позволяет максимально распараллелить все вычисления ввиду отсутствия авторегрессивности в вычислениях.

Вместе с тем в данной работе в качестве учителя выбирается не случайная модель, а оригинальная авторегрессивная модель WaveNet, которая в 2016 году позволила достичь state-of-the-art качества генерации аудио. Эта модель является не нормализующим потоком, а обыкновенной авторегрессивной моделью, которая обучается предсказывать следующий кусочек аудио (фрейм) длиной в несколько миллисекунд.

Таким образом, с помощью IAF и Probability Density Distillation авторам удалось ускорить генерацию более чем в 1000 раз без потери качества!

На картинке выше мы видим, что модель использует лингвистические признаки для генерации аудио. Эта задача является примером задачи условной генерации, где на вход модели подается спектрограмма, сгенерированная отдельной моделью по тексту, а на выход ожидается речь в аудио-формате (waveform). О том, как модель использует дополнительную информацию для обуславливания, поговорим в главе про Waveglow

Glow

Исследователи из OpenAI в 2018 году опубликовали работу Glow: Generative Flow with Invertible 1×1 Convolutions, которая значительно улучшает результаты модели RealNVP. Опишем два главых улучшения.

Во-первых, для перемешивания каналов Glow использует обратимые свертки с ядром 1x1 вместо фиксированной матрицы перестановок каналов в RealNVP;

Это нововведение является по-настоящему красивым, так как в нем предлагается способ вычисления якобиана 2D-свертки за . А именно, логарифм якобиана такой 1x1-свертки с числом каналов для тензора размера равен , где – матрица свёртки 1х1.

Авторы предлагают использовать следующий вариант LU-разложения для матрицы :

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

Нетрудно показать, что

Благодаря этому авторам удалось снизить сложность вычислений якобиана с до

Кроме того, для улучшения сходимости использовали собственно разработанный actnorm-слой (activation normalization). Поскольку нормализующие потоки требуют много вычислительных ресурсов, для обучения используются мини-батчи маленького размера, из-за чего батч-нормализация работает не очень хорошо. Авторы предлагают использовать следующий тип нормализации – actnorm:

• Нормализуем входной тензор (промежуточное изображение) по размерности каналов;
• Инициализируем параметры смещения и разброса статистиками с первого батча;
• Далее обучаем их в качестве обычных параметров.

Таким образом, один блок нормализующего потока выглядит так:

WaveGlow

Ссылка на статью

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

Как мы помним, эта задача также является примером задачи условной генерации:

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

В качестве обуславливающего фактора для WaveGlow мы имеем сгенерированную по тексту mel-спектрограмму, а на выходе ожидаем получить соответствующую тексту и спектрограмме аудио-запись. Как мы видим на изображении и в формулах ниже, mel-спектрограмма используется как дополнительный признак для нейросети, генерирующей параметры афинного преобразования. В качестве модели, которая производит параметры афинного преобразования, используется похожая на WaveNet архитектура с dilated-свертками.

Правая часть схемы ниже более подробно показывает строение слоя affine coupling:

Операция разделяет тензор пополам на два тензора меньшей размерности и для их последующего участия в слое аффинного связывания (affine coupling).

Пример генерации:
Источник

Ground truth	WaveNet	WaveGlow

Out-of-distribution detection

Может показаться, что способность точно и эффективно вычислять функцию правдоподобия может позволить без труда обнаруживать аномалии в данных, что может пригодиться во многих приложениях. Однако в работе Kirichenko et al. на примере задачи генерации изображений было показано, что нормализующие потоки выучивают отображение картинок в латентное пространство, основываясь на локальных корреляциях пикселей и графических деталях, а не на семантическом контенте. Из-за этого правдоподобие OOD-объектов может быть выше, чем правдоподобие in-distribution сэмплов.

Однако позже было предложено использовать ряд эвристик для того, чтобы улучшить способность к детекции аномалий за счет подсчета значения функции правдоподобия:

1. Использовать значение правдоподобия второй модели потока, обученного на отличном от исходного датасете (например, ImageNet при исходном CelebA). А затем вычислять отношенение этих двух значений для вынесения вердикта об аномальности объекта. Schirrmeister et al.
2. В работе Serrà et al. показали, что проблема качества нормализующих потоков в задаче детекции аномалий связана с чрезмерным влиянием сложности входных данных на значение функции правдоподобия. Поэтому авторы предложили использовать в качестве поправки размер сжатого изображения с помощью одного из алгоритмов компрессии (JPEG2000/PNG).

Сравнение с другими типами генеративных моделей

Обратимся к статье Bond-Taylor et al., в которой приводится количественный анализ всех существующих семейств генеративных моделей в задаче генерации изображений из датасета CIFAR-10.

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

1. скорость обучения;
2. скорость генерации;
3. число обучаемых параметров;
4. разрешение генерируемого изображения;
5. ограничение на форму якобиана;
6. возможность вычислять правдоподобие объекта;
7. FID (Fréchet Inception Distance) тестовой выборки;
8. Отрицательный логарифм правдоподобия тестовой выборки.

За расшифровкой обратимся к таблице ниже:

Подведя итог, можно сказать, что нормализующие потоки:

• требуют очень много времени на обучение, так как при обучении проводятся нетривиальные неоптимизированные вычисления;
• имеют скорость генерации, сравнимую с GAN-ами;
• менее эффективны по соотношению качество/число параметров, чем GAN-ы и диффузионные модели;
• позволяют быстро вычислять точное значение функции правдоподобия объекта;
• обладают сравнительно неплохим качеством генерации, проигрывающим GAN-ам и диффузионным моделям.

Итак, нормализующие потоки явно выделяются среди других семейств генеративных моделей своими свойствами – обратимостью и способностью вычислять правдоподобие объекта. Но если для решения задачи они не требуются, то имеет смысл попробовать другие модели – в первую очередь, GAN-ы и диффузионные модели.