Реализация разреженности с помощью функции Top-K
Основная операция, выполняемая Q-Sparse, - это применение функции разреженности Top-K к входному тензору.
Конкретно, архитектура Transformer использует линейные слои nn.Linear (умножение матриц) для проекции в слоях внимания и прямой связи, что можно представить как Y=X·W^T. (где X - входной тензор, W - его веса, Y - выходной тензор)
В Q-Sparse для входного тензора активации X сначала вычисляется его абсолютное значение |X| и выполняется сортировка, чтобы найти K элементов с наибольшим абсолютным значением.
K здесь - заранее заданный гиперпараметр, определяющий степень разреженности.
Затем Q-Sparse создает бинарный тензор-маску M той же формы, что и X, устанавливая соответствующие позиции в M равными 1 для K элементов с наибольшим абсолютным значением в |X|, а остальные позиции - равными 0.
Затем входной тензор X умножается на тензор-маску M поэлементно (произведение Адамара), получая разреженный тензор X_sparse.
В процессе прямого распространения разреженный тензор X_sparse заменяет исходный входной тензор X в последующих вычислениях (например, умножении матриц).
Поскольку большинство элементов в X_sparse установлены в ноль, это значительно снижает вычислительные затраты и требования к пропускной способности памяти.
В процессе обратного распространения Q-Sparse использует прямой оценщик (Straight-Through Estimator, STE) для вычисления градиента функции Top-K.
В традиционном обучении обычно необходимо вычислять градиент функции потерь по отношению к параметрам сети и использовать градиентный спуск для обновления параметров с целью минимизации потерь.
Однако когда в сети присутствуют недифференцируемые операции, такие как квантование или Top-K, возникают проблемы с вычислением градиентов, поскольку градиент выхода по отношению к входу для этих операций равен 0 в большинстве точек, что приводит к неэффективному распространению градиента.
STE решает проблему исчезновения градиента, напрямую передавая градиент тензору перед разреженностью.
В обычном обратном распространении градиент функции потерь L по x равен ∂L/∂x=∂L/∂y⋅∂y/∂x, но его нельзя вычислить напрямую из-за недифференцируемости.
Решение STE заключается в вычислении только градиента функции потерь по отношению к разреженному тензору y, а затем копировании его напрямую в исходный тензор x, то есть использовании ∂L/∂y в качестве оценки для ∂L/∂x.
Для слоев прямой связи Q-Sparse использует квадратичную функцию ReLU вместо обычной функции активации ReLU, квадратичная операция может дополнительно повысить разреженность активации (⊙ обозначает произведение Адамара).
Кроме того, для адаптации к квантованным моделям Q-Sparse квантует входной тензор перед применением разреженности Top-K, чтобы обеспечить совместимость операции разреженности с квантованным представлением. Функция представлена следующим образом:
Здесь ε - небольшая константа, используемая для избежания деления на ноль.
В частности, для 1-битного квантования весов Q-Sparse использует следующую функцию квантования, где α - среднее абсолютное значение тензора весов W.
60% активационных параметров достигают того же эффекта
Сравнительные эксперименты показывают, что Q-Sparse значительно превосходит предыдущие методы ReLU как по степени разреженности, так и по производительности модели.
Для оценки конкретного эффекта Q-Sparse авторы оценили его производительность в задачах обучения с нуля, продолжения обучения и тонкой настройки.
Эксперименты по обучению с нуля использовали модель Llama, результаты показали, что на моделях 700M и 7B Q-Sparse с 70% top-K (т.е. 40% общей разреженности) может достичь потерь при обучении, сопоставимых с плотным базовым уровнем.
Цель продолжения обучения - разрежение плотных моделей, объектом эксперимента была Mistral-7B.
В результате, при активационных параметрах 2,9B и 3,8B, оценки модели на наборах данных ARC, MMLU и других не показали значительного снижения.
В экспериментах по тонкой настройке для моделей Qwen-7B и Mistral-7B Q-Sparse показал результаты, аналогичные продолжению обучения, достигнув производительности, очень близкой к плотным моделям, с использованием около 60% активационных параметров.
Эти результаты означают, что при одинаковой производительности разреженные активационные модели могут значительно сократить количество активационных параметров в процессе вывода по сравнению с плотными моделями, тем самым снижая количество потребляемых FLOPS.
Для квантованных моделей команда применила Q-Sparse к своей собственной модели BitNet b1.58 и провела обучение и оценку на нескольких наборах данных.
Можно видеть, что при масштабах 700M и 7B скорость сходимости и конечное значение функции потерь квантованных моделей с использованием Q-Sparse сопоставимы с квантованными моделями без использования Q-Sparse (BitNet b1.58).
Это показывает, что Q-Sparse может быть беспрепятственно интегрирован в квантованные модели без значительного влияния на обучение и сходимость модели.
На основании этого авторы считают, что сочетание Q-Sparse с методами квантования может дополнительно повысить эффективность больших языковых моделей на этапе вывода.
Обнаружение нового "Scaling Law" для оптимизации вывода
Помимо оценки производительности этих моделей при использовании разреженной активации, авторы также исследовали взаимосвязь между производительностью модели, масштабом и степенью разреженности, и сделали некоторые новые открытия.
Закон масштабирования производительности для моделей с разреженной активацией: Авторы обнаружили, что, подобно плотным моделям, производительность моделей с разреженной активацией также следует степенному закону масштабирования.
Конкретно, при заданной степени разреженности S значение функции потерь L(N,S) при сходимости модели может быть аппроксимировано следующей формулой: