Софтмакс Гумбеля: как устроен и для каких нейронных сетей полезен

2024-08-08 11:45

Всем привет! Меня зовут Николай Лысенко, я занимаюсь рекомендательными системами в Яндекс Маркете. Сегодня хочу затронуть интересную тему: что делать, если в графе вычислений (aka нейронная сеть) возникает дискретное место, через которое не проходит градиент. Как многие знают, для решения этой проблемы есть такие методы, как REINFORCE и софтмакс Гумбеля (Gumbel-Softmax trick). О последнем и пойдёт речь.

Хотя про софтмакс Гумбеля уже много написано, ценность этой статьи в том, что вам не придётся ничего искать в интернете и не потребуется делать выкладки на бумаге. Я постарался собрать всю нужную информацию и расписать все промежуточные вычисления.

Начнём с примера

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

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

Но допустим, что в истории действий пользователя есть информация, что он недавно искал на маркетплейсе холодильник, но пока его не купил. Для первой партии рекомендаций нужно выбрать 30 товаров. Показать ли 30 холодильников? Всё таки нет, так как не исключено, что пользователь уже купил холодильник в другом месте. А что лучше: показать пять холодильников подряд, а потом 25 других товаров, или показывать то холодильник, то пять разных товаров?

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

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

Проблема здесь вот в чём. На -м шаге имелись непрерывные вероятности товаров, но на i+1 -м и последующих шагах вместо выхода -го шага используется просэмплированный из него дискретный идентификатор товара. А ведь через своё влияние на последующие рекомендации -е действие тоже вносит вклад в оптимизируемую ошибку. Чтобы такую нейронную сеть можно было обучить, нужно научиться делать обратный проход через операцию сэмплирования из категориального распределения.

Всё так сложно потому, что, во-первых, награда выдаётся за совокупность действий. Во-вторых, каждое действие зависит от предыдущих. Если отказаться хотя бы от одного из этих условий, никакой софтмакс Гумбеля будет не нужен. Например, в GPT горизонт планирования — ровно одно слово. Благодаря этому (по крайней мере, на претрейне) нет никакого обучения с подкреплением, а есть просто классификация, где нужно угадывать следующее слово. И каким-то чудом так получается, что, двигаясь каждый раз на одно слово вперёд, модель всё равно выдаёт связный текст, а не набор слов. Однако в рекомендациях нельзя показать пользователю товар, посмотреть, как он на него отреагирует, а только потом показать следующий товар. А вот влиянием соседних рекомендаций на текущую можно пренебречь, решив вышеописанную проблему 30 холодильников как-нибудь по-другому: например, через DPP.

Так или иначе, проблема вычисления градиентов при наличии сэмплирования из категориального распределения достаточно общая. Она может возникать не только в обучении с подкреплением, но и в генеративных моделях (наподобие GAN). Чтобы её решить, используют приём с распределением Гумбеля. Его и разберём, но сначала остановимся на пререквизитах.

Вероятностные распределения

Стандартным распределением Гумбеля G(0, 1) (далее просто ) называется распределение с кумулятивной функцией:

$F_G(x) = e^{-e^{-x}}.$

Или, что эквивалентно, — распределение с функцией плотности:

$f_G(x) = e^{-(x + e^{-x})}.$

Сэмплировать величины из стандартного распределения Гумбеля можно методом обратного преобразования:

просэмплировать из распределения, равномерного на отрезке , величину $y = F_G(x) = e^{-e^{-x}}$ ;
вычислить по формуле $x = F_G^{-1}(y) = -log(-log(y))$ .

Какой-то смысл в распределении Гумбеля пока искать не нужно. Просто мы ввели некое распределение, из которого умеем сэмплировать, а позже увидим, что это умение пригодится.

Также далее будет фигурировать экспоненциальное распределение. У экспоненциального распределения с интенсивностью lambda кумулятивная функция равна 0 при отрицательных . А при x ge 0 :

$F_E(x; lambda) = 1 - e^{-lambda x}.$

В то же время функция плотности равна 0 при отрицательных , а на неотрицательной полуоси имеет такой вид:

$f_E(x; lambda) = lambda e^{-lambda x}.$

Связь между этими распределениями такова. Если x sim F_E(x; lambda) , то y = -log(lambda x) sim F_G(y) . Это вытекает из формулы замены переменных. Согласно этой формуле, плотность равна:

$f_E(x; lambda) leftlvert frac{dx}{dy} ight vert= lambda e^{-lambda x} leftlvert frac{d frac{1}{lambda} e^{-y}}{dy} ight vert= lambda e^{-e^{-y}} leftlvert -frac{1}{lambda} e^{-y} ight vert= f_G(y).$

