Broadcasting и векторизация

Зачем это нужно

Посмотрите на любую строку кода трансформера, и вы увидите операции над массивами разных форм. К активациям формы (batch, seq, d) прибавляется сдвиг формы (d,). Из каждого вектора-токена вычитается его среднее (нормализация, урок 3.4.1). К скорам внимания (batch, heads, seq, seq) применяется маска (seq, seq). Softmax превращает каждую строку чисел в вероятности (урок 1.3.2).

Ни одна из этих операций не пишется циклом. Их делают два механизма: агрегации по осям и broadcastingbroadcasting (растягивание форм)EN broadcastingПравило, по которому NumPy и PyTorch выполняют операции над массивами разных форм без копирования данных. — правило, по которому NumPy и PyTorch «растягивают» массив меньшей формы до большей, не копируя данные. Добавьте к ним матричное умножение @ и einsum, и у вас будет весь словарь, на котором написан attention.

Агрегации по осям

sum, mean, max, argmax, std можно считать по всему массиву или вдоль одной осиось (измерение массива)EN axisАгрегация по оси axis=k «съедает» эту ось, если не указать keepdims=True.:

import numpy as np

X = np.array([[1., 2., 3.],
              [4., 5., 6.]])        # shape (2, 3)

X.sum()             # 21.0 — всё
X.sum(axis=0)       # [5. 7. 9.] — вдоль строк, по каждому столбцу → shape (3,)
X.sum(axis=1)       # [ 6. 15.] — вдоль столбцов, по каждой строке → shape (2,)
X.max(axis=1)       # [3. 6.]
X.argmax(axis=1)    # [2 2] — индекс максимума в каждой строке (так выбирают самый вероятный токен)

Правило, которое снимает все вопросы: axis=k — это ось, которая исчезает. У X форма (2, 3). Сумма по axis=0 убирает ось длины 2 и оставляет (3,), сумма по axis=1 убирает ось длины 3 и оставляет (2,). Для тензора (batch, seq, d) запись x.mean(axis=-1) значит «среднее по каждому вектору-токену», результат — (batch, seq).

Среднее и разброс: mean и std

mean — обычное среднее арифметическое: сложить и поделить на количество. std — стандартное отклонениестандартное отклонениеEN standard deviationКорень из дисперсии; измеряется в тех же единицах, что и сама величина.: насколько числа в среднем разбросаны вокруг своего среднего. Если все числа одинаковые, std равно 0; чем сильнее они разбегаются, тем std больше. Посчитаем его вручную для чисел [2, 4, 4, 4, 5, 5, 7, 9]:

1. Среднее: $(2 + 4 + 4 + 4 + 5 + 5 + 7 + 9) / 8 = 40 / 8 = 5$.
2. Отклонения от среднего: $-3, -1, -1, -1, 0, 0, 2, 4$.
3. Квадраты отклонений: $9, 1, 1, 1, 0, 0, 4, 16$. Их среднее $32 / 8 = 4$ называют дисперсией.
4. Корень из дисперсии: $\sqrt{4} = 2$. Это и есть std.

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

v = np.array([2, 4, 4, 4, 5, 5, 7, 9])
print(v.mean(), v.std())              # 5.0 2.0
print(np.array([5, 5, 5, 5]).std())   # 0.0 — разброса нет

std, как и sum, считается по осям: X.std(axis=0) — разброс каждого столбца. Он понадобится в практике этого урока и в следующем уроке, где мы будем измерять шум.

keepdims: оставить ось длины 1

С keepdims=True ось не исчезает, а остаётся длины 1:

X.mean(axis=1)                       # [2. 5.]         shape (2,)
X.mean(axis=1, keepdims=True)        # [[2.] [5.]]     shape (2, 1)
X - X.mean(axis=1, keepdims=True)    # [[-1. 0. 1.] [-1. 0. 1.]] — каждая строка центрирована

«Центрировать» значит вычесть из чисел их среднее: после этого среднее каждой строки равно 0.

Почему последняя строка работает, а X - X.mean(axis=1) падает с ошибкой operands could not be broadcast together with shapes (2,3) (2,)? Отвечают правила broadcasting.

Правила broadcasting

Когда NumPy видит операцию над массивами разных форм, он:

1. Выравнивает формы по правому краю. Если осей меньше, слева дописываются единицы.
2. Сравнивает размеры по каждой оси. Они совместимы, если равны или один из них равен 1.
3. Ось длины 1 растягивается до размера другого массива. Данные при этом не копируются: NumPy просто читает одну и ту же строку много раз.

Если хотя бы по одной оси размеры разные и ни один не равен 1, будет ошибка.

a        (3, 1)          a        (2, 3)          a       (2, 3)
b           (4,)   →     b           (3,)    →    b          (2,)
          (1, 4)                   (1, 3)                  (1, 2)
-------------------      -------------------      -------------------
a + b    (3, 4)  ✓       a + b    (2, 3)  ✓       a + b   3 ≠ 2  ✗ ошибка

Вот как растягивается первый пример, столбец (3, 1) плюс строка (4,):

a = np.array([[0], [10], [20]])   # (3, 1)
b = np.array([0, 1, 2, 3])        # (4,)
print(a + b)

[[ 0  1  2  3]
 [10 11 12 13]
 [20 21 22 23]]

	b: 0	b: 1	b: 2	b: 3
a: 0	0	1	2	3
a: 10	10	11	12	13
a: 20	20	21	22	23

Столбец a скопирован вправо 4 раза, строка b — вниз 3 раза. Таблица сложения без единого цикла. Так же устроено внешнее произведениевнешнее произведениеEN outer productСтолбец на строку — u vᵀ, матрица ранга 1. В китайских текстах «внешним произведением» иногда называют векторное произведение (cross product), смотрите по контексту. x[:, None] * y[None, :] — таблица умножения «каждый на каждого»: для x = [1, 2, 3] и y = [10, 20] получится [[10, 20], [20, 40], [30, 60]].

Теперь понятно, зачем keepdims. Среднее по строкам формы (2,) выравнивается по правому краю, то есть пытается сопоставиться со столбцами (их 3). Форма (2, 1) выравнивается правильно: 2 строки против 2 строк, а единица растягивается на 3 столбца.

flowchart LR
A["X (2, 3)"] --> R{"Выровнять по правому краю"}
M["mean (2, 1)"] --> R
R --> C1["ось 0: 2 и 2 — равны ✓"]
R --> C2["ось 1: 3 и 1 — единица растягивается ✓"]
C1 --> OUT["X - mean: (2, 3)"]
C2 --> OUT

Как NumPy проверяет совместимость форм

Broadcasting в трансформере

B, T, d, H = 2, 4, 8, 2          # батч, число токенов, длина вектора токена, число «голов» внимания

h = np.zeros((B, T, d))
bias = np.arange(d)              # (d,) → (1, 1, d)
(h + bias).shape                 # (2, 4, 8) — сдвиг прибавлен к каждому токену каждого примера

scores = np.zeros((B, H, T, T))
causal = np.tril(np.ones((T, T), dtype=bool))      # (T, T): tril оставляет диагональ и всё под ней, выше — False
masked = np.where(causal, scores, -np.inf)         # (T, T) → (1, 1, T, T) → (B, H, T, T)
print(causal.astype(int))

[[1 0 0 0]
 [1 1 0 0]
 [1 1 1 0]
 [1 1 1 1]]

Одна маска (T, T) применяется ко всем примерам батча и всем головам сразу. np.where(условие, a, b) из прошлого урока берёт число из a там, где условие верно, и из b — где нет. Токен $i$ видит только токены $j \le i$, а на месте будущих стоит $-\infty$ (-np.inf — особое значение float «минус бесконечность», меньше любого числа), который softmax превратит в ноль. Это причинная маска GPT (модуль 4.4).

Тихая ловушка: broadcasting без ошибки

Broadcasting опасен тем, что ошибка формы часто не вызывает исключения. Классика — подсчёт MSE (mean squared error, средний квадрат разности между предсказанием и правильным ответом, урок 2.2.1). Предсказания модели имеют форму (n, 1), а правильные ответы — (n,).

