Google снова пытаются изобрести альтернативу трансформерам: на этот раз они показали новую архитектуру Mixture-of-Recursions

Идея: рекурсивно переиспользовать одни и те же слои для повторной обработки некоторых токенов. На каждой новой итерации мы оставляем только те токены, которые требуют дополнительного рассмотрения. Получается какой-то новый вид ризонинга, при этом модель несколько раз думает только над наиболее сложными частями ответа, экономя компьют.

Теперь о том, как это работает:

➖ У нас есть стек слоев обычного трансформера, которые мы хотим переиспользовать несколько раз. Для каждого токена принимается решение: выпустить его сейчас или продолжить обрабатывать. Чем больше номер итерации – тем меньше токенов в ней остается. В конце получается, что над самыми сложными частями модель думала много раз, а над самыми простыми – один или пару.

➖ Есть варианты, как шерить слои (скрин 5). Можно в каждой итерации оставлять идентичные блоки вообще без изменений, а можно вставлять уникальные первый и последний слой, а середину оставлять одинаковой. В статье показали, что вторая вариация работает приятнее.

➖ Аналогично, можно по-разному решать, сколько итераций пройдет каждый токен (скрин 4). Вариант первый: в самом начале с помощью классификатора назначать каждому токену число N его итераций. Вариант два: в начале каждой итерации вставлять роутер, который будет решать, кто пойдет дальше. В первом случае обучение чуть стабильнее, но и там и там – свои танцы с бубном регуляризациями.

Ну и да, за счет того, что KV кэш мы не пересчитываем, вся эта архитектура остается достаточно эффективной. Например, по сравнению с ванильным трансформером, MoR: (1) выбивает те же метрики с в два раза меньшим числом параметров; (2) в два раза бодрее на инференсе; (3) снижает потребление FLOPs на 25 % на тех же объемах данных.

Получается довольно симпатичный баланс, и в статье даже пишут, что это новая Парето-оптимальность для трансформерных архитектур: и затраты, и качество оказываются одновременно лучше предшественников (скрин 6).

Радуемся и очень ждем уже на проде

Код | Статья