Трансформеры без нормализации: как обойтись без сложных слоев

Я тут наткнулся на интересную статью от довольно именитых авторов: Jiachen Zhu, Xinlei Chen, Kaiming He, Yann LeCun и Zhuang Liu. Вот ссылка на саму статью: Transformers without Normalization. Скажу сразу — работа поразила меня своим подходом!

Итак, говорят, что в трансформерах можно избавиться от стандартных слоев нормализации. Вместо этого можно использовать поэлементный гиперболический тангенс — Dynamic Tanh (DyT), который выглядит как DyT(x)=tanh(αx), где α — это обучаемый параметр. И, внимание, производительность при этом не падает, никаких особых манипуляций с гиперпараметрами не требуется!

Когда в 2015 году на сцену вышла BatchNorm, она действительно произвела фурор. Вроде как задумывалась одним образом, а сработала по-другому, но все же как-то улучшила результаты. Примерно в том же духе позже появились и другие варианты нормализации, такие как LayerNorm и RMSNorm, которые теперь заполонили мир трансформеров. Все эти механизмы помогают справляться с внутренними изменениями распределения, насыщением активаций и тому подобным вещам, хотя на самом деле их реальные функции все еще можно оспаривать. Но их эффективность не может не впечатлять: качественные результаты, лучшее сходимость и более глубокие архитектуры — кто с этим не согласится?

Вкратце, нормализация работает таким образом: убирает среднее, делит на дисперсию (иногда игнорируя что-то из этого), а затем добавляет обучаемые параметры для «сдвига» и «масштабирования». Но стоит ли это того, когда можно попробовать что-то совершенно иное?

Не думаю, что стоит сбрасывать со счетов и более радикальные идеи, такие как SELU, хотя и со своим звоном они как будто не прошли тест на популярность.

Возвращаясь к исследованию, авторы подчеркивают, что нормализация и не нужна. DyT демонстрирует, что можно обойтись без сложных процедур, причем без необходимости считать какую-либо статистику активаций между слоями. Этот подход намного проще, поскольку не требует хранить информацию для инференса, как это было в случае с BatchNorm. Черт побери, зачем нам все это, если можно сделать чуть уму непостижимо проще?

Авторы проанализировали показатели трансформеров с LayerNorm — ViT-B, wav2vec 2.0 и DiT-XL. Открыли довольно интересные зависимости: на первых слоях вход и выход связаны линейно, а на глубоких слоях наблюдаются S-образные кривые с широкой линейной областью. Вот тут-то и пришла идея реализовать DyT. Полная формула выглядит так: DyT(x) = γ * tanh(αx) + β, где α, γ и β — обучаемые параметры.

В этом подходе можно рассматривать DyT как функцию активации, но авторы сосредоточили внимание исключительно на замене LN слоев. Провели ряд экспериментов на различных моделях, и результаты, согласно им, были весьма сопоставимыми. Маленькие колебания, как правило, были незначительными: чуть хуже, чуть лучше, но в целом все на уровне. Интересный аспект заключался в том, что инференс на DyT происходил почти в два раза быстрее на уровне слоя!

В залипательных абляциях, которые они провели, заменяли tanh на другие функции, такие как hardtanh и sigmoid, но так и не смогли превзойти tanh. Без обучаемого параметра α тоже результаты уменьшались.

Также наблюдалась интересная зависимость параметра α в процессе обучения — он явно следовал зависимости 1/std от активаций. После обучения у них сохранялся сильный корреляционный статус с 1/std входных активаций, и, что любопытно, у более глубоких слоев активации демонстрировали больее высокую дисперсию.

Всё это привело их к сравнительному анализу с другими методами, такими как Fixup, SkipInit и σReparam. На ViT и MAE DyT показывал великолепные результаты, превосходя оставшиеся методы.

Что касается инициализации α... Х