y = np.array([1., 2., 3., 4.])                     # (4,)
y_pred = np.array([[1.1], [1.9], [3.2], [3.9]])    # (4, 1) — так часто выдаёт модель
print((y_pred - y).shape)                          # (4, 4) — не то, что вы хотели!
print(((y_pred - y) ** 2).mean())                  # 2.4425 — «ошибка» модели
print(((y_pred[:, 0] - y) ** 2).mean())            # ≈ 0.0175 — настоящий MSE (последние знаки — погрешность float)

Вместо 4 разностей получилась таблица 4×4 «каждый с каждым», и loss завышен в 140 раз. Модель при этом будет как-то учиться, и баг можно искать неделю. Вторая версия ловушки — квадратная матрица: S - S.mean(axis=1) для S формы (3, 3) не упадёт, а молча вычтет среднее не из той оси.

Защита

Печатайте формы и ставьте проверки: assert y_pred.shape == y.shape, (y_pred.shape, y.shape). Для средних по строкам всегда используйте keepdims=True. В тестах берите неквадратные размеры (например, 3×5), тогда перепутанные оси сразу дадут ошибку.

Матричное умножение: * против @

A * B — поэлементное произведение: перемножаются числа на одинаковых местах (с broadcasting). A @ B — матричное умножениематричное умножениеEN matrix multiplication (matmul)(m×n)·(n×k) → (m×k), элемент cᵢⱼ = Σₜ aᵢₜbₜⱼ. Соответствует композиции преобразований. Стоит ≈ 2·m·n·k операций., и устроено оно иначе. Вспомните скалярное произведение из практики урока 0.2.3 (v @ w у класса Vector) и x @ x из прошлого урока: перемножить числа на одинаковых местах и сложить. Матричное умножение — это много скалярных произведений сразу: число в строке $i$ и столбце $k$ результата — скалярное произведение строки $i$ матрицы A на столбец $k$ матрицы B.

A = np.array([[1, 2, 3],
              [4, 5, 6]])     # (2, 3)
B = np.array([[1, 2],
              [0, 1],
              [2, 0]])        # (3, 2)
print(A @ B)

[[ 7  4]
 [16 13]]

Число результата	Строка A	Столбец B	Скалярное произведение
[0, 0]	[1, 2, 3]	[1, 0, 2]	1·1 + 2·0 + 3·2 = 7
[0, 1]	[1, 2, 3]	[2, 1, 0]	1·2 + 2·1 + 3·0 = 4
[1, 0]	[4, 5, 6]	[1, 0, 2]	4·1 + 5·0 + 6·2 = 16
[1, 1]	[4, 5, 6]	[2, 1, 0]	4·2 + 5·1 + 6·0 = 13

Отсюда правило форм: $(m, n) @ (n, k) \to (m, k)$. Строка A и столбец B должны быть одной длины $n$ — внутренние размеры совпадают, — а в результате столько строк, сколько у A, и столько столбцов, сколько у B. Умножить $(2, 3)$ на $(2, 3)$ нельзя: строка длины 3 и столбец длины 2 не перемножить, будет ошибка. Что матричное умножение значит геометрически, разберём в уроке 1.1.2, а здесь нам важны формы. Кстати, A @ B — это вызов A.__matmul__(B): тот же механизм магических методов, что у вашего класса Vector.

Линейный слой нейросети — это X @ W + b:

X = np.ones((32, 768))     # батч из 32 векторов размерности 768
W = np.ones((768, 3072))   # веса слоя: 768 → 3072
b = np.ones(3072)          # сдвиг, broadcasting по батчу
(X @ W + b).shape          # (32, 3072)

Если у массивов больше двух осей, @ умножает матрицы по двум последним осям, а все ведущие оси считает «батчем» и применяет к ним broadcasting. Пример — attention-скоры $QK^\top$ ($K^\top$ — транспонированная матрица K, урок 0.3.1): для каждой пары токенов число «насколько запрос токена $q$ подходит к ключу токена $k$» (урок 4.4.2). Для всех примеров и всех голов они считаются одной строкой:

