Суть эксперимента: от Squared ReLU к «Untied»
В стандартном трансформере функция FFN (Feed-Forward Network) часто использует ReLU. Исследователи Amy_ и Michael Bukatin предложили модификацию, вдохновленную работой Primer, где квадратичная активация ReLU(xW)² показала лучшую эффективность обучения. Вместо возведения в квадрат одного выхода, авторы разделили процесс на два независимых потока (branches) с разными весовыми матрицами Wa и Wb, которые затем перемножаются поэлементно:
Этот подход, названный Untie Squared ReLU, является более общей формой квадратичной активации. Если Wa = Wb, вариант становится эквивалентным классическому Squared ReLU. Идея продолжает исследования мультипликативных нейронных единиц как компактного способа представления полиномиальных взаимодействий.
Настройки и метрики
Эксперимент проводился на архитектуре modded-nano-gpt (12 слоев, d_model=768, 6 attention heads) на датасете FineWeb-10B. Обучение выполнялось на двух GPU NVIDIA H100 с использованием оптимизатора Muon. Ключевые параметры:
- Шаги оптимизации: 1,750
- Размер скрытого слоя FFN: 3,072 (для базовой линии) и 2,048 (для сравнения по параметрам)
- Количество параметров FFN: Базовая линия (Squared ReLU) имеет 4.72M параметров, новый вариант — 7.08M (рост в 1.5 раза из-за добавления второго up-projection).
Результаты: Потери, сходимость и вычисления
После 1,750 шагов обучения вариант Untie Squared ReLU достиг уровня валидационных потерь, практически идентичного базовой линии Squared ReLU. Однако это потребовало на ~18% больше машинного времени (wall-clock time) на шаг.
Интересный феномен наблюдался в косинусном сходстве матриц Wa и Wb. Изначально случайная инициализация в высокоразмерном пространстве давала значения, близкие к ортогональности. В ходе обучения сходство быстро росло и стабилизировалось вокруг ~0.77. Однако последующий эксперимент с намеренной инициализацией (0 и 0.999) показал, что нет универсального значения сходимости: при высокой начальной схожести матрицы «разъединялись» в процессе обучения.
Профилирование FLOPs подтвердило теоретические расчеты: добавление третьего умножения матриц (вместо двух в базовой линии) увеличило вычислительную нагрузку ровно в 1.5 раза.
| Метрика / Параметр | Squared ReLU (Baseline) | Untie Squared ReLU (Variant) |
|---|---|---|
| Параметры FFN (на слой) | 4.72M | 7.08M (+50%) |
| Вычислительная сложность (FLOPs) | 1.392 TFLOPs | 2.088 TFLOPs (+50%) |
| Wall-clock time (на шаг) | База (100%) | +18% медленнее |
| Косинусное сходство Wa/Wb | N/A (одна матрица) | Сходится к ~0.77 (при случайной инициализации) |
| Валидационная потеря (после 1750 шагов) | Базовый уровень | Практически идентична базовому уровню |
Почему это важно?
Результаты показывают, что «развязывание» квадратичной активации не дает преимущества в скорости или качестве сходимости на коротких дистанциях (1.75k шагов), но требует значительных вычислительных ресурсов. Однако, поскольку базовая линия была предварительно настроена (co-adapted learning rate, warmup), а новый вариант обучался без тонкой настройки гиперпараметров, окончательный вердикт о эффективности архитектуры требует более длительных экспериментов. Тем не менее, работа демонстрирует интересный механизм самонастройки весовых матриц в мультипликативных блоках.
Источник: LessWrong ↗