FFF — это FF, только быстрее
Основная часть параметров LLM относится к слоям прямого распространения (FF, feedforward). Но для выхода не нужны все нейроны этих слоев. В реальности на инференсе играет роль только небольшая доля нейронов FF. Отсюда возникает идея оптимизировать каким-то образом этот процесс и не тратить лишние вычисления на нейроны, которые потом никак не повлияют на выход.
Авторы из ETH Zurich предложили сеть UltraFastBERT, вариант BERT, который используют 0,3% его нейронов, но даёт сравнимые результаты. На инференсе каждого слоя используется всего 12 из 4095 нейронов.
UltraFastBERT построена почти полностью на фреймворке crammedBERT. Единственное изменение — структура внутренних слоев. FF трансформера заменили на быстрые слои прямого распространения (FFF, fast feedforward).
Сами FFF предложили те же авторы в прошедшем августе. Цель была разбить ставшие слишком тяжеловесными FF слои в трансформерах на модули. Первоначальная идея смеси экспертов тоже была как раз про это. Там FF разделили на отдельные блоки-эксперты и дополнительные гейты, которые определяли, выходы каких экспертов брать с какими слоями. То есть итоговый инференс становится быстрее за счет того, что используются только какая-то часть блоков-экспертов. Ускорение получается линейным. FFF же обещает экспоненциальное ускорение.
FFF разделяет вход на отдельные области с помощью дерева и параллельно обучает блоки, отвечающие за эти области. Отдельная часть узлов отвечает за то, какая смесь блоков (листьев дерева) пойдет в выход. Изначально границы областей размытые, но по мере обучения становятся всё более жёсткими и в конце концов выбирается только один лист. Причем за логарифмическое время.
Переход к жестким границам происходит последовательно при выборе одного из двух “детей” родительского узла. В зависимости от конечной функции потерь, оптимайзер обновляет веса детей и родителя. Со временем веса становятся всё более различающимися, а выбор всё более черно-белым. FFF глубиной d и шириной листа l обучается на всех скрытых нейронах (их получается l*2^d), а чтобы рассчитать выход, нужно будет пройти по одной ветке, то есть использовать всего d нейронов.
В каком то смысле листья в FFF это аналог экспертов в MoE, но ключевая разница вот в чем: FFF использует тот факт, что отдельные части входа активируют отдельные листья. Дерево FFF разделяет пространство входа, а не просто случайно выбирает самого “громкого” эксперта. Разделение входа и само по себе может быть полезным. Например, для интерпретируемости.
Чтобы опробовать FFF в деле, авторы взяли BERT-подобный фреймворк из 3072 нейронов в FF (остальные слои отбросили, так как их вклад в скорость не сильно заметен). 3072 не получится поделить на бинарное дерево, поэтому авторы взяли 4095 нейронов в расчете на дерево с максимальной глубиной 11.
Исходные FFF для этой задачи немного упростили. Во-первых, стерли разницу между листовыми и нелистовыми узлами. Для всех узлов используется одна и также GeLu, у всех узлов одинаковые веса. Во-вторых, зафиксировали размер листа равный одному. В-третьих, добавили возможность использовать параллельно несколько деревьев. Результаты отдельных деревьев суммируются.
Авторы построили 13 разных UltraFastBERT с разной шириной листьев и глубиной деревьев — от 3072х0 до 1х11. На бенчмарке GLUE результаты разных для этих разных сетей, но все примерно совпадают с crammedBERT (попробовали и 3072 и 4095 нейронов). А если нет разницы, зачем платить больше?
Все сети обучали в течение одного дня на одном A6000 GPU. Последнюю модель, с приставкой -long обучали в два раза дольше. С увеличением глубины результаты становились хуже.
UltraFastBERT использует для инференса только 1/341 нейронов BERT c 4095 нейронами. Но ускоряет работу “всего” в 78 раз. Почему не в 341? Ответ в некоторой степени парадоксальный. Потому что используется условное умножение матриц — строки входа умножаются на столбцы весов по очереди, а столбец выбирается в зависимости от предыдущего выхода. И хотя умножение плотных матриц, которое лежит в основе стандартных FF, гораздо более трудоемкий процесс, он очень хорошо оптимизирован и ускорен алгоритмически. А условное умножение матриц слишком редко встречается, поэтому приходится проводить его бесхитростно, напрямую, за счет снижения скорости.