Q = np.ones((2, 8, 10, 64))           # (B, H, T, d_head)
K = np.ones((2, 8, 10, 64))
# swapaxes(-1, -2) меняет местами две последние оси — транспонирует каждую матрицу (T, d) внутри батча
scores = Q @ K.swapaxes(-1, -2)       # (B, H, T, d) @ (B, H, d, T) → (B, H, T, T)
scores.shape                          # (2, 8, 10, 10)

einsum: формы, записанные словами

Когда осей много, легко запутаться, что с чем умножается. einsumeinsum (нотация Эйнштейна)EN Einstein summation (einsum)Запись тензорных операций строкой индексов, например 'bhqd,bhkd->bhqk'. позволяет записать операцию строкой индексов:

$$C_{ik} = \sum_j A_{ij} B_{jk}$$

A = np.arange(6).reshape(2, 3)
B = np.arange(12).reshape(3, 4)
C = np.einsum('ij,jk->ik', A, B)   # [[20 23 26 29] [56 68 80 92]] — то же, что A @ B

Формула над кодом читается так: число $C$ в строке $i$ и столбце $k$ равно сумме произведений $A_{ij} B_{jk}$ по всем $j$. Маленькие буквы внизу — номера строки и столбца: $A_{ij}$ — это A[i, j]. Знак $\sum_j$ («сумма по $j$») — это цикл, который перебирает $j$ и складывает (подробно о нём — урок 1.0.3). Получается ровно правило «строка на столбец» из раздела выше. Проверим первое число ответа: $C_{00} = A_{00}B_{00} + A_{01}B_{10} + A_{02}B_{20} = 0 \cdot 0 + 1 \cdot 4 + 2 \cdot 8 = 20$.

Та же формула тремя циклами — медленно, зато видно, что делает einsum:

C = np.zeros((2, 4), dtype=int)
for i in range(2):                   # i — строка ответа
    for k in range(4):               # k — столбец ответа
        for j in range(3):           # j есть во входах, но нет в выходе — по нему складываем
            C[i, k] += A[i, j] * B[j, k]
print(C)                             # [[20 23 26 29] [56 68 80 92]]

Правила:

1. Каждая буква — ось. Входы перечисляются через запятую, выход — после ->.
2. Буква, которая есть во входах, но нет в выходе, суммируется.
3. Одинаковые буквы в разных входах означают «перемножить соответствующие элементы».

Операция	einsum	Эквивалент
сумма всех элементов	'ij->'	A.sum()
суммы по столбцам	'ij->j'	A.sum(axis=0)
транспонирование	'ij->ji'	A.T
скалярное произведение	'i,i->'	v @ w
внешнее произведение	'i,j->ij'	v[:, None] * w[None, :]
матричное умножение	'ij,jk->ik'	A @ B
батчевое умножение	'bij,bjk->bik'	A @ B для 3D
attention-скоры	'bhqd,bhkd->bhqk'	Q @ K.swapaxes(-1, -2)

Q = np.random.default_rng(0).standard_normal((2, 8, 10, 64))
K = np.random.default_rng(1).standard_normal((2, 8, 10, 64))
att = np.einsum('bhqd,bhkd->bhqk', Q, K)
print(att.shape, np.allclose(att, Q @ K.swapaxes(-1, -2)))   # (2, 8, 10, 10) True

Строку 'bhqd,bhkd->bhqk' можно прочитать вслух: «для каждого примера $b$ и головы $h$ сложи по $d$ произведения запроса $q$ на ключ $k$». Транспонировать ничего не нужно: einsum сам понимает, какие оси сопоставить.

Когда einsum, а когда @

@ обычно чуть быстрее и привычнее. einsum хорош, когда осей больше трёх или их нужно переставить по ходу дела. Для цепочек из трёх и более тензоров в np.einsum передавайте optimize=True: NumPy подберёт порядок умножений.

Векторизация на практике: попарные расстояния

