Tanh Attention: мини-эксперимент⁠⁠

Введение и обсуждение идеи. Выявление проблем обычного внимания.

Недавно решил попробовать обучить MultiheadAttention способом стохастического градиентного спуска на последовательностях длинной 512 с целью убедиться, что я прав в утверждении, что обычное мультиголовое внимание не сможет обучиться из-за градиентов и некоторых конфликтов. В итоге, после примерно двух часов работы за компьютером, точность (accuracy) выросла с 0,0 до 0,002. По привычке я первым делом полез проверять градиенты и обнаружил, что их значения находятся в диапазоне 0.0002 - 0,000008.

Однако с архитектурой всё в порядке: она полностью соответствует стандарту - Pre-Norm и Skip Connection на месте. Тогда я решил поискать проблему в самом MultiheadAttention. После недолгих размышлений я решил, что корень проблемы — в SoftMax: а именно, в способе распределения вероятностей по токенам и в расчёте градиента.

Во-первых, давайте разберёмся, как обычный MultiheadAttention нормализует входные данные перед SoftMax и помогает ли это в принципе. В стандартном MultiheadAttention нормализация - это просто деление на квадратный корень от размера головы. Это делается именно так, потому что при вычислении произведения QK.T (веса которого распределены в диапазоне [−1;1]) каждое результирующее значение по статистике будет примерно равно корню из размерности головы. Это аналогично тому, что при 60 подбрасываниях монетки выпадет примерно по 30 орлов и решек - чистая статистика. Однако деление на √d_k - это лишь "костыль", поскольку на практике значения QK.T могут быть как больше, так и меньше этой величины. А так как в процессе обучения веса постоянно меняются, то вместе с ними масштабируются и эти значения. Здесь и кроется проблема: если одни элементы произведения становятся больше, а другие - меньше, то SoftMax распределяет основную вероятность лишь на несколько "больших" элементов, оставляя для остальных исчезающе малые доли. Когда SoftMax выдаёт ничтожно малые значения для тех или иных токенов, градиент для них стремится к нулю (по формуле p⋅(1−p)), и при последующем умножении на другие градиенты это полностью обнуляет градиент для данного токена. Когда градиент, на который мы обновляем веса, оказывается значительно меньше самих весов, обучение практически останавливается - мы «стоим на месте». Более того, когда модель опирается лишь на несколько токенов, присваивая остальным исчезающе малые вероятности, она фактически теряет контекст, а вместе с ним и потенциально важные детали. Это серьёзно усугубляет проблему обучения: другие значимые токены вносят слишком мелкий вклад, который попросту подавляется доминирующими. В итоге, такое внимание неспособно эффективно обучаться стандартным SGD - ограничением выступает сам градиент.

Но что же тогда делать? И почему на практике все работает?

На практике всё работает именно так, потому что никто не использует одну голову для длинных последовательностей. Например, в DeepSeek V4 PRO используется 128 голов и множество слоёв внимания. Да, такая модель тоже игнорирует много токенов. Да, у неё тоже при обучении градиенты с большой вероятностью где-то затухали. Но никто не использует SGD. А несколько слоёв внимания компенсируют потерю, создавая дополнительные параметры и, грубо говоря, дополнительно выявляют признаки из предыдущего слоя внимания. То есть все используют Adam / AdamW, который подходит к обучению "с умом", а не просто умножает градиент на learning rate. И именно благодаря ему все трансформеры работают отлично. Я же использую SGD только в рамках эксперимента.

А если хочется стабильных градиентов? А если компьютер не тянет Adam / AdamW, но хочется такое внимание, которое по качеству будет похоже на MultiheadAttention или даже превзойдёт его, но при этом будет стабильно обучатся и с SGD? Для начала давайте выявим самую главную проблему, почему обычное внимание не обучается методом SGD. А вот и причина: как я говорил ранее, SoftMax делает градиенты для токенов с низкой вероятностью очень маленькими. И сеть просто не способна учиться стабильно из-за этих мелких градиентов. А что если на других токенах внимание просто смотрит на два этих токена, а на остальные плюет? Это будет ужасным конфликтом для обучения генеративных моделей, классификаторов и для многих других задач.

Немного подумав, я сначала пришёл к выводу поставить Sigmoid вместо SoftMax, а вместо деления произведения QK.T на √d_k делить на std(QK.T) отдельно для каждого вектора, независимо от других. Как раз Sigmoid не зависит от всех чисел в векторе, а работает локально с каждым элементом отдельно, что как раз и исправляет проблему затухания градиента, свойственную Softmax. Но сразу же откинул эту идею по следующим причинам:

1. Sigmoid (по моим расчётам) имеет рабочий диапазон [−0,5; 0,5], в котором градиент почти максимален, а выходы не крутятся вокруг 0.5. Но деление произведения QK.T на std(QK.T) приводит к тому, что среднее отклонение становится равным единице, что вдвое превышает рабочий диапазон Sigmoid. В результате градиенты Sigmoid оказываются далеки от своего максимума.

