🔬 Метод

Чтобы не учить на триллионах токенов с нуля берут за основу предобученную модель, в данном случае из семейства Qwen-2.5.

Определение позиций блоков, где необходим attention

На первом этапе учат суперсеть (как в ProxylessNAS) с разными опциями блоков - как исходным attention, так и линейными по длине альтернативами (Mamba, GLA, DeltaNet, Gated Delta Net, RWKV7).

Суперсеть дистиллируется на логитах исходной модели, на каждом forward pass сэмплируется один из возможных путей.

Затем с помощью beam search ищут конфигурацию, оптимизирующую некоторую целевую метрику - лосс или точность на бенчмарке. И эту модель дообучают на большем количестве данных.

Замечают следующее:

📌 Для MMLU/retrieval полный attention необходим лишь в паре блоков.
📌 Важные блоки для MMLU/retrieval, вообще говоря, не пересекаются
📌 Обь единения блоков для MMLU/retrieval почти достаточно для задач по математике

Выводы справедливы для разных архитектур линейных слоев.

Выбор блока

Из рассмотренных вариантов по соотношению скорость/качество лушче всего себя показывает Gated DeltaNet, и далее используется в качестве основного.

Новый attention блок

Авторы предлагают новый attention блок - JetBlock. В его основе динамическая (зависящая от входа) свертка, примененная к value, и статические свертки на query и key. И нечто под названием time-mixing + gating в конце.

Этот блок по скорости такой же примерно, как альтернативы, но лучше по качеству.

Зачем дотюнивают параметры размерности ключей/значений / числа голов JetBlock для максимизации скорости ⚡️.