Добро пожаловать в CAMELoT

Большие языковые модели (LLM) сталкиваются с трудностями при обработке длинных входных последовательностей из-за высоких затрат памяти и времени выполнения. Модели с расширенной памятью стали многообещающим решением этой проблемы, но текущие методы ограничены её объёмом и требуют дорогостоящего повторного обучения для интеграции с новой LLM. В этой статье мы познакомимся с модулем ассоциативной памяти, который может быть связан с любой предварительно обученной LLM без повторного обучения, что позволяет ему обрабатывать произвольно длинные входные последовательности.

В отличие от предыдущих методов этот модуль ассоциативной памяти объединяет представления отдельных токенов в непараметрическую модель распределения. Эта модель управляется динамически путём надлежащего балансирования новизны и свежести входящих данных. Извлекая информацию из консолидированной ассоциативной памяти, базовый LLM на стандартных тестах достигает лучших результатов. Эта архитектура называется CAMELoT (Consolidated Associationive Memory Enhanced Long Transformer). Она демонстрирует превосходную производительность даже при крошечном контекстном окне в 128 токенов, а также обеспечивает улучшенное контекстное обучение с гораздо большим набором демонстраций.

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

Ассоциативная память (AM) — важный класс систем памяти. Основная функция AM — это связывание запроса с представлениями, хранящимися в памяти. В частности, для любого заданного запроса AM должна выбрать объединенный слот памяти с представлением, которое наилучшим образом соответствует запросу. Эти сохранённые представления должны кратко суммировать прошлый опыт и предоставлять ценные подсказки для будущих действий. В последнее время наблюдается значительный интерес к современным реинкарнациям сетей ассоциативной памяти. Существует большой объём литературы по консолидации памяти в нейронных сетях, а также различные локальные правила обучения, которые в вычислительном отношении дешевле, чем сквозное обучение методом обратного распространения.

В то же время большие языковые модели приобрели важное значение для многих практических приложений (чат-боты, генерация текста, доказательство теорем, количественное рассуждение и т. д.). Одним из важнейших параметров для LLM является допустимая длина входного контекста L.

Поддержка более длинного контекста позволяет включать более подробную фактическую информацию и работать с длинными документами или большими репозиториями кода. Это повышает производительность LLM в случаях, когда этот более длинный контекст доступен во время вывода. К сожалению, обычный механизм внимания трансформера масштабируется квадратично L² с увеличивающимся числом токенов, что делает увеличение длины контекста вычислительно сложным.

Кроме того, обучение современных LLM — это дорогостоящее мероприятие, которое обычно требует больших ресурсов. По этой причине оно обычно доступно только крупным конторам. LLM с открытым исходным кодом, хотя и хорошо работают на практике, часто имеют ограниченные контекстные окна, сталкиваясь с трудностями при выполнении задач NLP, включающих длинные документы. Эти ограничения поднимают вопрос: возможно ли разработать PnP модуль, который можно будет соединить с предварительно обученной LLM, что позволит ему обрабатывать неограниченно длинный контекст далеко за пределами L? При этом модуль должен быть менее затратным в вычислительном плане и не должен требовать переобучения или тонкой настройки базовой программы LLM.

Такой модуль был разработан в недрах IBM Research. Его функция заключается в консолидации информации о предыдущем контексте далеко за пределами текущего контекстного окна (длиной L) при аппроксимации вычисления внимания (для следующего предсказания токена) на полном контексте. Информация записывается в модуль AM с использованием локального правила записи, которое является вычислительно дешёвым и динамически отслеживает новизну и актуальность входящей информации.

Консолидация предполагает сжатие информации для сохранения. Так работает мозг, поскольку он имеет конечное количество нейронов, и то же самое верно для LLM, поскольку размеры LLM также должны быть ограничены.

Консолидированный контекст моделируется как набор непараметрических распределений, по одному на ключевое пространство каждого слоя LLM. Режимы этих распределений динамически консолидируются по мере того, как окно контекста проносится по длинному общему контексту, новые режимы создаются в соответствии с их обнаруженной новизной, а устаревшие режимы удаляются. Затем внимание к длинному контексту просто аппроксимируется путём извлечения режимов, наиболее близких к текущим скрытым состояниям контекста, и добавляет их в виде кэша ключ-значение.

LLM с поддержкой ассоциативной памяти (AM)

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

