Батчевые операции

Батчевая ось — это ось, по которой ничего не сворачивается

Из урока про einsum: индекс, который есть во входах и в выходе, не суммируется. Он просто повторяет операцию.

$$\texttt{'}\htmlData{k=batch}{bh}\htmlData{k=contract}{q d}\texttt{,}\htmlData{k=batch}{bh}\htmlData{k=contract}{k d}\texttt{->}\htmlData{k=batch}{bh}qk\texttt{'}$$

Отсюда полное определение: батчевая ось — это та, что присутствует во всех операндах и в результате. Контракция — исчезающий индекс, по которому берётся сумма.

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

Правило batched matmul

torch.matmul для тензоров порядка выше двух работает так:

1. последние две оси трактуются как матрица;
2. все остальные — как батчевые и подчиняются broadcasting (урок 020);
3. свёртка идёт по последней оси первого и предпоследней второго.

A	B	результат
$(b, n, k)$	$(b, k, m)$	$(b, n, m)$
$(b, h, n, k)$	$(b, h, k, m)$	$(b, h, n, m)$
$(b, n, k)$	$(k, m)$	$(b, n, m)$ — B растянута
$(b, 1, n, k)$	$(h, k, m)$	$(b, h, n, m)$ — обе растянуты

Третья строка — самый частый случай на практике: линейный слой применяется к батчу, и матрица весов $(k, m)$ растягивается по всем батчевым осям без копирования.

Четвёртая строка — там, где стоит быть внимательным: broadcasting по батчевым осям означает, что пропущенная ось не вызовет ошибки, если формы окажутся совместимыми. То же предупреждение, что в уроке 020, но теперь внутри matmul.

Что видно в спецификации

линейный слойбатчевое умножениедве батчевые осиусреднение по позициям

bnk, km → bnm

индекс	размер	роль
b	8	батчевая или свободная
n	64	батчевая или свободная
k	512	контракция
m	512	батчевая или свободная

форма результата

[8, 64, 512]

элементов

262,144

умножений-сложений

134,217,728

операций на элемент

512

Сравните первые две спецификации: 'bnk,km->bnm' и 'bnk,bkm->bnm'. Работа одинаковая, а вес второй операции в $b$ раз больше по памяти — потому что во втором случае у каждого элемента батча своя матрица. Это буквально разница между nn.Linear и групповой свёрткой, и в нотации она видна как наличие или отсутствие одной буквы.

Свёртка как батчевая операция

Полезное упражнение для интуиции: свёртка $1\times1$ — это в точности линейный слой, применённый к каждой позиции.

$\texttt{'bchw,dc->bdhw'}$

Каналы сворачиваются ($c$), пространственные оси ведут себя как батчевые ($h$, $w$). Отсюда и практическое наблюдение: свёртки $1\times1$ в архитектурах вроде ResNet — это не «упрощённые свёртки», а способ смешать каналы, ничего не делая с пространством.

Настоящая свёртка $3\times3$ в эту нотацию не укладывается, потому что требует перекрывающихся окон, а einsum умеет только независимые индексы. Реализуется она через unfold (материализация окон) или специализированные ядра — и это единственное место в блоке, где нотация индексов не покрывает операцию.

Форма как документация

Практический вывод из всего урока. Формы — самая дешёвая доступная документация, и её стоит писать явно:

# q, k, v: (batch, heads, seq, head_dim)
scores = einsum('bhqd,bhkd->bhqk', q, k) / math.sqrt(head_dim)
assert scores.shape == (batch, heads, seq, seq), scores.shape

Три строки, из которых:

• комментарий фиксирует смысл осей, чего форма не сообщает;
• спецификация фиксирует, что именно сворачивается;
• assert ловит несовпадение там, где оно произошло.

Из уроков 010 и 020 известно, что ошибки в осях чаще всего не падают. Эти три строки — то, что превращает молчаливую ошибку в громкую.