Трансформеры, несмотря на всю их мощь, всё ещё легко заболтать до беспамятства, то есть дать им на вход такой длинный контекст, что они не смогут удержать его весь. В целом, это наверное главный, если не единственный их недостаток и потому улучшить память трансформеров — главная, если не единственная мечта их авторов. И тут пускают в ход всё, меняют механизм внимания, создают всевозможные гибриды с RNN в поисках успешного подхода. 

Один из вариантов — Sliding Window Attention, в котором при расчете следующего токена модель обращается к некоторому количеству токенов из “окна” заданной ширины. Вариантов задания этого окна может быть много, некоторые описаны в этой статье 2020 года. 

Transformer-XL предлагает дополнить трансформеры рекуррентностью на уровне блоков. В статье описывается улучшение памяти на 80% по сравнению с RNN  и на 450% по сравнению с vanilla Transformer (в смысле длины запоминаемых связей). Этого добиваются с помощью дополнительных скрытых состояний, которые рассчитываются для каждого блока и передаются как дополнительный контекст при расчете следующего блока того же слоя. Важное замечание, что при этом переходе искусственно останавливают градиент, чтобы не перегрузить используемую память. 

В некоторых других вариантах (статья 1, статья 2) дополнительные скрытые состояния формируются на верхних слоях трансформеров и передаются на нижние. 

В апреле авторы из Google предложили архитектуру TransformerFAM (Feedback Attention Memory). В её основе — ещё один вариант рекруррентности в трансформере, хотя сами авторы используют совершенно другую терминологию — из области нейробиологии и психологии. О самой архитектуре ниже расскажем подробнее, но пока сделаем вслед за авторами отступление в нейробиологию.

Долгосрочная память человека — это аналог весов LLM. То, что выучено когда-то раз и навсегда и теперь используется. Однако для исполнения конкретных задач человек использует так называемую рабочую память. Это не очень емкая система, где мы храним временную информацию, которая должна быть под рукой именно сейчас. Рабочую память мы используем, например, при устном счете или обдумывании логической задачки. От кратковременной памяти рабочая отличается именно этим “под рукой”. В кратковременную память помещается информация, которую нужно просто хранить, а в рабочую — которой нужно манипулировать. Активация рабочей памяти происходит с помощью импульсов, которые ходят по одному и тому же кругу нейронов (звучит как RNN, но авторы такого сравнения не делают). Такого механизма в трансформерах нет. Именно эту рабочую память и попытались создать авторы TransformerFAM. 

В целом архитектура TransformerFAM не отличается каким-то принципиальным образом от Sliding Window Attention, а ещё очень напоминает Transformer-XL, только без остановки градиента. В классическом sliding window определяется конечное число токенов, к которым обращается внимание. Есть еще вариант Block Sliding Window Attention, в котором последовательность токенов разбивается на блоки (тоже фиксированной конечной величины), и внимание обращается не к какому-то числу токенов, а к какому-то числу блоков. По факту это означает, что ширина “окна” будет всякий раз разной, но из какого-то набора, определяемого геометрией архитектуры.

Архитектуру BSWA авторы последней статьи взяли за основу и в качестве сравнения, и в качестве иллюстрации работы. В FAM токены всё так же обрабатываются блоками. Но к каждому блоку теперь добавляются специальные токены, которые, собственно и считаются “памятью”. Именно они дали название всей архитектуре, FAM. Теперь обращаться при работе с очередным токеном система будет не только к заданному количеству токенов своего блока, но и к этим токенам памяти предыдущего блока. В них содержится архив обработанной прошлой информации. Рассчитанные токены памяти FAM в свою очередь передаются в следующий блок того же слоя. 

 Схема из статьи выглядит довольно запутанно, поэтому есть смысл посмотреть на формулы, откуда сравнение с BSWA более наглядно. Впрочем, это, конечно, дело вкуса. 

Не вдаваясь в подробности обратим внимание на различие, просто чтобы понять устройство новой архитектуры. Справа теперь Q, K, и V внимания рассчитываются не только для каждого блока, но и для специальных токенов памяти FAM (те же матрицы, но со знаком F). В архитектуре TransformerFAM нету отдельной сети или даже отдельного компонента для расчета памяти. Функции QKV, SelfAttention и FF остаются теми же.

Авторы провели большую серию экспериментов, сравнили FAM с 8 другими моделями на нескольких заданиях.

FAM действительно показывает себя хорошо, где-то даже отлично по сравнению с другими LLM. Но авторы, к сожалению, не указывают какие LLM спрятаны за условными обозначениями. Неизвестно чем это объясняться, но очень жаль. Зато мы можем сравнить в явном виде FAM и BSWA. На тех же заданиях FAM почти неизменно лучше, как в большом, так и малом размере.

Ещё один интересный результат касается размера самого блока FAM, то есть количества токенов памяти. Сначала при увеличении их количества результат сети растет. Но 64 токена памяти оказываются оптимальным значением на двух из трех заданиях. Если увеличивать размер FAM блока дальше, то результат будет хуже. Но интересно здесь то, какой вывод делают сами авторы. В первой версии статьи (она доступна на том же arxiv) они приводят закон Миллера из психологии. Согласно ему, человек может запоминать в “рабочей памяти” 7 плюс-минус два элемента. Отсюда следовало заключение, что пик в 64 токена — это не баг, а фича. Впрочем от этого вывода авторы отказались и в последней версии его уже нет. 

Больше наших обзоров AI‑статей — на канале Pro AI.