2. После того, как Sigmoid выдал свои выходы, мы их умножаем на V. Вот тут и кроется вторая проблема. Так как Sigmoid выдает значения в диапазоне (0; 1) и никогда не равен ровно 0 или 1, то V умножается строго на положительные числа. Но главная беда в том, что сумма выходов Sigmoid не ограничена единицей - в отличие от Softmax. И эти значения будут расти с длиной последовательности. Чем длиннее вход, тем больше токенов вносят свой вклад, и итоговые значения в выходной матрице V становятся огромными. Конечно, с градиентами для самого внимания проблем не будет, но для следующего слоя - будут. Это исправляется тем, что мы ставим LayerNorm или RMSNorm перед следующим слоем, но все же - как проблема, это проблема.

3. И пожалуй, самая главная проблема - максимально возможный градиент функции. Sigmoid имеет максимальный градиент = 0,25. А по моим расчётам лучший диапазон для градиента очень маленький - [-0.5; 0.5].

Немного подумав и перебрав несколько вариантов замены Sigmoid я остановился на Tanh, чему и посвящена эта мини-статья.

Взял я именно его по следующим причинам:

1. Tanh в отличие от Sigmoid по моим расчётам имеет больший диапазон - [-1; 1]. В этом диапазоне градиент почти в два раза больше максимально возможного градиента Sigmoid (0,25). За счёт большего в два раза диапазона он полностью подходит под нашу нормализацию, которую я поставил вместо QK.T.

2. Также Tanh в отличие от Sigmoid имеет максимальный градиент в 4 раза больше, чем максимальный градиент Sigmoid. Это очень хорошее дополнение. Если брать рабочий диапазон Tanh (по моим расчётам) [-1; 1], то минимальный градиент будет примерно равен 0,5, а максимальный - 1. А у Sigmoid в рабочем диапазоне (по моим расчётам) [-0,5; 0,5] минимальный градиент примерно равен 0,23, а максимальный - 0,25. Лучше брать тот вариант, где градиент больше. Особенно для SGD.

3. Больших значений после произведения выходов Tanh на V будет намного меньше, чем с Sigmoid. Это можно объяснить тем, что у Tanh выход имеет диапазон [-1; 1] - там есть отрицательные числа, а у Sigmoid [0; 1] - отрицательных чисел быть не может. То есть теперь мы умножаем на V не только положительные числа, но и отрицательные, которые смогут добавить дополнительные варианты для подбора весов и помогут при сложении помешать огромному выбросу, либо сделать его не таким огромным. Но, это не отменяет важность LayerNorm / RMSNorm на выходе из внимания.

Была мысль взять ReLU или SiLU, но сразу же я ее откинул, так как ReLU и SiLU хоть и имеют хорошие градиенты на положительных входах, но это гарантированно приведет к огромным значениям в V. Также это может понизить качество модели в плане того, что мы просто обнуляем или "гасим" мелкие значения, тем самым уменьшая варианты для подбора весов градиентом.

Разобравшись, почему именно Tanh - давайте окончательно соберем весь этот пазл.

Как я уже говорил, рабочий диапазон Tanh — [-1; 1]. В этом диапазоне градиент находится в идеальном положении. А так как деление на √d_k - это просто костыль, который может не работать (как я уже объяснял), мы будем делить на std(QK.T), где std вычисляется отдельно для каждого вектора в матрице QK.T. Далее - просто матрично умножаем выходы Tanh на V. И после - конкатенируем результаты всех голов. Следующее действие очень важно для стабильности градиентов: после конкатенации голов необходимо поставить LayerNorm или RMSNorm, так как даже при использовании Tanh, хоть вероятность взрывов и меньше, они всё ещё возможны.

Итак, формула для одной головы принимает следующий вид:

А формула выхода из Tanh Attention примет такой вид:

Тестовое обучение Tanh Attention методом SGD на последовательности длинной 512. Подводим итоги.

После того как я реализовал Tanh Attention, я собрал две модели:

1. Модель, с стандартным MultiheadAttention (1 слой внимания, 4 головы)

2. Модель с Tanh Attention (1 слой внимания, 4 головы)

Архитектуры, датасет, данные для обучения и код абсолютно одинаковы в обоих экспериментах. Обучение проводилось методом SGD с одинаковым learning rate = 0.001. Ниже представлены результаты, а также графики распределения градиентов для Tanh Attention и стандартного MultiheadAttention.

Внимание-Accuracy-Время обучения-LOSS-Диапазон градиентов

Tanh Attention-21.22 %-2 часа 37 минут-6.58-0.00604-0.029

MultiheadAttention-0.002 %-2 часа 39 минут-9.84-0.000008-0.0002

Фото:

Итог: в этой статье мы разобрали проблемы стандартного механизма внимания и предложили способ их решения. Эксперименты показали, что обычное внимание вообще не обучается с помощью SGD, тогда как Tanh Attention стабильно сходится.

Дополнительно я провёл эксперимент с оптимизатором Adam. На Tanh Attention он показал такую же скорость сходимости, как и SGD. Обычное внимание с Adam обучалось почти так же хорошо, как и Tanh Attention, хотя и немного медленнее — но это некритично.

Общий вывод: замена Softmax на Tanh улучшает градиенты и добавляет вариативность модели. Кроме того, Tanh Attention с высокой вероятностью лучше удерживает контекст за счёт более равномерного распределения весов.

P.S. статью так же написал на хабре, но с модерацией не известно.