Как работает softmax и почему он вызывает проблемы с градиентами при больших logits?

Краткий тезис

Softmax — это функция активации, преобразующая вектор произвольных вещественных чисел (logits) в распределение вероятностей. При больших значениях logits (например, 100) softmax становится практически one-hot (одно значение близко к 1, остальные к 0), что приводит к исчезающим градиентам (gradients|vanishing gradients) — производные стремятся к нулю, и обучение замедляется или останавливается. Основные решения: temperature scaling (деление logits на температуру T > 1), вычитание максимального logit для численной стабильности и использование log-softmax для работы в логарифмическом пространстве.

1. Определение и формула softmax

Softmax (нормализованная экспонента) — функция, которая для вектора z = [z₁, z₂, ..., z_K] вычисляет:

[ [text](/wiki/text){softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} ]

Свойства:

• Сумма всех выходов равна 1.
• Каждый выход в интервале (0, 1).
• Сохраняет порядок: больший z_i → больший softmax.

Logits — это сырые выходы последнего линейного слоя нейронной сети перед softmax. Они могут быть любыми вещественными числами (от -∞ до +∞).

2. Поведение при больших logits

Рассмотрим logits: [100, 90, 80]. Вычислим экспоненты:

• e¹⁰⁰ ≈ 2.688 × 10⁴³
• e⁹⁰ ≈ 1.220 × 10³⁹
• e⁸⁰ ≈ 5.540 × 10³⁴

Сумма практически равна e¹⁰⁰, поэтому:

• softmax первого элемента ≈ 1
• softmax остальных ≈ 0

Результат распределение становится one-hot (одна единица, остальные нули). Это крайне нежелательно для обучения, так как теряется информация о относительной уверенности модели.

3. Проблема градиентов: исчезающие градиенты

Градиент softmax по logit z_i для класса i:

[ \frac{\partial [text](/wiki/text){softmax}(z_i)}{\partial z_i} = [text](/wiki/text){softmax}(z_i) (1 - [text](/wiki/text){softmax}(z_i)) ]

Для больших logits softmax(z_i) ≈ 1, тогда производная ≈ 1·(1-1) = 0.

Для других классов j ≠ i:

[ \frac{\partial [text](/wiki/text){softmax}(z_i)}{\partial z_j} = -[text](/wiki/text){softmax}(z_i) \cdot [text](/wiki/text){softmax}(z_j) ]

Если softmax(z_i) ≈ 1, а softmax(z_j) ≈ 0, то производная ≈ 0.

Итог все градиенты стремятся к нулю → веса перестают обновляться → обучение замирает. Это классический случай vanishing gradient problem.

4. Математический анализ: почему градиенты малы

Рассмотрим производную softmax для доминирующего класса:

[ \frac{\partial p_i}{\partial z_i} = p_i (1 - p_i) ]

Где p_i = softmax(z_i). Функция p_i(1-p_i) достигает максимума 0.25 при p_i = 0.5. При p_i → 1 или p_i → 0 производная → 0.

Для больших logits p_i → 1, поэтому градиент → 0. Для очень маленьких logits (отрицательных) p_i → 0, градиент также → 0. Таким образом, softmax «чувствителен» только в области, где p_i около 0.5, то есть когда logits соизмеримы.

5. Temperature scaling (масштабирование температуры)

Temperature scaling — деление logits на параметр T (температура) перед softmax:

[ [text](/wiki/text){softmax}T(z_i) = \frac{e^{z_i / T}}{\sum{j} e^{z_j / T}} ]

• T = 1 — стандартный softmax.
• T > 1 — «смягчает» распределение: разница между вероятностями уменьшается, градиенты становятся ненулевыми.
• T < 1 — «усиливает» контраст, приближает к one-hot.

Применение в дистилляции знаний (knowledge distillation) используют T > 1 для передачи «мягких» меток; в генеративных моделях (LLM) temperature контролирует креативность.

6. Вычитание максимального logit (subtract max)