Сначала — что такое расстояние между точками. На плоскости у точки две координаты, например $a = (1, 2)$ и $b = (4, 6)$. По теореме Пифагора расстояние между ними — корень из суммы квадратов разностей координат: $\sqrt{(4 - 1)^2 + (6 - 2)^2} = \sqrt{9 + 16} = \sqrt{25} = 5$. В $d$ измерениях всё так же, только слагаемых $d$. Это евклидово расстояниеевклидово расстояниеEN Euclidean distanceL2-норма разности двух векторов ‖a − b‖₂ — обычное расстояние «по прямой».. Корень часто не извлекают и работают с квадратом расстояния (здесь 25): чем ближе точки, тем меньше и расстояние, и его квадрат, а считать проще.

Задача: даны точки $A$ формы $(n, d)$ и $B$ формы $(m, d)$ (каждая строка — точка из $d$ координат), нужна матрица квадратов расстояний $D$ формы $(n, m)$: $D[i, j]$ — квадрат расстояния от точки A[i] до точки B[j]. Это ядро поиска похожих документов для LLM (RAG, ступень 8) и метода ближайших соседей kNN (ступень 2). Три решения:

def loops(A, B):                      # двойной цикл Python
    D = np.empty((len(A), len(B)))
    for i in range(len(A)):
        for j in range(len(B)):
            D[i, j] = ((A[i] - B[j]) ** 2).sum()
    return D

