Внимание забывает.
Память — не обязана.

MAC (Memory-as-Context, «память как контекст») — это трансформер, у которого окно внимания намеренно маленькое, а рядом работает нейронная память, переносящая информацию, которую окно уже не видит. Стоимость внимания остаётся ограниченной, какой бы длинной ни была последовательность.

100%
точность, с которой модель вспоминает факты, лежащие за пределами того, что она вообще видит (без памяти — почти случайное угадывание, 3%)
O(n·w)
во столько считается внимание в MAC — вместо O(n²), когда модель сравнивает каждое слово со всеми предыдущими
−57%
настолько точнее модель предсказывает текст по сравнению с самой первой версией, за то же время обучения

основа

Как выглядит трансформер

Трансформер — это стопка одинаковых блоков. Текст нарезается на токены, токены превращаются в векторы-эмбеддинги, и дальше каждый блок дважды уточняет их: слой внимания — единственное место, где токены обмениваются информацией друг с другом, — и слой FFN, который обрабатывает каждый токен по отдельности. После N таких блоков модель называет вероятность следующего токена.

токены эмбеддинги блок × N внимание токены обмениваются информацией FFN каждый токен по отдельности следующий токен сюда MAC добавляет второй путь — память

Всё дорогое (и всё интересное) происходит внутри блока — именно его MAC и переустраивает, не трогая остальную схему. Схема выше намеренно упрощена.

резидуальные связи: вход прибавляется к выходу слоя x norm внимание h голов · RoPE · маска norm FFN каузальная маска: токен видит только прошлое · RoPE кодирует позиции внутри внимания после N блоков: финальный norm → проекция в словарь → softmax → вероятности следующего токена

Это уже точная структура (pre-norm вариант, как в mac.py): перед каждым слоем — нормализация, выход слоя прибавляется к входу через резидуальную связь ⊕, внимание состоит из h параллельных голов. Именно к этому блоку MAC добавляет путь памяти — его чтение вливается во вторую сумму ⊕ тем же резидуальным способом.

проблема

Компромисс: стоимость вычислений против дальности контекста

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

Полное внимание

Доступен весь контекст, но стоимость O(n²): при n = 100 000 квадратичный член доминирует в стоимости шага.

каждый токен сравнивается со всеми прошлыми — заполнена вся нижняя половина матрицы
Оконное внимание

Стоимость O(n·w), но зависимости длиннее w структурно невыразимы: между токенами, разнесёнными дальше окна, в графе вычислений нет пути.

только узкая полоса последних w токенов; бледное — забыто навсегда
Окно + память — подход MAC
полоса окна та же, а забытые клетки накрывает память

MAC разделяет функции: локальные зависимости обрабатывает окно внимания (та же дешёвая полоса, что у оконного внимания), а дальние — нейронная память с фиксированной стоимостью на токен. Синие клетки — не точное сравнение токенов, а сжатый конспект: дальняя зависимость снова выразима, хоть и с потерями.

архитектура

Один блок — два пути информации

Каждый слой проводит информацию двумя путями. Путь внимания отвечает за близкое: скользящее окно плюс несколько обучаемых персистентных токенов. Путь памяти отвечает за далёкое: каждый токен записывает связку «ключ→значение» в бегущее состояние быстрых весов, а любой последующий токен может обратиться к нему с запросом — по сути, справочная таблица, едущая вместе с последовательностью.

токены память быстрых весов запись k→v · чтение по запросу оконное внимание локальное окно + перс. токены ⊕ сумма FFN ×N состояние переносится через всю последовательность
путь внимания — локальный, точный, ограниченный путь памяти — глобальный, сжатый, рекуррентный

