Модуль 5. Современные механизмы внимания

Решение. На каждом слое для одного сохранённого токена нужны и ключ, и значение:

$2H_{kv}d_h\cdot2\ \text{байт}=2\cdot8\cdot64\cdot2=2048\ \text{B/token}.$

Поэтому кэш полных слоёв равен $36\cdot131072\cdot2048=9$ GiB. В гибриде 18 полных слоёв занимают $18\cdot131072\cdot2048=4.5$ GiB, а 18 локальных хранят только окно: $18\cdot128\cdot2048=4.5$ MiB. Сумма: 4.5043945 GiB. Число голов запросов в этой формуле не участвует: K/V хранятся по $H_{kv}$.

Решение. Для первого блока берём максимум $m_1=2$:

$\ell_1=e^{2-2}+e^{1-2}=1+e^{-1}=1.367879,$ $n_1=1e^{0}+2e^{-1}=1.735759.$

Для второго блока $m_2=3$, поэтому $\ell_2=1+e^{-2.5}=1.082085$ и $n_2=3+4e^{-2.5}=3.328340$. Общий максимум $M=3$; статистики первого блока нужно перенести в эту шкалу множителем $e^{m_1-M}=e^{-1}$:

$\ell=e^{-1}\ell_1+\ell_2=1.585300,\qquad n=e^{-1}n_1+n_2=3.966890.$

Итоговый выход внимания $O=n/\ell=\mathbf{2.502297}$. Перенормировка и позволяет объединять блоки без хранения всей матрицы логитов.

Решение. Добавим к обычным логитам $z_i=q^\top k_i/\sqrt{d_h}$ ещё один логит $s$ и соответствующее значение $v_s=0$. Тогда

Этот дополнительный столбец логитов может получить часть вероятностной массы, не добавляя собственного вклада в выход. У фиксированного ключа логит равен $q^\top k_s/\sqrt{d_h}$ и меняется вместе с запросом. Вектор $k_s=0$ даёт только константу 0; произвольную обучаемую константу $s$ для всех $q$ нужно задавать отдельным скаляром или смещением.

Решение. Размер группы равен $g=H_q/H_{kv}=4$, поэтому голова запроса $h$ использует KV-голову $\lfloor h/4\rfloor$: головы 0–3 получают KV0, а 4–7 — KV1. Если обозначить их значениями [10, 20], правильное расширение равно

repeat_interleave(4): [10, 10, 10, 10, 20, 20, 20, 20]

Это и есть групповая формула. tile(4) даёт [10, 20, 10, 20, 10, 20, 10, 20]: размер тензора верен, но головы запросов связываются с другими KV-головами. Следовательно, одной проверки формы массива недостаточно.

Демо к этому проектному упражнению: русскоязычный demo-notebook модуля 5, разделы «Демо 1» и «Демо 2» как отправная точка для проектной работы модуля. Он показывает компактный воспроизводимый пример на малых данных, а не готовое решение полномасштабного проекта.

Решение. Причинное ограничение сравнивает абсолютные позиции: ключ $k$ разрешён при $k\le5$. Поэтому полная строка маски равна [1,1,1,1,1,1]. Окно ширины 3 добавляет условие $5-3+1\le k$, то есть оставляет позиции 3, 4 и 5: [0,0,0,1,1,1].

У tril(ones(1,6)) единственная строка имеет локальный индекс 0, поэтому результат [1,0,0,0,0,0]. Эта операция ничего не знает о пяти уже обработанных позициях. Чтобы получить нужную маску, нужно сравнивать абсолютные индексы или передать соответствующее смещение. В ноутбуке повторите рассуждение для нескольких новых запросов: у каждой строки будет своя абсолютная позиция.

Решение. Соберём запросы строками $Q\in\mathbb R^{3\times6}$, а латенты — строками $C\in\mathbb R^{7\times4}$. В прямом варианте $K=CA^\top$ имеет форму $7\times6$, а $V=CB^\top$ — $7\times5$. Содержательные логиты равны

$QK^\top=QAC^\top.$

Справа сначала вычисляется $QA$ формы $3\times4$, затем матрица логитов формы $3\times7$. В обеих ветвях добавим одну и ту же позиционную часть и применим одинаковые масштабирование, маску и softmax. Получим одинаковые веса $P\in\mathbb R^{3\times7}$. Для выхода:

$PV=P(CB^\top)=(PC)B^\top.$

Прямой путь сначала восстанавливает семь значений ширины 5. Переставленный путь смешивает латенты в $PC$ формы $3\times4$ и лишь затем расширяет результат до $3\times5$. При выходной матрице $O\in\mathbb R^{5\times d_{out}}$ можно заранее объединить $B^\top O$; результат сохраняется.

Если RoPE вращает восстановленный ключ, эффективный запрос становится $A^\top R_{n-m}^\top q$: он зависит от позиции ключа $n$. Это не одно преобразование, которое можно сохранить для всех ключей. Отдельная позиционная ветвь обходит именно эту зависимость. Опыт §8 ноутбука проверяет логиты, выход и выходную проекцию с допуском $10^{-12}$; контрпример с другой позицией ключа проверяет предел перестановки.

Для GPU-проекта сначала сравните точность на малой последовательности с одной головой и эталоном FP64. Затем измеряйте выбранные большие размеры: одинаковые входы, маски и точности, прогрев, no_grad, синхронизация CUDA перед замером и после него, сброс пикового счётчика. Указывайте, входят ли проекции, создание маски и кэш в измерение; фиксируйте выбранное ядро и различайте выделенную и зарезервированную память. Уменьшите размер при нехватке памяти; не стройте большой FP64-эталон лишь ради проверки скорости.