Теория, анализ и лучшие практики сигмоидального self-attention

Внимание является ключевой частью архитектуры трансформера. Это отображение последовательности в последовательность, которое преобразует каждый элемент последовательности во взвешенную сумму значений. Веса обычно получают как функцию softmax скалярных произведений между ключами и запросами. В недавних работах исследовались альтернативы функции внимания softmax в трансформерах, такие как ReLU и сигмоидная активация. В этой работе мы возвращаемся к сигмоидной функции внимания и проводим углубленный теоретический и эмпирический анализ. Теоретически мы доказываем, что трансформеры с сигмоидной функцией внимания являются универсальными аппроксиматорами функций и выигрывают от улучшенной регулярности по сравнению с функцией внимания softmax. С помощью подробного эмпирического анализа мы выявляем стабилизацию больших начальных норм внимания на ранних этапах обучения как решающий фактор для успешного обучения моделей с сигмоидной функцией внимания, превосходящих предыдущие попытки. Мы также представляем FLASHSIGMOID, аппаратно-ориентированную и эффективную с точки зрения памяти реализацию сигмоидного внимания, обеспечивающую 17% ускорение ядра вывода по сравнению с FLASHATTENTION2 на графических процессорах H100. Эксперименты в области языка, зрения и речи показывают, что правильно нормализованное сигмоидное внимание демонстрирует такую же высокую производительность, как и внимание softmax, в широком диапазоне областей и масштабов, чего не удавалось в полной мере достичь предыдущим попыткам реализации сигмоидного внимания. Наша работа объединяет существующие решения и устанавливает лучшие практики для сигмоидного внимания в качестве прямой замены softmax в трансформерах. https://arxiv.org/abs/2409.04431