Что происходит с каждым токеном, по шагам:

  1. Чтение из памяти. Токен формирует запрос и спрашивает память быстрых весов: «что известно о таком ключе?» Ответ — вектор, собранный из всего, что записали все предыдущие токены, даже за тысячи позиций отсюда.
  2. Оконное внимание. Тот же токен внимательно смотрит на последние w соседей — здесь решается синтаксис, согласование, локальный смысл. Персистентные токены (обучаемые, одни и те же для любого текста) дают вниманию стабильную «рабочую доску».
  3. Сумма путей. Ответ памяти прибавляется к выходу внимания резидуально — как поправка, а не как лишние токены в контексте. Поэтому память не удорожает само внимание ни на один токен.
  4. Запись в память. Токен оставляет след для будущих: пара «ключ→значение» вписывается в состояние памяти. Следующие токены прочтут её на шаге 1.
  5. FFN и следующий слой. Обычная полносвязная сеть — и всё повторяется в каждом из N слоёв, у каждого слоя своя память.

Наведите курсор на шаг — соответствующий блок подсветится на схеме.

три резидуальные прибавки: вход каждого шага прибавляется к его выходу через ⊕ x norm память чтение · запись norm внимание окно + перс. токены norm FFN память стоит только на выбранных слоях — у остальных этот шаг просто пропускается окна внимания чередуются по слоям: короткие и длинные, последний слой всегда видит весь контекст во втором стенде (nanochat) порядок другой: сначала внимание, потом память — на выводы это не влияет

Это порядок операций из кода, на котором получены результаты страницы: чтение памяти идёт первым, и её ответ вливается в общий поток ещё до внимания. То есть пути не сливаются в одной сумме, как на упрощённой схеме выше, — это три последовательные прибавки ⊕ в один общий поток, каждая со своей нормализацией. Память сама многоголовая, как и внимание; при обучении она считается сразу кусками по многу токенов, при генерации — по одному, и оба способа дают одинаковый результат.

механизм

Как матрица чисел «запоминает» факты

Память быстрых весов — это одна матрица S на слой. «Быстрые веса» — потому что она меняется на каждом токене (в отличие от обычных, «медленных» весов, которые меняет только обучение). Запись — сложить внешнее произведение ключа и значения; чтение — умножить запрос на матрицу:

запись S S + k vT токен добавляет в матрицу ассоциацию «ключ → значение»
чтение ответ = qTS / qTz запрос вынимает то, что было записано похожими ключами

Аналогия: модель ведёт конспект, а не перечитывает текст. Если запрос q похож на ключ k, записанный когда-то раньше, произведение qTS вернёт примерно то самое значение v — независимо от того, 10 или 10 000 токенов прошло с момента записи. Нормировка на z (сумму ключей) превращает это в средневзвешенный ответ, как в настоящем внимании, только без хранения всех токенов.

Цена — сжатие с потерями: вся последовательность кодируется матрицей фиксированного размера, ёмкость которой конечна. Ограничение воспроизводится в симуляции:

точность извлечения

Симуляция того же вычисления, что в модели: случайные единичные ключи размерности 32, S = Σ kvT, извлечение — ближайшее значение к qTS. Пока пар меньше размерности ключа, они почти ортогональны и не мешают друг другу; дальше следы начинают затирать друг друга.

В репозитории есть и вторая, более простая реализация памяти — слотовая: фиксированный набор обучаемых «ячеек», в которые куски текста записываются усреднением. Она проще устроена, но на запоминании фактов проигрывает быстрым весам впятеро — усреднение размывает факты друг о друга (см. график ниже).

преимущество 1 · извлечение

Достаёт факты, до которых внимание доказуемо не дотягивается

Как проверить, что память выполняет извлечение, а не просто сглаживает статистику? Тест MQAR (multi-query associative recall) устроен как контролируемый эксперимент. В начале последовательности объявляются пары «ключ–значение» — например, A→7, B→3. Дальше идёт случайный наполнитель. В конце модель спрашивают: A→? Симуляция прохода по последовательности:

память быстрых весов: пусто
зелёная рамка — что видит окно внимания (последние 4 токена)

К моменту запроса пары A→7 и B→3 находятся за пределами окна — для оконного внимания ответ структурно недостижим, и его точность не может превышать случайную. Единственный канал — записать пары в состояние памяти при прочтении и извлечь при запросе:

только внимание (окно 32)
3,2% — случайный уровень
+ слотовая память
21,5%
+ память быстрых весов
100%

Условия: 8 пар «ключ–значение» и 4 вопроса по ним внутри текста длиной 128 токенов; окно внимания — 32 токена, то есть пары заведомо вне его досягаемости. Случайно угадать значение — 3%. Полное внимание (без окна) эту задачу решает — значит, дело действительно в архитектуре, а не в том, что задача сама по себе непосильна.

Почему это важно за пределами синтетики: «вспомнить значение по ключу» — это скелет реальных навыков. Имя, названное главой раньше; переменная, объявленная в начале файла; факт из середины документа, нужный для ответа в конце. Модель, проваливающая MQAR, будет проваливать и их — просто на реальном тексте это труднее заметить.

Та же проверка, но пары спрятаны намного дальше — до ~2000 токенов от вопроса вместо ~100. Здесь важно честное условие: полное внимание (без окна) должно уметь решить задачу за отведённое время обучения, иначе сравнивать с ним нечестно — оно просто не успело обучиться, а не «не смогло в принципе».

ДистанцияКонфигурацияRecallШагов
≈488полное внимание (контроль)19,3%6855
окно, без памяти3,7%6674
окно + линейная память99,7%4514
≈1000полное внимание (контроль)11,5%2342
окно, без памяти3,1%2460
окно + линейная память20,7%1777
≈2020полное внимание (контроль)10,8%725
окно, без памяти2,9%703
окно + линейная память10,3%609
сравнение пока нечестное

Случайно угадать — 3%. На дистанции ~1000 и ~2020 полное внимание (контроль) само не решает задачу (11,5% и 10,8% — почти на уровне случайного угадывания). Значит, эти две строки сравнивать пока нельзя: за отведённое время ни у одной архитектуры не получилось толком обучиться на такой дистанции — чем длиннее текст, тем меньше шагов обучения помещается в тот же бюджет времени (6855 → 2342 → 725), и всем банально не хватило времени.

А вот на дистанции ~488 сравнение честное и результат впечатляющий: окно с памятью достигает 99,7%, решительно обходя полное внимание (19,3% — оно решает задачу частично, но явно лучше случайного угадывания, значит сравнение уже имеет смысл). Это интересно само по себе: даже полному вниманию, которое технически видит весь текст, трудно научиться связывать настолько далёкие друг от друга слова за отведённое время. А памяти, которая просто записывает пару «ключ→значение» и потом читает её напрямую, для этого нужно в разы меньше времени на обучение.

Просто дать больше времени не помогло — и это отдельная находка. Увеличение бюджета в 2,5–6,7 раза дало кратно больше шагов обучения, но точность контроля почти не сдвинулась:

ДистанцияБюджетШаговRecall контроля
≈488180с → 450с6855 → 1732919,3% → 21,0%
≈1000180с → 700с2342 → 945111,5% → 21,2%
≈2020180с → 1200с725 → 485810,8% → 16,3%

Результат застревает далеко от решения, а не медленно ползёт к 100% — значит, дело не просто в нехватке времени на обучение. Маленькой модели трудно самостоятельно нащупать, как связывать слова за 500–2000 токенов друг от друга среди кучи ненужного текста между ними — сколько времени ей ни давай. Похоже на настоящий тупик в обучении, а не на бюджетную проблему.

Сравнение памяти на этих дистанциях пока нечестное: нужна либо модель покрупнее, либо другой способ обучения — например, сперва тренировать на коротких расстояниях, потом постепенно увеличивать их.

Решение нашлось в тот же день — учить постепенно. Вместо того чтобы сразу бросать модель на трудную дистанцию, одну и ту же модель обучали по нарастающим этапам: сначала совсем короткие тексты, потом чуть длиннее, и так далее — 128 → 256 → 512 → 1024 → 2048 токенов (то есть дистанции ~119 → ~2039), с тем же суммарным временем обучения, что раньше давало всего 16,3%:

