Softmax: от логитов к вероятностям

Зачем он нужен

Классификатор на $K$ классов в конце выдаёт $K$ чисел — логитов (logits): например, $(2;\,1;\,0)$ для «кошка», «собака», «птица». Это просто баллы: они бывают отрицательными, не складываются в 1, и по ним не видно, насколько модель уверена. Softmax превращает такой вектор в распределение вероятностей: все числа положительны, в сумме дают 1, а порядок сохраняется — у кого балл больше, у того и вероятность больше.

Где он встречается:

• выход многоклассового классификатора — от логистической регрессии до ResNet и ViT: $\hat y=\operatorname{softmax}(Wx+b)$;
• языковые модели — на каждом шаге LLM выдаёт логиты для всех токенов словаря (десятки и сотни тысяч), softmax делает из них распределение следующего токена, а температура регулирует «смелость» генерации;
• механизм внимания (attention) в трансформерах — веса $\operatorname{softmax}(QK^\top/\sqrt d)$ решают, на какие токены смотреть;
• смесь экспертов (MoE) — маршрутизатор распределяет токен между экспертами через softmax;
• обучение с подкреплением — больцмановская политика выбирает действие $a$ с вероятностью, пропорциональной $e^{Q(a)/T}$;
• статистическая физика — это распределение Больцмана — Гиббса $p_i\propto e^{-E_i/(kT)}$; отсюда и слово «температура».

Интуиция: голосование с усилением

Представьте, что классы — кандидаты, а логиты — их баллы. Чтобы получить доли «голосов», баллы нужно сделать положительными и нормировать. Просто разделить на сумму нельзя: баллы бывают отрицательными, а сумма — нулевой. Экспонента решает обе проблемы: $e^{z}$ всегда положительна, а разница баллов превращается в отношение: каждая единица преимущества делает долю кандидата в $e\approx2{,}718$ раза больше доли соперника.

Потяните верхушки столбиков-логитов слева и посмотрите, как меняются вероятности справа. Потом нажмите «+1 ко всем логитам»: столбики слева поднимутся, а справа ничего не изменится.

Формула по косточкам

$z=(z_1,\dots,z_K)$

вектор логитов — сырые выходы последнего слоя, любые действительные числа;

$e^{z_i}$

экспонента: делает число положительным и превращает разности логитов в отношения;

$\sum_{j}e^{z_j}$

нормирующая сумма (в физике — статистическая сумма, partition function $Z$);

$s_i$

вероятность $i$-го класса: $s_i\in(0;\,1)$ и $\sum_i s_i=1$.

Пройдём сквозной пример шпаргалки — логиты $z=(2;\,1;\,0)$:

Температура

Часто логиты перед softmax делят на число $T>0$ — температуру:

Таблица показывает, что температура делает с тем же вектором $(2;\,1;\,0)$:

Два предела объясняют название. При $T\to0$ выигрывает только наибольший логит: softmax стремится к one-hot вектору $\operatorname{argmax}$ (если максимум единственный). При $T\to\infty$ разности исчезают, и распределение становится равномерным, $s_i\to1/K$. Softmax — «мягкий» (soft) argmax: промежуточное состояние между жёстким выбором и полным безразличием.

Свойства

• Распределение: $s_i>0$ и $\sum_i s_i=1$ при любых логитах. Ровно 0 или 1 не получается никогда — только в пределе (а на компьютере — из-за округления).
• Сохраняет порядок: $z_i>z_j\iff s_i>s_j$, поэтому $\operatorname{argmax}s=\operatorname{argmax}z$. Чтобы просто предсказать класс, softmax можно не считать.
• Отношения: $s_i/s_j=e^{z_i-z_j}$ — разность логитов равна логарифму отношения вероятностей. Обратно, $z_i=\ln s_i+C$: логиты восстанавливаются по вероятностям с точностью до общей константы.
• Инвариантность к сдвигу: $\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$ для любого числа $c$. Из $K$ логитов значимы только $K-1$ разностей.
• Нет инвариантности к масштабу: $\operatorname{softmax}(az)$ при $a>1$ острее, при $0<a<1$ — площе. Это та же температура, $a=1/T$.
• Конкуренция: рост одного логита увеличивает его вероятность и уменьшает все остальные — масса перераспределяется, а не создаётся.

Два класса — это сигмоида

Для двух классов softmax — это сигмоида разности логитов, поэтому бинарному классификатору хватает одного выхода. Похожая формула верна и для $K$ классов, если смотреть на один класс против всех остальных:

Каждая вероятность — сигмоида от отрыва своего логита от «гладкого максимума» остальных. Чтобы проверить, разделите числитель и знаменатель $s_i$ на $e^{z_i}$.

Геометрия: симплекс вероятностей

Все распределения на $K$ классах образуют вероятностный симплекс:

Для $K=3$ это равносторонний треугольник. Вершины — уверенные ответы $(1;\,0;\,0)$, $(0;\,1;\,0)$, $(0;\,0;\,1)$, центр — равномерное распределение $(\tfrac13;\,\tfrac13;\,\tfrac13)$, а точка внутри — «смесь» вершин с весами $s_i$ (барицентрические координаты).