def broadcast(A, B):                  # (n, 1, d) - (1, m, d) → (n, m, d) → сумма по d
    return ((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=-1)

def via_matmul(A, B):                 # ||a - b||² = ||a||² + ||b||² − 2·a·b
    return (A ** 2).sum(1)[:, None] + (B ** 2).sum(1)[None, :] - 2 * A @ B.T

Для $n = 500$, $m = 400$, $d = 64$ на машине автора (один поток процессора):

Способ	Время	Во сколько раз быстрее цикла
двойной цикл	≈ 500 мс	1×
broadcasting	≈ 52 мс	≈ 10×
через @	≈ 2.4 мс	≈ 200×

Третий вариант опирается на школьную формулу $(a - b)^2 = a^2 + b^2 - 2ab$, применённую к каждой координате и сложенную по всем координатам: квадрат расстояния = (сумма квадратов координат a) + (сумма квадратов координат b) − 2 × (скалярное произведение a и b). Запись $\lVert a \rVert^2$ в комментарии и есть «сумма квадратов координат» — квадрат длины вектора, norm() из урока 0.2.3. Проверим на тех же точках: $\lVert a \rVert^2 = 1 + 4 = 5$, $\lVert b \rVert^2 = 16 + 36 = 52$, $a \cdot b = 1 \cdot 4 + 2 \cdot 6 = 16$, итого $5 + 52 - 2 \cdot 16 = 25$ — тот же ответ. А скалярные произведения сразу всех пар даёт одно матричное умножение A @ B.T.

Broadcasting убрал цикл, но создал промежуточный тензор $(500, 400, 64)$ — 12.8 млн чисел, 102 МБ во float64. При $n = m = 10^4$ этот тензор занял бы уже 51 ГБ: обычный компьютер столько памяти не имеет. Вариант через @ раскрывает квадрат разности и сводит задачу к матричному умножению, для которого есть оптимизированная библиотека BLAS. Векторизация — это не только «убрать цикл», но и сформулировать задачу через матричные операции. Ровно этим и занимаются авторы эффективных ядер внимания вроде FlashAttention и FlashMLA.

Практика

Четыре функции, ни одного цикла. Каждая — комбинация приёмов этого урока:

• normalize_rows — поделить каждую строку на её длину (норму, как norm() у класса Vector из урока 0.2.3), чтобы длина стала 1: строка [3, 4] длины 5 превращается в [0.6, 0.8]. Длины всех строк сразу — это корень np.sqrt из суммы квадратов по axis=1 с keepdims=True.
• standardize_columns — из каждого столбца вычесть его среднее и поделить на его std. После этого у каждого столбца среднее 0 и std 1, и столбцы в разных единицах (метры, рубли) становятся сравнимы. Столбец [2, 4, 4, 4, 5, 5, 7, 9] (среднее 5, std 2) превратится в [-1.5, -0.5, -0.5, -0.5, 0, 0, 1, 2]. Понадобится на ступени 2.
• pairwise_sq_dists — квадраты расстояний между всеми парами точек из раздела выше. Нужны в kNN и RAG.
• cosine_similarity — косинусное сходствокосинусное сходствоEN cosine similarityКосинус угла между векторами — скалярное произведение, делённое на произведение длин. Для нормированных векторов равно скалярному произведению.: насколько два вектора смотрят в одну сторону, независимо от их длины. Это скалярное произведение векторов, заранее приведённых к длине 1. Ответ лежит от −1 до 1: [1, 0] и [3, 0] дают 1 (одно направление), [1, 0] и [0, 2] — 0 (перпендикулярны), [1, 0] и [-1, 0] — −1 (противоположны). Так ищут похожие по смыслу эмбеддинги (урок 1.1.1).

Одна ловушка: у нулевой строки длина 0, а деление на ноль даст nan («не число»). Поэтому длину сначала «подпирают» снизу крошечным числом: np.maximum(norms, 1e-12) поэлементно берёт большее из двух.

Нормализация строк и попарные расстояния без цикловnumpy · нужен интернет7 тестов

import numpy as np


def normalize_rows(X):
    """Делит каждую строку X формы (n, d) на её длину (L2-норму), чтобы длина стала 1.

    Длина строки — корень из суммы квадратов её чисел: у [3, 4] это sqrt(9 + 16) = 5,
    поэтому [3, 4] -> [0.6, 0.8]. Нулевые строки должны остаться нулевыми (без деления на 0 и без nan).
    Без циклов: длины строк — сумма по оси 1 с keepdims=True (форма (n, 1)), дальше broadcasting.
    """
    # TODO: 1) norms = корень (np.sqrt) из суммы квадратов X по axis=1, keepdims=True;
    #       2) верните X, делённый на np.maximum(norms, 1e-12) — так нулевая строка не даст деления на 0.
    ...


def standardize_columns(X, eps=1e-8):
    """Каждый столбец X (n, d): вычесть среднее столбца и поделить на (std столбца + eps).

    После этого у каждого столбца среднее 0 и std 1. Пример для одного столбца:
    [2, 4, 4, 4, 5, 5, 7, 9] (среднее 5, std 2) -> примерно [-1.5, -0.5, -0.5, -0.5, 0, 0, 1, 2].
    eps — крошечное число, чтобы не делить на 0, если все числа в столбце одинаковые.
    Входной массив X менять нельзя — верните новый.
    """
    # TODO: mean и std по оси 0 (по столбцам), оба с keepdims=True — форма (1, d);
    #       затем (X - среднее) / (std + eps).
    ...


def pairwise_sq_dists(A, B):
    """Квадраты расстояний между всеми парами строк: A (n, d), B (m, d) -> D (n, m).

    Каждая строка — точка с d координатами. D[i, j] — сумма квадратов разностей координат
    точек A[i] и B[j]: для (0, 0) и (3, 4) это 3*3 + 4*4 = 25. Без циклов Python.
    """
    # TODO: 1) diff = A[:, None, :] - B[None, :, :] — форма (n, m, d): разности «каждая с каждой»;
    #       2) возведите в квадрат и сложите по последней оси (axis=-1) — останется (n, m).
    ...


def cosine_similarity(A, B):
    """Косинусное сходство всех пар строк: A (n, d), B (m, d) -> S (n, m).

    S[i, j] — скалярное произведение строк A[i] и B[j], заранее приведённых к длине 1.
    Ответ от -1 до 1: [1, 0] и [3, 0] -> 1 (одно направление), [1, 0] и [0, 2] -> 0,
    [1, 0] и [-1, 0] -> -1. Формула: cos(a, b) = a·b / (длина a · длина b).
    """
    # TODO: нормализуйте строки обеих матриц функцией normalize_rows и перемножьте их через @;
    #       у второй матрицы возьмите .T, чтобы её строки стали столбцами.
    ...


# Проверьте себя кнопкой «Запустить»:
X = np.array([[3.0, 4.0], [0.0, 0.0], [1.0, 0.0]])
print(normalize_rows(X))
print(pairwise_sq_dists(X, X))

ЗапуститьПроверитьПодсказкаСбросить

Китайский угол: einsum в коде DeepSeek-V3

Этот блок забегает далеко вперёд: если половина слов непонятна — это нормально, вернитесь к нему после модуля 5.2. Сейчас достаточно увидеть, что строки einsum из этого урока встречаются в коде настоящей большой модели.

В репозитории deepseek-ai/DeepSeek-V3 лежит эталонный код инференса — файл inference/model.py. В классе внимания MLA (Multi-head Latent Attention, модуль 5.2) есть режим, где attention-скоры считаются прямо через einsum. Упрощённо (без части, отвечающей за RoPE — способ сообщить модели позицию токена, урок 4.5.2) строки выглядят так:

q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
scores = torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos])

Теперь вы можете их прочитать. b — батч, s — позиции запросов, t — позиции ключей, h — головы, d — размерность головы, c — размерность сжатого латентного вектора (в V3 он равен 512). Первая строка умножает каждый запрос на матрицу восстановления ключей. Её «поглощают» в запрос, чтобы не разворачивать кэш обратно в полные ключи. Вторая строка считает скоры сразу по сжатому KV-кэшу (это ключи и значения уже прочитанных токенов, которые модель хранит при генерации, урок 4.6.4): индекс c повторяется и суммируется.

Память MLA экономит тем, что хранит в KV-кэше не полные ключи и значения всех голов, а один сжатый вектор длины 512 на токен (и ещё 64 числа для RoPE). А чтобы не распаковывать этот кэш на каждом шаге генерации, произведение матриц вычисляют в другом порядке — «переставляют скобки». einsum делает такую перестановку читаемой. Подробно разберём в уроке 5.2.2.

В китайском учебнике d2l-zh broadcasting называется «механизмом вещания» (как радиовещание: одно значение рассылается на все позиции), а суммы по осям — «суммой с понижением размерности» и «суммой без понижения» (наш keepdims=True). Пригодится, когда будете читать переводы китайских разборов кода моделей.

Проверь себя

Проверь себя6 вопросов·

1 Какой формы будет a + b, если a.shape == (8, 1, 6, 1) и b.shape == (7, 1, 5)?

(8, 7, 6, 5)(8, 1, 6, 5)Ошибка: разное число осей(8, 7, 1, 5)

2 X имеет форму $(3, 4)$, v — $(3,)$. Что даст X + v?

Форму $(3, 4)$: v прибавится к каждому столбцуФорму $(3, 4)$: v прибавится к каждой строкеОшибку несовместимых формФорму $(3, 3, 4)$

3 Какие выражения правильно центрируют каждую строку X формы $(n, d)$ (вычитают среднее строки)?Можно выбрать несколько ответов

X - X.mean(axis=1, keepdims=True)X - X.mean(axis=1)[:, None]X - X.mean(axis=1)(X.T - X.mean(axis=1)).T

4 A — $(1000, 128)$, B — $(500, 128)$. Сколько миллионов элементов в промежуточном тензоре A[:, None, :] - B[None, :, :]?

5 y_pred.shape == (100, 1), y.shape == (100,). Какой формы будет y_pred - y?

(100,)(100, 1)(100, 100)Ошибка

6 Запишите строку einsum для матричного умножения A (оси i, j) на B (оси j, k). Без кавычек.

Проверить ответы

Итоги

• axis=k — ось, которая исчезает при агрегации. keepdims=True оставляет её длины 1 для дальнейшего broadcasting.
• Broadcasting: выровнять формы справа, единицы растягиваются, иначе ошибка. Данные при этом не копируются.
• Главная опасность — тихий broadcasting: (n, 1) - (n,) даёт (n, n). Проверяйте формы через assert.
• * — поэлементно, @ — матрично по двум последним осям, ведущие оси работают как батч.
• einsum записывает операцию строкой индексов. 'bhqd,bhkd->bhqk' — это attention-скоры. Им же написан MLA в коде DeepSeek-V3.

Дальше: случайные числа, воспроизводимость и графики — как смотреть на данные глазами.