Конфигурация (на ~2039)RecallШагов (всего)
полное внимание (контроль)99,99%32 216
окно(32) + линейная память100,00%20 014

Оба варианта решают задачу почти идеально с тем же бюджетом времени, что раньше давал 16,3%. Это подтверждает догадку: дело было не в возможностях архитектуры, а в том, что модели было слишком трудно нащупать решение, начиная сразу с сложной задачи. Стоит выучить приём на лёгкой дистанции — он почти бесплатно переносится на дистанцию в 17 раз длиннее.

Побочная находка: память решает задачу за 38% меньше шагов обучения, чем полное внимание — даже там, где по чистому времени на GPU память в целом проигрывает (см. следующий раздел), сам приём извлечения у неё получается быстрее.

преимущество 2 · стоимость

Длинный контекст без квадратичного счёта

При окне размера w внимание стоит O(n·w) вместо O(n²). Соотношение объёмов вычислений для конкретных n и w:

полное внимание, пар токенов
окно MAC, пар токенов

внимание MAC считает пар токенов

Шкала логарифмическая — на линейной оконный столбик был бы не виден. Память при любом значении ползунков добавляет одну матрицу фиксированного размера на слой — её стоимость от длины текста не зависит вообще.

Измеренная цена пути памяти: после автоматической компиляции его вычислений — около 6% скорости обучения (до неё было от 12 до 50% в зависимости от настроек; как налог сняли — см. раздел «Честные границы»). Лишних слов во «внимании» модели при этом не прибавляется — память добавляет свой ответ поверх обычного, а не удлиняет тот текст, который модель обрабатывает вниманием.

преимущество 3 · проверяемость

Каждое утверждение выше — это тест или бенчмарк

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

Этап (тот же 5-минутный бюджет)val bpbЧто изменилось
Исходная версия2.168память работала неточно, обновление весов мешало обучению
Устранено подглядывание в будущее2.040починена внутренняя проводка данных
Заменён способ обновления весов1.226теперь модель может научиться запоминать факты
Убрано лишнее кодирование позиции слов0.951меньше параметров, лучше результат
Добавлено оконное внимание0.922ограниченная стоимость, лучший результат

Первый шаг устранил подглядывание в будущее. Второй — заменил способ обновления весов на более простой и надёжный, который не мешал модели формировать связки «записал → потом использовал». Третий — убрал лишний способ кодирования позиции слова (оказалось, что встроенного в само внимание достаточно) — при этом параметров стало меньше, а качество выросло.

Итог: 2.168 → 0.922 бита на байт при неизменной задаче и одном и том же времени на обучение. Все значения — по одному прогону на конфигурацию.

методика

Условия всех измерений

честные границы

Где память окупается — а где пока нет

Память окупается на извлечении. Там, где ответ лежит за пределами окна, она — разница между случайным угадыванием и стопроцентной точностью.

На языковом моделировании она пока не окупается. Причина прозаична: при фиксированном бюджете времени выигрывает тот, кто успевает сделать больше шагов обучения, а путь памяти стоит 15–20% скорости. Если зависимости текста и так помещаются в окно, памяти нечего доносить — и её цена не отбивается. Лучший результат и на коротких рассказах, и на корпусе книг — обычный оконный трансформер без памяти.

Эксперимент на 200 МБ книг (30 минут на конфигурацию): окно 128 токенов превзошло даже полное внимание — модель такого размера не извлекает пользы из контекста дальше окна, и разрыв, который память должна была бы закрывать, попросту отсутствует. Но если считать не по времени, а по числу шагов обучения, память выглядит лучше: слотовый вариант почти сравнялся с окном, потратив на 15% меньше шагов. То есть память проигрывает не потому, что хуже учится, а потому, что на каждый шаг уходит больше времени.

Конфигурация (книги, 30 мин)val bpbшагов
окно 128, без памяти1.46346 332
окно 128 + слотовая память1.47139 198
полное внимание, без памяти1.47747 220
окно 128 + память быстрых весов1.48437 208