Softmax отображает всё пространство логитов $\mathbb R^K$ на внутренность симплекса: границы и вершины недостижимы. Все точки прямой $z+c\mathbf 1$ переходят в одну точку, поэтому по сути softmax — взаимно однозначное соответствие между $(K-1)$-мерным пространством разностей логитов и внутренностью симплекса. Обратное отображение — логарифм: $z=\ln s$ с точностью до сдвига.

Производная: матрица Якоби

У softmax $K$ входов и $K$ выходов, поэтому его производная — не число, а матрица $K\times K$: элемент $J_{ij}=\partial s_i/\partial z_j$ показывает, как $i$-я вероятность откликается на $j$-й логит.

Свойства этой матрицы:

• диагональ $s_i(1-s_i)$ — ровно производная сигмоиды: не больше $\tfrac14$, максимум при $s_i=\tfrac12$;
• вне диагонали все элементы отрицательны: подняв один логит, мы отнимаем вероятность у остальных;
• симметрична, и каждая строка (и столбец) в сумме даёт 0: $J\mathbf 1=s-s\,(s^\top\mathbf 1)=0$. Это производная равенства $\sum_i s_i=1$, а вектор $\mathbf 1$ — направление сдвига, который ничего не меняет. Поэтому $J$ вырождена;
• неотрицательно определена: $v^\top Jv=\sum_i s_iv_i^2-\bigl(\sum_i s_iv_i\bigr)^2$ — дисперсия величины $v$ при распределении $s$, а дисперсия не бывает отрицательной;
• затухает при уверенности: когда $s$ близок к вершине симплекса, все элементы $J$ близки к нулю — насыщение, как у сигмоиды.

Обратный проход без матрицы

В backprop нужна не сама $J$, а произведение $J^\top g$, где $g=\partial L/\partial s$ пришёл сверху. Благодаря структуре $\operatorname{diag}(s)-ss^\top$ это стоит $O(K)$ операций вместо $O(K^2)$:

Для словаря LLM с $K\approx10^5$ токенов это разница между $10^5$ и $10^{10}$ чисел на каждую позицию. А в паре с кросс-энтропией всё ещё проще: $\partial L/\partial z=s-y$ (вывод — на странице кросс-энтропии).

3D-картина: softmax над плоскостью логитов

Зафиксируем $z_3=0$ — это не теряет общности, ведь сдвигом любой логит можно сделать нулём, — и нарисуем $s_1$ как функцию двух переменных $z_1$ и $z_2$. Получается «рампа»: поверхность близка к 1 там, где $z_1$ больше и $z_2$, и нуля, и близка к 0 там, где первый класс проигрывает. Поперёк границы она поднимается по сигмоиде — это формула «один против всех» в действии.

Вид сверху делит плоскость логитов на три области. Границы — множества равных логитов: луч $z_1=z_2\ge0$ разделяет классы 1 и 2, луч $z_1=0$, $z_2\le0$ — классы 1 и 3, луч $z_2=0$, $z_1\le0$ — классы 2 и 3. Все три луча выходят из начала координат, где $s=(\tfrac13;\,\tfrac13;\,\tfrac13)$.

Log-sum-exp: гладкий максимум

У softmax есть «первообразная» — функция log-sum-exp:

LSE — гладкий максимум: наибольшая экспонента доминирует в сумме, и логарифм возвращает почти её показатель. Точнее,

Левое неравенство верно, потому что сумма не меньше своего наибольшего слагаемого, правое — потому что каждое из $K$ слагаемых не больше наибольшего. Для $(2;\,1;\,0)$: $\operatorname{LSE}=\ln11{,}107\approx2{,}408$ — между $2$ и $2+\ln3\approx3{,}099$. С температурой $\operatorname{LSE}_T(z)=T\ln\sum_j e^{z_j/T}$ зазор сжимается до $T\ln K$: при $T\to0$ LSE становится обычным максимумом, а его градиент $\operatorname{softmax}(z/T)$ — one-hot вектором argmax.

Softmax в обучении

Простейшая модель с softmax — многоклассовая логистическая (softmax-) регрессия: $K$ линейных функций дают логиты, softmax — вероятности, потеря — кросс-энтропия:

Здесь $y$ — one-hot вектор правильного класса. Вся матрица Якоби в паре с логарифмом сократилась до разности «предсказание − правда».

Softmax внутри трансформера: внимание

Каждая строка $\operatorname{softmax}(\cdot)$ — распределение внимания одного токена по токенам последовательности (включая его самого), сумма весов 1. Деление на $\sqrt d$ — это температура: скалярное произведение $d$-мерных векторов с независимыми компонентами нулевого среднего и единичной дисперсии имеет дисперсию $d$. Без масштабирования логиты внимания велики, softmax насыщается до почти one-hot, а матрица Якоби — и с ней градиенты — почти обнуляется.

Температура при генерации текста