Для численной стабильности при вычислении softmax на практике вычитают максимальный logit:

[ [text](/wiki/text){softmax}(z_i) = \frac{e^{z_i - \max(z)}}{\sum_{j} e^{z_j - \max(z)}} ]

Это не меняет результат (так как вычитание константы из всех z_i эквивалентно делению числителя и знаменателя на e^{max(z)}), но предотвращает переполнение экспоненты (overflow). Например, e¹⁰⁰ — огромное число, а e⁰ = 1 — безопасно.

Проблема градиентов не решается — распределение остаётся one-hot, но вычисления становятся стабильными.

7. Log-softmax и его преимущества

Log-softmax — логарифм от softmax:

[ \log[text](/wiki/text){softmax}(z_i) = z_i - \log\sum_{j} e^{z_j} ]

Часто используется в функциях потерь (например, cross-entropy loss с log-softmax) по двум причинам:

• Численная стабильность log(exp(большое)) может дать NaN, но формула с вычитанием max решает это.
• Градиенты производная log-softmax по z_i равна (1 - softmax(z_i)), что при больших logits даёт ~0, но при обратном распространении через loss (cross-entropy) градиенты становятся (softmax(z_i) - target_i), что не затухает, если target_i = 1 для правильного класса. Однако если модель слишком уверена (softmax ≈ 1), градиент всё равно мал.

Важно log-softmax сам по себе не решает проблему vanishing gradients, но в комбинации с temperature scaling даёт хорошие результаты.

8. Практические рекомендации для обучения

1. Инициализация весов избегать слишком больших начальных logits. Использовать инициализацию Xavier/Glorot или He.
2. Batch normalization: стабилизирует распределение активаций, уменьшает разброс logits.
3. Label smoothing заменяет one-hot метки на «смягчённые» (например, [0.9, 0.05, 0.05]), что не позволяет softmax становиться слишком уверенным.
4. Gradient clipping ограничивает градиенты по норме, но не решает причину.
5. Temperature scaling добавлять learnable температуру или использовать T > 1 на ранних этапах обучения.

9. Связь с другими функциями активации

Softmax — это многоклассовое обобщение sigmoid. Проблема vanishing градиентов характерна для всех «насыщающихся» функций активации.

10. Пример кода: визуализация градиентов

import numpy as np
import matplotlib.pyplot as plt

def softmax(z, T=1.0):
    z = z / T
    e_z = np.exp(z - np.max(z))
    return e_z / np.sum(e_z)

def softmax_grad(z, T=1.0):
    p = softmax(z, T)
    # диагональные элементы: p_i*(1-p_i)
    grad = np.diag(p * (1 - p))
    # недиагональные: -p_i * p_j
    for i in range(len(z)):
        for j in range(len(z)):
            if i != j:
                grad[i, j] = -p[i] * p[j]
    return grad

# Пример
logits = np.array([100, 90, 80])
print("Softmax:", softmax(logits))
print("Градиент по первому logit:", softmax_grad(logits)[0, 0])  # ≈ 0

Вывод: градиент практически нулевой.

Пет-проект для закрепления

Задача Реализовать эксперимент, демонстрирующий влияние температуры на градиенты softmax.

Инструменты Python, NumPy, Matplotlib.

Шаги:

1. Сгенерировать набор logits (например, от -10 до 10 с шагом 0.5).
2. Для каждого logits вычислить softmax и его градиент по доминирующему классу.
3. Повторить для T = 0.5, 1, 2, 5.
4. Построить графики зависимости градиента от logits для разных T.

Ожидаемый результат График покажет, что при T=1 градиент быстро падает до нуля при |logits| > 5, а при T=5 градиент остаётся значительным в широком диапазоне.

Дополнительно Реализовать обучение простой двухслойной сети на MNIST с разными температурами и сравнить скорость сходимости.

Связь с другими вопросами

Навигация

• Предыдущий: 660
• Следующий: 662
• Индекс: 00. Индекс разборов