May 13, 2024

TransformerFAM: Feedback attention is working memory

В статье разработана новая архитектура трансформеров TransformerFAM (Feedback Attention Memory) на примере обратной связи марковской цепочки в RNN, реализуют обратную связь с помощью механизма attention.

Source: Arxive, YouTube YannicKilcher


Для начала приведем схему работы методов Sliding Window Attention (SWA) и Block Sliding Window Attention (BSWA) на картинках ниже для проведения аналогии с новым методом и понимания новизны архитектуры TransformerFAM.

В работе предлагают новый метод сбора контекста, основываясь на текущем блоке токенов и на двух memory segments с информацией о прошлом контексте.

Если более просто объяснять работу TransformerFAM, то это происходит следующим образом - есть блоки токенов текста (красные), с которых получают токены query (зеленые). Каждый query токен ссылается на текущий токен текста и прошлый (в зависимости от значения memory segment). Но в работе вводят еще один токен (назовем его feedback токен и он голубого цвета), который является обучаемым (!) и тоже вбирает в себя информацию токенов текста от всего блока. Далее процесс повторяется, но следующий feedback ссылается уже не только на токены текста, но и на прошлый feedback токен. Feedback токен по сути является hidden state токеном (пасхалка RNN) и вбирает в себя всю прошлую информацию текста. Используя backprop, по сути feedback токен позволяет держать в себе бесконечную информацию контекста (не путать с контекстным окном!).

Ниже приведены алгоритмы работы TransformerBSWA и TransformerFAM для более подробного сравнения.

Novelty

Без сомнения, алгоритм является новым в своей сфере, однако в последнее время наблюдается массовое возвращение к теме рекуррентных сетей, поэтому стоит ожидать дальнейших попыток применения аналогичных методов в трансформерах, механизме attention и других технологиях.

Results

TransformerFAM имеет близкую к 100% точность на PassKey Retrieval тесте. В этом тесте в начале предъявляется ключ доступа, затем следует очень длинный заполняющий контекст. Наконец, задается вопрос о ключе. По результатам, TransformerFAM выдает верный ответ даже при длине контекста выше 200 000 токенов.

TransformerFAM имеет более высокие scores на Long Context Tasks, чем TransformerBSWA с разными значениями memory segments.