Языковая модель выбирает следующий токен из $\operatorname{softmax}(z/T)$. При $T\to0$ (жадный выбор) берётся самый вероятный токен, при $T<1$ текст предсказуемее, при $T>1$ — разнообразнее и рискованнее. Температура $T>1$ работает и при дистилляции знаний (Хинтон и др., 2015): «мягкие» вероятности учителя несут больше информации, чем one-hot метки.

В коде: устойчивость, log-softmax и PyTorch

Наивная формула ломается на больших логитах: $e^{1000}\approx2\cdot10^{434}$ не помещается ни во float32, ни во float64 и превращается в бесконечность, а дальше получается $\infty/\infty=\text{NaN}$. Во float32 (и в bfloat16 с тем же диапазоном порядков) предел ближе: $e^{x}$ переполняется уже при $x>88{,}72$ (наибольшее число float32 $\approx3{,}4\cdot10^{38}$). На очень отрицательных логитах беда обратная: экспоненты обнуляются, и выходит $0/0$.

Лекарство — инвариантность к сдвигу: вычтем $m=\max_j z_j$. Наибольший показатель станет равен 0, $e^0=1$, знаменатель окажется между 1 и $K$ — ни переполнения, ни деления на ноль:

pythonimport numpy as np
np.set_printoptions(precision=3)

def softmax(z, axis=-1, T=1.0):
    """Устойчивый softmax: вычитаем максимум, результат не меняется."""
    z = np.asarray(z, dtype=float) / T
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

def log_softmax(z, axis=-1):
    # log s_i = z_i − logsumexp(z): без log(0) даже для крошечных вероятностей
    z = np.asarray(z, dtype=float)
    m = z.max(axis=axis, keepdims=True)
    return z - m - np.log(np.exp(z - m).sum(axis=axis, keepdims=True))

def softmax_backward(s, g):
    # J^T g = s * (g − <g, s>): O(K) вместо матрицы K×K
    return s * (g - (g * s).sum(axis=-1, keepdims=True))

z = np.array([2.0, 1.0, 0.0])
print(softmax(z))            # [0.665 0.245 0.09 ]
print(softmax(z + 1000))     # [0.665 0.245 0.09 ] — сдвиг не важен
print(softmax(z, T=0.5))     # [0.867 0.117 0.016] — острее
print(log_softmax([0.0, -800.0]))  # [   0. -800.], а np.log(softmax(...)) даст -inf

pythonimport torch
import torch.nn.functional as F

z = torch.tensor([2.0, 1.0, 0.0], requires_grad=True)
s = F.softmax(z, dim=-1)          # tensor([0.6652, 0.2447, 0.0900])
s[0].backward()                   # первая строка матрицы Якоби
print(z.grad)                     # tensor([ 0.2227, -0.1628, -0.0599])

# Для обучения softmax отдельно не нужен:
logits = torch.randn(4, 3)        # батч: 4 примера, 3 класса
target = torch.tensor([0, 2, 1, 0])
loss = F.cross_entropy(logits, target)   # внутри — log_softmax + NLL
log_p = F.log_softmax(logits, dim=-1)    # если нужны логарифмы вероятностей

Частые ошибки

1. Softmax перед nn.CrossEntropyLoss. Эта потеря ждёт логиты и сама применяет log-softmax. Двойной softmax зажимает вероятности: для трёх классов второй softmax не даёт больше $e/(e+2)\approx0{,}576$, потеря не опускается ниже $0{,}551$, градиенты искажены.
2. Не та ось. Для батча формы $(N,K)$ нормировать нужно по классам (dim=-1), а не по примерам (dim=0). Исключения не будет — модель просто учит бессмыслицу.
3. Наивная экспонента без вычитания максимума — NaN на больших логитах, особенно в смешанной точности: во float16 $e^{x}$ переполняется уже при $x>11{,}09$.
4. Softmax для multi-label. Если у объекта может быть несколько меток сразу (кот и собака на фото), классы не конкурируют — нужна сигмоида на каждом выходе и BCE.
5. Вероятность softmax как гарантия. $s_{\max}=0{,}99$ не значит «прав в 99% случаев»: глубокие сети часто переуверены. Калибровка — отдельная задача; простейший способ — temperature scaling, подбор одного $T$ на валидации.
6. Сравнение логитов разных примеров. Логиты определены с точностью до сдвига, поэтому «логит 5 у одного примера больше, чем 3 у другого» ничего не значит. Сравнивайте вероятности или разности логитов внутри одного примера.

Связи с другими формулами

• Сигмоида: $\operatorname{softmax}(z_1,z_2)_1=\sigma(z_1-z_2)$, а диагональ якобиана $s_i(1-s_i)$ — производная сигмоиды.
• Кросс-энтропия: естественная пара softmax; вместе дают градиент $s-y$, а $\mathrm{CE}=\operatorname{LSE}(z)-z_y$.
• BCE: частный случай кросс-энтропии для двух классов, где softmax сводится к сигмоиде.
• Focal Loss: домножает $-\ln p_t$ на $(1-p_t)^\gamma$, где $p_t$ — softmax-вероятность правильного класса.
• ReLU: softplus $\ln(1+e^{x})=\operatorname{LSE}(x,0)$ — гладкий вариант $\max(x,0)$.

Итоги

Проверь себя