Как нейросеть помнит факты: линейный слой, MLP и самовнимание как память

Разбор в блоге alexwang.ai отвечает на вопрос, как большие языковые модели помнят столько информации и сколько могут хранить. Автор предлагает смотреть на линейные отображения, MLP и самовнимание как на разные виды памяти «ключ-значение» и из этой картины выводит эвристические оценки ёмкости каждого слоя. Это не точные результаты, а прикидки при упрощающих допущениях.

Базовые понятия. Память хранит M воспоминаний; каждое воспоминание, пара из ключа (по нему ищут, сравнивая по сходству) и значения (его хотят получить). Запрос x возвращает выход f(x). Хорошая память выдаёт значение v_i и на ключ k_i, и на похожий вектор. Память, которая вспоминает по сходству, в литературе называют ассоциативной. Главная величина всего текста, отношение помех к сигналу: насколько остальные воспоминания мешают вспомнить нужное. Ёмкость, число воспоминаний, при котором помехи вырастают до размера сигнала. По словам автора, при 64-мерных ключах переход от линейного слоя к MLP и затем к самовниманию, это скачок примерно от 64 воспоминаний к примерно 2700 и затем к более чем миллиону.

Простейшая линейная память копит воспоминания через внешние произведения: W = сумма v_i·k_iᵀ. Чтобы записать новое воспоминание, к весам просто прибавляют vkᵀ. В отступлении автор показывает, что шаг градиентного спуска на линейном слое, это и есть запись воспоминания с ключом x и значением −η·ȳ, где ȳ, градиент потерь по выходу слоя. Читать такую память можно умножением входа на матрицу: это сумма значений, взвешенных сходством запроса с каждым ключом.

При чтении возникают помехи. Если ключи нормированы до длины 1, то f(k_j) равно значению целевого воспоминания (сигнал) плюс сумма вкладов всех остальных воспоминаний с похожими ключами (помеха). Если бы ключи были ортонормированными, помех не было бы, но в пространстве размерности d_k лишь d_k ортогональных направлений, поэтому при M > d_k идеальное воспоминание невозможно. Для оценки ёмкости автор берёт случайные единичные ключи и случайные значения схожей длины: сходство ключей порядка ±1/√d_k, M−1 слагаемых помехи складываются как случайное блуждание, и отношение норм помехи и сигнала порядка √(M/d_k). Приравняв его к 1, получаем ёмкость M ≈ d_k: когда воспоминаний столько же, сколько измерений ключа, помеха заглушает сигнал. Даже ниже этого порога помех у такой памяти много.

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

Оптимальная линейная память. Веса подбирают так, чтобы сразу подогнать все воспоминания, то есть решают задачу линейной регрессии (метод наименьших квадратов). При M > d_k решение: W = VᵀK(KᵀK)⁻¹. Это та же внешняя память VᵀK, но запрос предварительно умножается на (KᵀK)⁻¹, что «отбеливает» ключи, и память перестаёт переоценивать направления, общие для многих ключей. Если M ≤ d_k и ключи линейно независимы, регрессия подгоняет все воспоминания точно, то есть помех нет, даже если ключи не ортогональны. При M > d_k для случайных ключей и значений отношение помех к сигналу равно примерно √(M/d_k − 1) и достигает 1 при M = 2d_k: ёмкость вдвое больше, чем у внешнего произведения. В отступлении автор выводит это через матрицу проекции P = K(KᵀK)⁻¹Kᵀ: след матрицы равен d_k, поэтому диагональные элементы P_jj ≈ d_k/M, а сумма квадратов недиагональных элементов равна P_jj − P_jj².

Численная проверка: M воспоминаний со случайными единичными ключами размерности d_k = 64. Память на внешних произведениях достигает ёмкости при M ≈ d_k = 64. Оптимальная память вспоминает точно вплоть до M = d_k, затем помехи растут от нуля как √(M/d_k − 1) и достигают ёмкости около M = 2d_k = 128, как и предсказано. Поскольку оптимальная память снимает фиксированное число «воспоминаний» помех, разрыв между двумя видами памяти с ростом M сокращается. Но такая память всё ещё линейна по сходствам: она масштабирует малые и большие сходства одинаково, чтобы улучшить результат, нужно менять сам способ чтения.

MLP как нелинейная память. Если пропускать каждое сходство k_iᵀx через нелинейную функцию σ перед взвешенной суммой значений, получится f(x) = Vᵀσ(Kx): вклад слабо похожих ключей подавляется и помех становится меньше. Автор отмечает, что так «заново открыт» двухслойный MLP: строки первой весовой матрицы, это ключи, столбцы второй, значения. Сходство между слоями прямой передачи и памятью, по словам автора, отмечали Sukhbaatar и соавторы (а также Geva и соавторы; в имеющемся тексте фраза обрывается). Дальше по плану идёт самовнимание, которое, как заявляет автор, хорошо вспоминает и способно хранить потенциально неограниченное число воспоминаний.

Ключевые факты

• Автор рассматривает линейные слои, MLP и самовнимание как ассоциативную память «ключ-значение» и оценивает ёмкость каждого через отношение помех к сигналу; ёмкость, число воспоминаний, при котором помехи сравниваются с сигналом.
• Простейшая линейная память на внешних произведениях вмещает порядка M ≈ d_k воспоминаний (для случайных единичных ключей и случайных значений); отношение помех к сигналу примерно √(M/d_k).
• Оптимальная линейная память (метод наименьших квадратов) вспоминает без помех до M = d_k и достигает ёмкости при M = 2d_k; в численной проверке с d_k = 64 это около 64 и около 128 воспоминаний.
• Если пропустить сходства ключей через нелинейность, получается двухслойный MLP: строки первой матрицы, ключи, столбцы второй, значения. По утверждению автора, при 64-мерных ключах ёмкость растёт от примерно 64 до примерно 2700 (MLP) и более миллиона (самовнимание).
• Шаг градиентного спуска на линейном слое автор трактует как запись воспоминания с ключом x и значением −η·ȳ.

Почему это важно

Вопрос, как модели хранят знания, лежит в основе многих споров о работе нейросетей. Разбор даёт простую единую рамку: линейный слой, MLP и самовнимание, разные способы записи и чтения пар «ключ-значение», а различие между ними сводится к тому, как они борются с помехами между воспоминаниями. Из неё получаются наглядные оценки ёмкости, которые автор подаёт как эвристики, а не как точные результаты.

Кому это важно

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

Как это применить

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

Можно ли доверять

Это независимый авторский разбор в блоге, не рецензируемая статья. Автор сам называет результаты эвристиками и опирается на идеализацию: случайные единичные ключи и случайные значения схожей длины. Численная проверка выполнена на случайных ключах с d_k = 64, экспериментов на настоящих больших языковых моделях в доступном тексте нет. Вывод цифр около 2700 (MLP) и более миллиона (самовнимание) в доступном фрагменте отсутствует: он обрывается в начале раздела про MLP, так что эти оценки приходится принимать на слово автора. Имя автора в видимом тексте не указано.

Риски и подводные камни

Не стоит читать оценки ёмкости как точные свойства реальных моделей: они получены для случайных ключей и значений, тогда как у настоящих моделей ключи и значения структурированы и обучаются. Цифры 64, около 2700 и более миллиона относятся к конкретному случаю 64-мерных ключей. Нужно помнить и о границах самого рассуждения: линейная память вмещает порядка d_k воспоминаний, а нелинейные варианты, по заявлению автора, существенно больше, но подробное обоснование этого в доступной части не приведено.

«Мы только что заново открыли двухслойный MLP из первых принципов!»

— автор поста на alexwang.ai