Альтернативный способ сэмплирования из категориального распределения

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

Чтобы сэмплировать из C(pi) , нам потребуется рассмотреть экспоненциальных распределений, у -го из которых интенсивность равна pi_i . Предположим, что случайная величина порождается так:

из каждого из этих распределений сэмплируется по числу ;
.

Убедимся, что такой способ сэмплирования совпадает с сэмплированием из C(pi) . Действительно, вероятность получить некое конкретное значение при этой процедуре равна:

$mathbb{P}[j] = int_0^{+infty} mathbb{P}[x_j = x] mathbb{P}[forall i e j ; x_i > x] dx= int_0^{+infty} f_E(x; pi_j) left(prod_{i e j} (1 - F_E(x; pi_i)) ight) dx= int_0^{+infty} pi_j e^{-pi_j x} left(prod_{i e j} e^{-pi_i x} ight) dx= pi_j int_0^{+infty} e^{-sum_i pi_i x} dx = pi_j.$

Продолжим представлять сэмплирование из C(pi) в альтернативных формах. Вместо сэмплирования x_i из экспоненциальных распределений с интенсивностями pi_i , просэмплируем y_i = -log(pi_i x_i) из стандартного распределения Гумбеля, а потом рассмотрим -log(x_i) = y_i + log(pi_i) . Тогда, поскольку -log(x) является убывающей функцией:

j = arg min_i x_i = arg max_i left(y_i + log(pi_i) ight).

Повторим это ещё раз в виде альтернативной процедуры сэмплирования из C(pi) :

для всех от 1 до независимым образом из стандартного распределения Гумбеля сэмплируются ;
.

При полученной процедуре сэмплирование из категориального распределения, зависящее как от вектора , так и от исхода случайности, оказывается перепараметризованным. Эта перепараметризация даёт то, что вектор теперь отделён от случайности, которая вся сосредоточена в векторе .

Решение проблемы вычисления градиентов

Несмотря на преимущества сэмплирования через распределение Гумбеля, одного его не хватит для вычисления градиентов. В нём фигурирует операция arg max , у которой нет информативного градиента по . Вместо неё можно использовать её дифференцируемый аналог: операцию mathrm{softmax} .

Введём следующие обозначения:

$J_{ extrm{one-hot}}$ — вектор длины , причём такой, что в нём везде нули кроме -й позиции, где ;
$J_{ extrm{continuous}}$ — вектор длины , который получается в результате операции софтмакс с температурой к вектору, -я компонента которого равна . Как известно, чем ближе к 0, тем ближе $J_{ extrm{continuous}}$ к $J_{ extrm{one-hot}}$ , но тем выше риски получить бесполезный градиент.

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

вычисление $J_{ extrm{continuous}}$ по и (температура может быть константным гиперпараметром, а может возвращаться каким-либо подграфом исходного графа);
one-hot-кодирование, превращающее $J_{ extrm{continuous}}$ в $J_{ extrm{one-hot}}$ .

Касательно прямого прохода это означает, что операция сэмплирования реализована через распределение Гумбеля, а функция потерь рассчитывается по $J_{ extrm{one-hot}}$ .
А вот при обратном проходе, казалось бы, лучше не стало. Несмотря на весь проделанный путь, мы снова упираемся в наличие операции, у которой нет информативного градиента. И вот тут уже приходится прибегнуть к некой нестрогости.

Искусственным образом объявляется, что у one-hot-кодирования такой же градиент, как у тождественной операции, то есть при обратном проходе её можно просто-напросто пропустить. Это допущение называется «дуболомным» обратным распространением (straight-through backpropagation). Хотя идея может звучать сомнительно, она описывалась классиками глубокого обучения: например, в статье Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation.

Итого получаем, что при обратном распространении:

от функции потерь до операции сэмплирования вычисляется $frac{partial extrm{Loss}}{partial J_{ extrm{one-hot}}}$ (и в этом нет нестрогости);
для весов , стоявших до сэмплирования при прямом проходе, вместо $frac{partial J_{ extrm{one-hot}}}{partial heta}$ используется $frac{partial J_{ extrm{continuous}}}{partial heta}$ .

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

Источники

Shixiang Gu, Eric Jang, Ben Poole - Categorical Reparametrization with Gumbel-Softmax, 2017
Chris Maddison, Andriy Mnih, Yee Whye Teh - The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables, 2016
Matej Balog, Nilesh Tripuraneni, Zoubin Ghahramani, Adrian Weller - Lost Relatives of the Gumbel Trick, 2017

Источник: habr.com

Софтмакс Гумбеля: как устроен и для каких нейронных сетей полезен

Комментарии: