MAC (Memory-as-Context, «память как контекст») — это трансформер, у которого окно внимания намеренно маленькое, а рядом работает нейронная память, переносящая информацию, которую окно уже не видит. Стоимость внимания остаётся ограниченной, какой бы длинной ни была последовательность.
основа
Трансформер — это стопка одинаковых блоков. Текст нарезается на токены, токены превращаются в векторы-эмбеддинги, и дальше каждый блок дважды уточняет их: слой внимания — единственное место, где токены обмениваются информацией друг с другом, — и слой FFN, который обрабатывает каждый токен по отдельности. После N таких блоков модель называет вероятность следующего токена.
Всё дорогое (и всё интересное) происходит внутри блока — именно его MAC и переустраивает, не трогая остальную схему. Схема выше намеренно упрощена.
Это уже точная структура (pre-norm вариант, как в mac.py): перед каждым слоем —
нормализация, выход слоя прибавляется к входу через резидуальную связь ⊕,
внимание состоит из h параллельных голов. Именно к этому блоку MAC добавляет путь памяти —
его чтение вливается во вторую сумму ⊕ тем же резидуальным способом.
проблема
Внимание — механизм, которым трансформер обращается к предыдущим токенам, — сравнивает
каждый новый токен со всеми прошлыми, поэтому число сравнений растёт квадратично
с длиной последовательности. Стандартная мера — скользящее окно: внимание ограничивается
последними w токенами. Стоимость становится линейной, но токены дальше
w выпадают из графа вычислений — зависимость от них невыразима в принципе.
Доступен весь контекст, но стоимость O(n²): при n = 100 000 квадратичный член доминирует в стоимости шага.
Стоимость O(n·w), но зависимости длиннее w структурно невыразимы: между токенами, разнесёнными дальше окна, в графе вычислений нет пути.
MAC разделяет функции: локальные зависимости обрабатывает окно внимания (та же дешёвая полоса, что у оконного внимания), а дальние — нейронная память с фиксированной стоимостью на токен. Синие клетки — не точное сравнение токенов, а сжатый конспект: дальняя зависимость снова выразима, хоть и с потерями.
архитектура
Каждый слой проводит информацию двумя путями. Путь внимания отвечает за близкое: скользящее окно плюс несколько обучаемых персистентных токенов. Путь памяти отвечает за далёкое: каждый токен записывает связку «ключ→значение» в бегущее состояние быстрых весов, а любой последующий токен может обратиться к нему с запросом — по сути, справочная таблица, едущая вместе с последовательностью.
Что происходит с каждым токеном, по шагам:
w соседей — здесь решается синтаксис, согласование, локальный смысл. Персистентные токены (обучаемые, одни и те же для любого текста) дают вниманию стабильную «рабочую доску».Наведите курсор на шаг — соответствующий блок подсветится на схеме.
Это порядок операций из кода, на котором получены результаты страницы: чтение памяти идёт первым, и её ответ вливается в общий поток ещё до внимания. То есть пути не сливаются в одной сумме, как на упрощённой схеме выше, — это три последовательные прибавки ⊕ в один общий поток, каждая со своей нормализацией. Память сама многоголовая, как и внимание; при обучении она считается сразу кусками по многу токенов, при генерации — по одному, и оба способа дают одинаковый результат.
механизм
Память быстрых весов — это одна матрица S на слой. «Быстрые веса» — потому что
она меняется на каждом токене (в отличие от обычных, «медленных» весов, которые меняет
только обучение). Запись — сложить внешнее произведение ключа и значения; чтение — умножить
запрос на матрицу:
Аналогия: модель ведёт конспект, а не перечитывает текст. Если запрос q похож
на ключ k, записанный когда-то раньше, произведение qTS
вернёт примерно то самое значение v — независимо от того, 10 или 10 000
токенов прошло с момента записи. Нормировка на z (сумму ключей) превращает
это в средневзвешенный ответ, как в настоящем внимании, только без хранения всех токенов.
Цена — сжатие с потерями: вся последовательность кодируется матрицей фиксированного размера, ёмкость которой конечна. Ограничение воспроизводится в симуляции:
Симуляция того же вычисления, что в модели: случайные единичные ключи размерности 32, S = Σ kvT, извлечение — ближайшее значение к qTS. Пока пар меньше размерности ключа, они почти ортогональны и не мешают друг другу; дальше следы начинают затирать друг друга.
В репозитории есть и вторая, более простая реализация памяти — слотовая: фиксированный набор обучаемых «ячеек», в которые куски текста записываются усреднением. Она проще устроена, но на запоминании фактов проигрывает быстрым весам впятеро — усреднение размывает факты друг о друга (см. график ниже).
преимущество 1 · извлечение
Как проверить, что память выполняет извлечение, а не просто сглаживает статистику?
Тест MQAR (multi-query associative recall) устроен как контролируемый эксперимент. В начале
последовательности объявляются пары «ключ–значение» — например,
A→7, B→3. Дальше идёт случайный наполнитель.
В конце модель спрашивают: A→? Симуляция прохода по последовательности:
К моменту запроса пары A→7 и B→3 находятся за пределами
окна — для оконного внимания ответ структурно недостижим, и его точность не может превышать
случайную. Единственный канал — записать пары в состояние памяти при прочтении и извлечь
при запросе:
Условия: 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 контроля |
|---|---|---|---|
| ≈488 | 180с → 450с | 6855 → 17329 | 19,3% → 21,0% |
| ≈1000 | 180с → 700с | 2342 → 9451 | 11,5% → 21,2% |
| ≈2020 | 180с → 1200с | 725 → 4858 | 10,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 считает пар токенов
Шкала логарифмическая — на линейной оконный столбик был бы не виден. Память при любом значении ползунков добавляет одну матрицу фиксированного размера на слой — её стоимость от длины текста не зависит вообще.
Измеренная цена пути памяти: после автоматической компиляции его вычислений — около 6% скорости обучения (до неё было от 12 до 50% в зависимости от настроек; как налог сняли — см. раздел «Честные границы»). Лишних слов во «внимании» модели при этом не прибавляется — память добавляет свой ответ поверх обычного, а не удлиняет тот текст, который модель обрабатывает вниманием.
Окно делает стоимость внимания на каждое слово постоянной — она больше не растёт квадратично с длиной текста.
Одна матрица на каждый слой, не растущая с длиной текста — модель в принципе может читать бесконечно длинный текст, не требуя всё больше памяти видеокарты.
При обучении память считается сразу большими блоками (быстро на видеокарте), а при генерации текста — по одному слову за раз (дёшево). Тесты подтверждают: оба способа дают один и тот же результат с точностью до пятого знака после запятой.
преимущество 3 · проверяемость
Архитектура поставляется вместе с инструментами, которые её проверяли, — и эти инструменты поймали настоящие ошибки, которые обычная метрика качества никогда бы не заметила: модель незаметно подглядывала в будущее вместо того, чтобы честно его предсказывать; способ обновления весов при обучении мешал ей учиться запоминать факты; а способ указывать позицию слова портил различимость слов на длинных текстах.
Включая тесты, которые буквально подменяют слова в конце текста и проверяют, что предсказания для более ранних слов не изменились ни на йоту — то есть модель никогда не подглядывает вперёд, в любой настройке.
Диагностика архитектуры: отличает настоящее запоминание фактов от простого угадывания по общей статистике текста — и с его помощью нашли две ошибки из трёх.
Бит на байт (bpb) — сколько бит модель тратит на предсказание одного байта текста: меньше — лучше, и метрику нельзя «накрутить» сменой токенизатора. Считается на одном и том же отложенном куске текста: то, на чём модель обучалась, может меняться, а линейка для оценки — нет.
| Этап (тот же 5-минутный бюджет) | val bpb | Что изменилось |
|---|---|---|
| Исходная версия | 2.168 | память работала неточно, обновление весов мешало обучению |
| Устранено подглядывание в будущее | 2.040 | починена внутренняя проводка данных |
| Заменён способ обновления весов | 1.226 | теперь модель может научиться запоминать факты |
| Убрано лишнее кодирование позиции слов | 0.951 | меньше параметров, лучше результат |
| Добавлено оконное внимание | 0.922 | ограниченная стоимость, лучший результат |
Первый шаг устранил подглядывание в будущее. Второй — заменил способ обновления весов на более простой и надёжный, который не мешал модели формировать связки «записал → потом использовал». Третий — убрал лишний способ кодирования позиции слова (оказалось, что встроенного в само внимание достаточно) — при этом параметров стало меньше, а качество выросло.
Итог: 2.168 → 0.922 бита на байт при неизменной задаче и одном и том же времени на обучение. Все значения — по одному прогону на конфигурацию.
методика
MAC-трансформер: dim 128, 4 слоя, 4 головы внимания (dim_head 64), ~1,6 млн параметров. Конфигурации в сравнениях различаются только размером окна и наличием/типом памяти; всё остальное фиксировано.
Токенизатор SentencePiece (словарь 2000), обучается на корпусе стенда. Метрика — бит на байт (bpb) на замороженных валидационных отрезках; она не зависит от выбора токенизатора. Для корпуса книг — чередующийся train/val-сплит, окна оценки по 1024 токена.
AdamW (β₁ 0,9, β₂ 0,99), lr 6·10−4 с расписанием, weight decay отделён и применяется только к весам. Бюджет фиксируется по астрономическому времени на GPU, поэтому сравнение смешивает качество шага и скорость шага — рядом с bpb всегда приводится число шагов.
Все числа — одиночные прогоны без усреднения по сидам; различия менее ~0,01 bpb не следует интерпретировать. Аппаратура: две GPU по 24 ГБ, по одной задаче на карту.
честные границы
Память окупается на извлечении. Там, где ответ лежит за пределами окна, она — разница между случайным угадыванием и стопроцентной точностью.
На языковом моделировании она пока не окупается. Причина прозаична: при фиксированном бюджете времени выигрывает тот, кто успевает сделать больше шагов обучения, а путь памяти стоит 15–20% скорости. Если зависимости текста и так помещаются в окно, памяти нечего доносить — и её цена не отбивается. Лучший результат и на коротких рассказах, и на корпусе книг — обычный оконный трансформер без памяти.
Эксперимент на 200 МБ книг (30 минут на конфигурацию): окно 128 токенов превзошло даже полное внимание — модель такого размера не извлекает пользы из контекста дальше окна, и разрыв, который память должна была бы закрывать, попросту отсутствует. Но если считать не по времени, а по числу шагов обучения, память выглядит лучше: слотовый вариант почти сравнялся с окном, потратив на 15% меньше шагов. То есть память проигрывает не потому, что хуже учится, а потому, что на каждый шаг уходит больше времени.
| Конфигурация (книги, 30 мин) | val bpb | шагов |
|---|---|---|
| окно 128, без памяти | 1.463 | 46 332 |
| окно 128 + слотовая память | 1.471 | 39 198 |
| полное внимание, без памяти | 1.477 | 47 220 |
| окно 128 + память быстрых весов | 1.484 | 37 208 |
Второй, независимый стенд (короткие рассказы, 5 минут на конфигурацию, обучение сразу на двух видеокартах): тот же вывод. Модель собрана из проверенных, открытых блоков проекта nanochat без единой правки в их коде, с той же памятью быстрых весов, добавленной сверху.
Что обучение на двух картах действительно работает правильно, проверили так: намеренно скормили каждой карте разные данные и убедились, что после каждого шага веса модели на обеих картах совпадают до бита — значит, карты честно обмениваются информацией, а не тренируются порознь.
| Конфигурация (TinyStories, 5 мин, DDP ×2) | val bpb | шагов |
|---|---|---|
| полное внимание, без памяти | 0.936 | 11 645 |
| окно 128, без памяти | 0.954 | 9 306 |
| полное внимание + память (4/4 слоя) | 1.041 | 4 738 |
| окно 128 + память (4/4 слоя) | 1.043 | 4 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.2199 | 43 527 | 1229 МБ |
| окно, без памяти | 1.2246 | 41 665 | 1229 МБ |
| окно + память быстрых весов | 1.2550 | 23 885 | 1848 МБ |
| полное внимание + память | 1.2794 | 20 240 | 1848 МБ |
Гипотеза не подтвердилась. Даже такая большая модель показывает почти одинаковый результат, смотрит ли она на весь текст целиком или только на ближайший кусочек (1.2246 против 1.2199 — разница ничтожна). Получается, ей пока и не нужен весь текст: то немногое полезное, что есть в дальнем контексте, она и так не успевает использовать. Поэтому лучший результат по-прежнему у самого простого варианта — окно без всякой памяти.
Но есть любопытная деталь: если память всё же добавить, то в паре с маленьким окном она работает заметно лучше, чем в паре с полным вниманием (1.2550 против 1.2794) — просто этого пока недостаточно, чтобы обогнать вариант без памяти вовсе.
Налог на скорость удалось снять. Всё это время память проигрывала в первую очередь потому, что каждый её шаг обучения стоил заметно дороже — она успевала сделать вдвое меньше шагов за то же время. Выяснилось, что дело не в самой памяти, а в том, как видеокарта исполняет её вычисления: тысячи мелких операций вместо нескольких крупных. Автоматическая компиляция (torch.compile) склеивает их в крупные — и путь памяти почти догоняет модель без памяти, при полном совпадении результатов вычислений. Перепрогон на том же корпусе книг:
| Конфигурация (книги, 30 мин, компиляция) | val bpb | шагов | VRAM |
|---|---|---|---|
| окно, без памяти | 1.2281 | 20 058 | 1285 МБ |
| окно + память быстрых весов | 1.2294 | 18 802 | 1497 МБ |
| полное внимание + память | 1.2352 | 18 575 | 1503 МБ |
Этот прогон шёл в другом режиме — на двух видеокартах одновременно (и с вдвое большей порцией текста на шаг), поэтому вариант без памяти прогнан заново в тех же условиях как точка отсчёта. По качеству он воспроизвёл предыдущий результат, так что выводы двух таблиц согласуются, а вот число шагов между таблицами сравнивать нельзя.
Впервые за весь проект память идёт вровень с лучшим вариантом. Разница 0.0013 — глубоко ниже порога различимости (0.01), а шагов обучения память теперь делает почти столько же (94% вместо прежних 57%). Выиграть ей по-прежнему не на чем: модель такого размера не извлекает пользы из дальнего контекста, доносить нечего. Но статус памяти изменился: раньше это была дорогая надстройка, теперь — почти бесплатная опция, которая ждёт масштаба, где дальний контекст начнёт приносить пользу.
changelog
Снят «налог на скорость» пути памяти: автоматическая компиляция склеивает его мелкие вычисления в крупные, и память впервые сравнялась по качеству с вариантом без памяти при том же времени обучения. Обновлены разделы «Стоимость» и «Честные границы».
Починен реальный баг с быстрым внутренним ядром для внимания (Flash Attention 3): оно не подключалось, и оконное внимание из-за этого считалось медленнее, чем должно было — теперь оба типа внимания сравниваются честно.
Масштабный тест: модель увеличена в разы, но окно и полное внимание всё ещё на равных — разрыв, на котором память могла бы себя проявить, пока не появился.
Проверка на дистанциях до ~2000 токенов: обнаружено, что модель застревает при обучении на трудных дистанциях, сколько бы времени ей ни давать. Решение нашлось в постепенном обучении — сначала на коротких дистанциях, потом на всё более длинных.
Модель впервые обучена сразу на двух видеокартах одновременно; независимый стенд на основе проекта nanochat подтвердил прежние выводы.
Масштабный эксперимент на 200 МБ книжного текста.
Найдены и исправлены три настоящих дефекта в модели (подглядывание в будущее, неудачный способ обновления весов, лишнее кодирование позиции слов) — качество предсказания текста выросло более чем вдвое при том же времени на обучение.