Как log_softmax () реализован для вычисления его значения (и градиента) с лучшей скоростью и числовой стабильностью?

Question

И M XNet, и PyTorch предоставляют специальную реализацию для вычисления журнала (softmax ()), которая быстрее и численно более стабильна. Однако я не могу найти фактическую реализацию Python для этой функции, log_softmax (), ни в одном из пакетов.

Может кто-нибудь объяснить, как это реализовано, или, что лучше, указать мне на соответствующий исходный код?

Answer 1

• Числовая ошибка:

>>> x = np.array([1, -10, 1000])
>>> np.exp(x) / np.exp(x).sum()
RuntimeWarning: overflow encountered in exp
RuntimeWarning: invalid value encountered in true_divide
Out[4]: array([ 0.,  0., nan])

Существует 2 способа избежать числовой ошибки при вычислении softmax:

• Нормализация опыта:

def exp_normalize(x):
    b = x.max()
    y = np.exp(x - b)
    return y / y.sum()

>>> exp_normalize(x)
array([0., 0., 1.])

• Лог Сумма Exp

def log_softmax(x):
    c = x.max()
    logsumexp = np.log(np.exp(x - c).sum())
    return x - c - logsumexp

Обратите внимание, что разумным выбором для обоих b, c в приведенной выше формуле является max (x). При таком выборе переполнение из-за exp невозможно. Наибольшее число, возведенное в степень после сдвига, равно 0.

Answer 2

Вы можете найти одну из реализаций ЦП здесь и векторизованную версию здесь (это версия журнала, которая вызывается из vec_host_softmax_lastdim).

Вы можете найти реализацию CUDA здесь , которая затем вызывает softmax_warp_forward.

Все они похожи, только синтаксис отличается. Как видите, обычно есть флаг, который определяет, будет ли softmax вычисляться с использованием журнала. То есть LogSoftMax вместо SoftMax.