Численная устойчивость

Катастрофическое сокращение

Из прошлого урока: шаг сетки пропорционален величине числа. Отсюда первое следствие — вычитание близких величин теряет значащие цифры.

Пример, который стоит увидеть, потому что он ломает и в двойной точности. Посчитаем дисперсию набора $[10^8,\ 10^8+1,\ 10^8+2]$ двумя способами:

формула	результат
$\mathbb{E}[x^2] - \mathbb{E}[x]^2$	$2$
$\frac1n\sum (x_i - \bar{x})^2$	$0.6667$

Верный ответ — $2/3$. Наивная формула ошибается втрое, и не в fp32, а в fp64: она вычитает два числа порядка $10^{16}$, различающиеся на единицу, и от результата остаётся мусор.

Отсюда общее правило: центрировать до возведения в квадрат, а не после. Тот же принцип объясняет, почему в блоке 3 формула дисперсии через отклонения от среднего — не педантизм, а необходимость.

$$\htmlData{k=cancel}{\frac{1}{n}\sum_i (x_i - \bar{x})^2} \quad\text{вместо}\quad \htmlData{k=naive}{\frac{1}{n}\sum_i x_i^2 - \bar{x}^2}$$

logsumexp

Второй источник беды — экспонента. $e^{1000}$ переполняется, $e^{-1000}$ обнуляется, и оба случая встречаются в логитах постоянно.

$\log \sum_i e^{z_i} \ = \ m + \log \sum_i e^{z_i - m}, \qquad m = \max_i z_i$

Тождество проверяется вынесением $e^m$ за скобку. Но численно два выражения совершенно разные:

$z$	наивно	со сдвигом
$[1000, 1001]$	$\infty$	$1001.313$
$[-1000, -1001]$	$-\infty$	$-999.687$

После сдвига наибольший показатель равен нулю, то есть $e^0 = 1$ — переполнение невозможно по построению. Остальные слагаемые не больше единицы, и если какое-то обнулилось, оно и было пренебрежимо мало.

Второй случай — обнуление, и он опаснее: наивная формула даёт $-\infty$, что дальше превращается в nan при первом же вычитании. Первый хотя бы шумит явно.

Стабильный softmax

Ровно тот же приём:

$\text{softmax}(z)_i = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}$

И у softmax есть свойство, которое делает это законным: инвариантность к сдвигу. Прибавление константы ко всем логитам не меняет результат, потому что $e^{z_i + c}$ сокращается в числителе и знаменателе.

Проверьте на виджете: сдвиньте все логиты вместе — распределение не изменится. Это не приближение, а тождество, и потому вычитание максимума бесплатно с точки зрения математики и обязательно с точки зрения арифметики.

логит 1

2

60.9%

логит 2

1

22.4%

логит 3

0.5

13.6%

логит 4

-1

3.0%

сдвинуть все логиты

0

температура T

1

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

1

максимум

60.9%

показать якобиан

Температура и затухание градиента

Температура делит логиты: $\text{softmax}(z / T)$. При $T \to 0$ распределение стремится к argmax, при $T \to \infty$ — к равномерному.

Практически важна не сама температура, а то, что происходит с градиентом. Якобиан softmax содержит $p_i(1 - p_i)$, и когда распределение схлопывается, эта величина обращается в нуль:

$T$	максимальная вероятность	масштаб градиента
$1$	$0.644$	$2.3 \cdot 10^{-1}$
$0.5$	$0.865$	$1.2 \cdot 10^{-1}$
$0.1$	$\approx 1$	$4.5 \cdot 10^{-5}$
$0.02$	$\approx 1$	$1.9 \cdot 10^{-22}$

Последняя строка — не «маленький градиент», а его отсутствие: $10^{-22}$ неотличимо от нуля даже в fp32. Отсюда прямая связь с уроком про внимание: масштабирование на $1/\sqrt{d}$ нужно именно затем, чтобы логиты не оказались в этой области.

Приёмы списком

проблема	решение
$\log \sum e^{z}$ переполняется	вычесть максимум
$\log \text{softmax}$	log_softmax, а не log(softmax(x))
$\log(1 + x)$ при малом $x$	log1p
$e^x - 1$ при малом $x$	expm1
произведение вероятностей	сумма логарифмов (блок 3, урок 120)
дисперсия	центрировать до квадрата, или алгоритм Уэлфорда
$\log(1 - \sigma(x))$	через logsigmoid(-x)
накопление длинной суммы	fp32-накопитель даже при bf16-данных

Вторая строка стоит отдельного внимания. log(softmax(x)) вычисляет экспоненты, нормирует, затем берёт логарифм — и если какая-то вероятность обнулилась, получается $-\infty$. log_softmax считает $z_i - \text{logsumexp}(z)$ напрямую: ни одной экспоненты в опасной области, ни одного логарифма от нуля. Именно поэтому в PyTorch кросс-энтропия принимает логиты, а не вероятности — она внутри вызывает log_softmax, и передача уже нормированных вероятностей ломает эту защиту.

Накопители

Последний приём, который часто забывают. Из прошлого урока: в bf16 прибавление $10^{-4}$ к единице не меняет ничего, потому что эпсилон равен $7.8 \cdot 10^{-3}$.

Значит, суммируя тысячу малых чисел в bf16, вы потеряете почти все: как только частичная сумма подрастёт, следующие слагаемые начнут округляться в неё бесследно. Поэтому:

• накопители держат в fp32, даже когда данные в bf16. Матричные ядра так и устроены: умножения в bf16, накопление в fp32;
• мастер-копия весов в fp32 — по той же причине, а не ради точности вычислений;
• для очень длинных сумм существует суммирование Кэхана, восстанавливающее потерянные младшие биты.

Здесь полезно вспомнить формат-виджет: сравните эпсилон bf16 и fp32 и прикиньте, сколько слагаемых нужно, чтобы разница стала заметной.

bf16fp32

раскладка битов

знак · 1

экспонента · 8

мантисса · 7

всего бит

16

эпсилон

7.81e-3

максимум

3.390e+38

мин. нормальное

1.18e-38

десятичных цифр

2.4

Последнее наблюдение из этого сравнения полезно само по себе: у bf16 и fp32 одинаковая экспонента. bf16 — это fp32 с отрезанной мантиссой, поэтому преобразование между ними сводится к отбрасыванию или дописыванию нулей, без всякой обработки диапазона. Это одна из причин, по которой bf16 удобен аппаратно.