Как мне кажется, Mixture-of-Depths – одна из самых интересных недавно вышедших статей, так как посвящена она довольно очевидной проблеме, что все токены в механизме оттеншена потребляют одинаковое количество компьюта, хотя польза от них не одинаковая. Конкретно эта статья фокусируется на том, как тратить меньше вычислений на токены, которые особенно и не нужны в предсказании

Идея очень простая и поэтому уважаемая: перед каждым трансформер-блоком (в котором находится аттеншн и MLP) работает роутер, который выдает для каждого токена входной последовательности какой-нибудь вес. Дальше в блок поступают и обрабатываются только top-k токенов c наибольшими весами, а все остальные через residual connection обходят блок без каких-либо изменений. Соответственно, механизм оттеншена в итоге требует k^2 вычислений, где k << числа токенов во входной последовательности

Поскольку k задается самим пользователем, размеры матриц внутри аттеншена нам изначально известны, и мы точно можем посчитать, сколько компьюта урежет заданное значение k

Как корректно заметил один человек в реплаях Твиттера, в таком подходе если какой-то токен был проигнорирован в конкретном трансформер-блоке, то дальше все последующие токены не смогут аттендиться к нему, так как этот токен не попадет в KV-кеш. То есть он навсегда остается таким “слепым пятном” в этом блоке, но в следующих блоках он все равно может попасть в top-k и сыграть роль там. В теории возможно, что какой-то совсем бесполезный токен занулится во всех блоках и не окажет вообще никакого влияния на генерацию

В экспериментах лучшая вариация MoD оказалась та, где k был равен 256 (12.5% от всех входных токенов попадали в блок) и где роутер стоял перед каждым вторым трансформер-блоком. Помимо того, что у этой модели лосс был даже ниже, чем у стандартной модели с таким же количеством параметров, она еще и на 66% быстрее совершала шаг инференса. Это по сути и подтверждает, что куча операций в ванильном трансформере излишняя и не дает прироста качества

Довольно приколькую идею скрестить Mixture-of-Depths и Mixture-of-Experts предложили также сами авторы – просто можно заменить одного эксперта на identity function и готово 👍 По-моему очень элегантно

Еще хорошую идею развития этого ресерча закидывают в конце в discussion – можно предположить, что некоторые токены очень полезны как keys, но не очень полезны как queries или наоборот. В общем кажется, что эта команда ресерчеров опубликует позднее что-то еще про более сложный роутинг