Консолидация (consolidation). Избыточная прошлая информация должна быть сжата и сохранена в блоке AM, при этом повторения должны сокращаться. Когда одно и то же понятие появляется в прошлом контексте несколько раз, расточительно хранить каждый отдельный экземпляр этого понятия в отдельном слоте памяти. Вместо этого все эти экземпляры должны быть объединены и сохранены только один раз.

Новизна (novelty). Новые понятия, не встречавшиеся модели в прошлом, должны быть обнаружены и сохранены в новом слоте памяти при их первом обнаружении. Новые слоты памяти могут быть впоследствии объединены с возможными будущими появлениями связанных понятий.

Актуальность (recency). В ситуациях, когда тема меняется, модель должна иметь возможность отбрасывать устаревшие слоты памяти, которые больше не являются полезными, если это требуется для включения дополнительных новых понятий, встречающихся после смены темы.

Модуль памяти оснащён операциями чтения и записи, поддерживающими извлечение информации из банка памяти и обновление банка памяти. С извлеченной информацией текущее контекстное окно LLM расширяется с помощью операции дополнения. Эти три желательных параметра являются основой CAMELoT.

Операция чтения

Когда контекстное окно длиной L обрабатывается через LLM, ключи и значения из каждого слоя (в более общем случае это может быть произвольное подмножество слоев) передаются в соответствующий модуль AM (по одному на слой с расширенной памятью). AM в каждом слое состоит из M слотов памяти, пронумерованных по индексу μ = 1, …, M. Каждый слот содержит две векторные переменные: ключи K_μ^mem и значения памяти V_μ^mem и две целочисленные скалярные переменные: счётчик c_μ (количество консолидированных случаев) и возраст τ_μ (сколько существует текущий слот с момента последнего обновления).

Когда набор ключей K_i и значений V_i (индекс i = 1, …, L перечисляет отдельные токены из текущего контекстного окна) передаётся в модуль AM для извлечения соответствующей информации, функция поиска определяет слоты памяти с наиболее сильной ассоциацией между ключом входного токена К и ключом слота памяти AM K_μ^mem:

Ключи и соответствующие им значения L, самые сильно связанные воспоминания (K и V), возвращаются для текущего L, токены передаются обратно в LLM в виде кэша ключ-значение. Подробности см. на рисунке 3.

Операция дополнения

Список извлеченных кэшей «ключ-значение» (K и V) передаются обратно в базовый LLM и используются в качестве префиксного контекста в каждом соответствующем слое с расширенной памятью. Они добавляются к ключам LLM и значениям текущих входных токенов. Список после расширения содержит 2L ключей и значений (длина текущего собственного контекста + длина извлечённых воспоминаний) и L запросов (только текущий контекст). Результатом работы механизма внимания являются расширенные скрытые состояния [h₁, ..., h_L], которые являются входными данными для следующего слоя, как показано в следующих уравнениях и на рисунке выше:

[a₁, …, a_L] = Attn (Q, K’, V’)

Q = [Q₁, Q₂, …, Q_L]

K’ = K¹⊕[K₁, …, K_L]

V’ = V¹⊕[V₁, …, V_L]

Операция записи

Состояние AM обновляется текущим контекстным окном в соответствии с операцией записи, состоящей из двух частей, описанных далее (см. иллюстрацию на рисунке 3).

Консолидация. Если сходство между текущим ключом маркера контекста и наиболее сильно связанным запомненным ключом велико (>R, где R является гиперпараметром), концепция, описанная этим токеном, объявляется знакомой. По этой причине её ключ и значение объединяются с ключом и значением, хранящимися в этом слоте памяти. В частности, слоты памяти обновляются в соответствии со следующими выражениями:

где с_μ отслеживает количество экземпляров, объединенных в слоте μ. Таким образом, консолидированные представления, хранящиеся в каждом слоте μ всегда являются средними арифметическими отдельных экземпляров, которые вошли в этот слот. Вводя частоту обновления ε_μ = 1/(с_μ+1), эти выражения можно переписать как инкрементные модификации существующих представлений, хранящихся в AM.

Новизна и актуальность. Если сходство с ближайшим запомненным ключом слабое (<R), концепция объявляется новой. В этом случае самый старый неиспользуемый слот памяти (с максимальным возрастом τ_μ) заменяется на K_i, V_i, и его возраст устанавливается на 0. После каждого слота μ(i) обновление и его возрастная статистика τ_μ(i) устанавливается на 0, возраст всех слотов, не имеющих соответствующего текущего скрытого состояния контекста, увеличивается на единицу.