Второй, независимый стенд (короткие рассказы, 5 минут на конфигурацию, обучение сразу на двух видеокартах): тот же вывод. Модель собрана из проверенных, открытых блоков проекта nanochat без единой правки в их коде, с той же памятью быстрых весов, добавленной сверху.

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

Конфигурация (TinyStories, 5 мин, DDP ×2)val bpbшагов
полное внимание, без памяти0.93611 645
окно 128, без памяти0.9549 306
полное внимание + память (4/4 слоя)1.0414 738
окно 128 + память (4/4 слоя)1.0434 509
не сравнивать напрямую с таблицей выше

Условия немного другие: здесь окно меньше по отношению к длине текста, и память стоит на всех слоях модели вместо половины из них, отсюда и более заметный налог на скорость (~50% шагов вместо 15–20%). Чем на большем числе слоёв стоит память, тем дороже она обходится — это отдельная настройка, которую стоит учитывать при сравнении. Плюс этот прогон сделан ещё до того, как мы починили медленный путь оконного внимания на этой инфраструктуре — окно здесь тоже, скорее всего, недополучило шагов обучения, просто слабее, чем в следующей таблице.

Масштабный тест сыгран. Гипотеза была: с ростом модели полное внимание должно обогнать окно, и тогда появится разрыв, на котором память сможет доказать свою цену. Модель увеличена в разы (dim 384, depth 8, ~14,5–16 млн параметров против dim 128/depth 4 выше) и обучена через тот же DDP-конвейер на тех же 200 МБ книг, 30 минут на конфигурацию:

Конфигурация (книги, dim 384/depth 8, 30 мин)val bpbшаговVRAM
полное внимание, без памяти1.219943 5271229 МБ
окно, без памяти1.224641 6651229 МБ
окно + память быстрых весов1.255023 8851848 МБ
полное внимание + память1.279420 2401848 МБ

Гипотеза не подтвердилась. Даже такая большая модель показывает почти одинаковый результат, смотрит ли она на весь текст целиком или только на ближайший кусочек (1.2246 против 1.2199 — разница ничтожна). Получается, ей пока и не нужен весь текст: то немногое полезное, что есть в дальнем контексте, она и так не успевает использовать. Поэтому лучший результат по-прежнему у самого простого варианта — окно без всякой памяти.

Но есть любопытная деталь: если память всё же добавить, то в паре с маленьким окном она работает заметно лучше, чем в паре с полным вниманием (1.2550 против 1.2794) — просто этого пока недостаточно, чтобы обогнать вариант без памяти вовсе.

Налог на скорость удалось снять. Всё это время память проигрывала в первую очередь потому, что каждый её шаг обучения стоил заметно дороже — она успевала сделать вдвое меньше шагов за то же время. Выяснилось, что дело не в самой памяти, а в том, как видеокарта исполняет её вычисления: тысячи мелких операций вместо нескольких крупных. Автоматическая компиляция (torch.compile) склеивает их в крупные — и путь памяти почти догоняет модель без памяти, при полном совпадении результатов вычислений. Перепрогон на том же корпусе книг:

Конфигурация (книги, 30 мин, компиляция)val bpbшаговVRAM
окно, без памяти1.228120 0581285 МБ
окно + память быстрых весов1.229418 8021497 МБ
полное внимание + память1.235218 5751503 МБ
не сравнивать напрямую с таблицей выше

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

Впервые за весь проект память идёт вровень с лучшим вариантом. Разница 0.0013 — глубоко ниже порога различимости (0.01), а шагов обучения память теперь делает почти столько же (94% вместо прежних 57%). Выиграть ей по-прежнему не на чем: модель такого размера не извлекает пользы из дальнего контекста, доносить нечего. Но статус памяти изменился: раньше это была дорогая надстройка, теперь — почти бесплатная опция, которая ждёт масштаба, где дальний контекст начнёт приносить пользу.

changelog

Как менялась эта страница