Градиент softmax и кросс-энтропии

Две функции

Softmax превращает произвольные числа (логиты) в распределение:

$p_i = \frac{e^{z_i}}{\sum_k e^{z_k}}$

Кросс-энтропия измеряет, насколько распределение $\mathbf{p}$ расходится с правильным ответом $\mathbf{y}$ (обычно one-hot):

$L = -\sum_i y_i \log p_i$

Для one-hot $\mathbf{y}$ с единицей на позиции $c$ это просто $L = -\log p_c$. Смысл разберём в блоке 4; здесь нас интересует только производная.

Покрутите логиты, сдвиг и температуру. Сдвиг ничего не меняет — это первое, что понадобится в выводе:

логит 1

2

60.9%

логит 2

1

22.4%

логит 3

0.5

13.6%

логит 4

-1

3.0%

сдвиг всех логитов

0

температура

1

сумма вероятностей

1

наибольшая

60.9%

строка якобиана

Инвариантность к сдвигу

$\frac{e^{z_i + c}}{\sum_k e^{z_k + c}} = \frac{e^c e^{z_i}}{e^c\sum_k e^{z_k}} = p_i$

Множитель $e^c$ сокращается. Два следствия:

• у softmax есть лишняя степень свободы: логиты определены с точностью до общего сдвига, поэтому у задачи с $K$ классами по существу $K-1$ параметров;
• отсюда берётся стандартный трюк устойчивости: вычесть $\max_k z_k$ перед экспонентой. Тогда наибольший показатель равен нулю, exp не переполняется, а ответ не меняется. Ровно это делает logsumexp, и к нему вернёмся в блоке 6.

Якобиан softmax

Дифференцируем $p_i$ по $z_j$. Случай $i = j$ и случай $i \ne j$ считаются по правилу частного и объединяются в одну формулу:

$$\frac{\partial p_i}{\partial z_j} = p_i\bigl(\htmlData{k=delta}{\delta_{ij}} - \htmlData{k=coupling}{p_j}\bigr)$$

Это матрица, а не вектор: softmax отображает $\R^K \to \R^K$. Диагональные элементы $p_i(1-p_i)$ положительны, внедиагональные $-p_ip_j$ отрицательны. Внедиагональный член и есть содержательная часть: вероятности связаны условием нормировки, и поднять одну можно только за счёт остальных.

Включите строку якобиана в виджете выше и проверьте знаки.

Сокращение

Теперь соберём градиент лосса по логитам. По правилу цепочки с суммой по путям — через каждую компоненту $\mathbf{p}$:

$\frac{\partial L}{\partial z_j} = \sum_i \frac{\partial L}{\partial p_i}\frac{\partial p_i}{\partial z_j}$

Первый множитель: $\partial L/\partial p_i = -y_i/p_i$. Подставляем якобиан:

$\frac{\partial L}{\partial z_j} = \sum_i \left(-\frac{y_i}{p_i}\right)p_i(\delta_{ij} - p_j) = -\sum_i y_i(\delta_{ij} - p_j)$

Множитель $p_i$ сократился — это первое чудо. Раскрываем:

$= -y_j + p_j\sum_i y_i = p_j - y_j$

потому что $\sum_i y_i = 1$ для распределения. Итого:

$$\frac{\partial L}{\partial \mathbf{z}} = \htmlData{k=res}{\mathbf{p} - \mathbf{y}}$$

Ни экспонент, ни делений, ни матрицы якобиана — просто разность предсказания и правильного ответа. Матрица $K \times K$ сократилась в вектор длины $K$.

Почему это важно, а не просто красиво

Численно. В $-\log p_c$ при малом $p_c$ и в $1/p_i$ теряется точность. В $\mathbf{p} - \mathbf{y}$ терять нечего. Именно поэтому фреймворки предоставляют cross_entropy(logits, target) одной функцией, а не композицию log(softmax(x)) с NLL: сокращение делается аналитически, до вычислений.

Если вы когда-нибудь получали NaN, написав log(softmax(x)) вручную, — вот объяснение. Это не хрупкость библиотеки, а отказ от готового сокращения.

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

Температура

$p_i = \frac{e^{z_i/T}}{\sum_k e^{z_k/T}}$

При $T \to 0$ распределение стремится к one-hot на максимуме, при $T \to \infty$ — к равномерному. Покрутите ползунок температуры: логиты не меняются, а распределение меняется от почти детерминированного до почти равномерного.

Где встречается: сэмплирование из языковых моделей, дистилляция знаний, Gumbel-softmax, энтропийная регуляризация в RL. Везде это один параметр — насколько распределению позволено быть уверенным.