Вероятностная интерпретация

В качестве альтернативы ключи и значения в слотах AM можно рассматривать как режимы непараметрической оценки Гауссова распределения смеси (Gaussian mixture distribution), аппроксимирующей многообразие ключ-значение окон прошлого контекста. Эта смесь принимает новые точки ключ-значение из текущего контекста через диагональное ядро ширины R. Средние значения мод смеси обновляются в соответствии с указанными выше правилами средних значений, при этом сохраняется необходимая достаточная статистика для вычисления средних значений в дальнейших обновлениях. Извлечение ближайших мод распределения к текущим скрытым состояниям контекста эффективно приближает полное внимание контекста. Для токенов, ключи и значения которых находятся за пределами радиуса R создаются моды новые, в то время как самые старые удаляются, сохраняя актуальность оценки распределения и её соответствие изменяющемуся контексту.

Сценарии оценки

CAMELoT оценивается по двум сценариям моделирования длинных текстов: задачам моделирования причинно-следственной связи (causal language modeling, CLM) и задачам контекстного обучения с небольшим количеством попыток (in-context learning, ICL).

В качестве базовой модели для CAMELoT использовалась LLaMa2-7b из Huggingface Library с 10k слотами памяти.

Оценка возможности моделирования длинного контекста CAMELoT производилась с использованием трёх ключевых наборов данных:

— Wiki-103, который содержит статьи из Википедии, охватывающие различные темы;

— Arxiv, коллекция академических статей, в первую очередь в области математики, компьютерных наук и физики. Этот набор данных известен своим высококачественным текстом и математическим содержанием, что делает его сложным эталоном для моделирования языка с длинным контекстом;

— PG-19, которая включает в себя полноформатные книги, предлагающие стандартный образец, широко используемый в долгосрочном моделировании естественного языка.

CAMELoT сравнивается с двумя известными трансформерами с расширенной памятью, которые продемонстрировали свою эффективность:

— Transformer-XL, использует подход, основанный на тонкой настройке, сохраняя фиксированную длину предыдущего ввода в кэше для улучшения текущего ввода. Примечательно, что он не использует поиск на основе сходства.

— Memorizing Transformer, сохраняет прошлые кэши циклическим образом. Таким образом, старые кэши заменяются новыми по мере заполнения банка памяти без консолидации.

Используя соответствующую прошлую информацию и новый процесс консолидации памяти, CAMELoT значительно повышает производительность модели (16.6% на PG-19, 29.7% на Arxiv, и 6.36% на wikitext-103, относительно базовой модели в среднем), превосходя другие методы с расширенной памятью.

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

В отличие от моделей без расширения памяти, CAMELoT демонстрирует относительно стабильную производительность при разной длине входных данных. Эта стабильность может быть обусловлена интеграцией дополнительных знаний в AM, сохранённых из предыдущих входных данных. Поскольку CAMELoT накапливает прошлую информацию, её видимый диапазон контекста выходит за рамки текущего ввода, что позволяет эффективно моделировать долгосрочные зависимости независимо от длины текущего ввода. Напротив, модель без расширения памяти полагается исключительно на локальный контекст текущего ввода, что приводит к колебаниям производительности в зависимости от его длины.

Модуль сохраняет свою эффективность даже при крошечных длинах ввода (например, 128), снижая потребность в аппаратных ресурсах, таких как мощные графические процессоры. Это позволяет трансформерам оперировать вниманием с более короткими вводами, но без ущерба для качества моделирования языка. Такое преимущество снижает барьеры для развёртывания больших языковых моделей в средах с ограниченным вычислительным бюджетом.

CAMELoT обладает независимой от модели конструкцией, что позволяет осуществлять бесшовную интеграцию в различные языковые модели. Экспериментальные результаты подтверждают его эффективность, при этом затруднения при моделировании языка с длинным контекстом значительно снижены (до 29,7%), а превосходная производительность последовательно достигается даже при крошечном окне ввода в 128 токенов. Будущие направления исследований, связывающие AM и LLM, включают улучшение конструкции AM (например, автоматическое обучение функции записи) или решение других задач моделирования длинных контекстов (например, сложение длинных документов).

НЛО прилетело и оставило здесь промокод для читателей нашего блога:
-15% на заказ любого VDS (кроме тарифа Прогрев) — HABRFIRSTVDS