Broadcasting

Правило целиком

Broadcasting — это одно правило, и его стоит выучить точно, потому что все ошибки с ним происходят от приблизительного знания.

$$\htmlData{k=align}{\text{оси сопоставляются с конца}}; \quad \text{совместимы, если равны или одна из них} = 1$$

Пошагово:

1. более короткую форму дополнить единицами слева;
2. пройти по осям и проверить: равны или одна равна $1$;
3. результат по каждой оси — максимум из двух.

Ключевое слово — справа. Именно это ломает интуицию:

форма A	форма B	результат
$(3, 4)$	$(4)$	$(3, 4)$ — работает
$(3, 4)$	$(3)$	ошибка

Второй случай выглядит естественнее первого — «прибавить по одному числу к каждой строке» — и именно он не работает. Вектор $(3)$ выравнивается по последней оси, то есть сопоставляется с $4$, и не подходит. Чтобы получить желаемое, нужно сказать это явно: $(3, 1)$.

Оси-единицы — это разрешение

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

$(3, 1) \ \text{и} \ (1, 4) \ \longrightarrow \ (3, 4)$

Обе оси растягиваются, каждая в своём направлении. Именно так пишут попарные расстояния, матрицы масок и позиционные разности — без единого цикла.

Полезно понимать, что растягивание не копирует данные: реализация ставит страйд $0$ по растянутой оси, и один и тот же элемент читается многократно. Поэтому broadcasting дешёв по памяти на входе — но результат-то материализуется полностью, и вот здесь и подстерегает главная опасность.

Как это молча создаёт гигантские тензоры

Самая дорогая ошибка с broadcasting — не падение, а успех.

$(8, 1, 512) \ \text{и} \ (100, 512) \ \longrightarrow \ (8, 100, 512)$

Это $409\,600$ элементов вместо $4096$ и $51\,200$ на входе. Формы совместимы, ошибки нет, и если вы ожидали поэлементной операции — получили попарную. При батче $256$ и длине последовательности $2048$ такая опечатка мгновенно даёт out-of-memory, и трассировка указывает на строку, где память кончилась, а не на ту, где потерялась ось.

Признак, по которому это узнаётся: тензор вырос, а вы этого не просили. Поэтому в рабочем коде полезно проверять формы утверждениями, а не глазами:

assert scores.shape == (batch, heads, seq, seq), scores.shape

Что делать вместо

Broadcasting неявен, и это его единственный настоящий недостаток. Явные альтернативы:

вместо	лучше
x[:, None] * y[None, :]	einsum('i,j->ij', x, y)
x.unsqueeze(1)	rearrange(x, 'b d -> b 1 d')
надежды на выравнивание	assert на форму

Первая строка — тема следующего урока: einsum делает то же самое, но каждая ось названа, поэтому пропущенную ось видно в спецификации, а не в счётчике памяти.

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

i,j->ijij,ij->ijbd,nd->bnbd,nd->bnd

i, j → ij

индекс	размер	роль
i	3	остаётся в выходе
j	4	остаётся в выходе

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

[3, 4]

элементов

12

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

12

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

1

Сравните две последние спецификации: они различаются одной буквой в выходе, а результат — в $512$ раз по объёму. В einsum это видно; в записи через None-индексы — нет.

Что стоит унести

• выравнивание идёт справа, и $(3,4)$ с $(3)$ не совместимы, хотя выглядят так;
• ось длины $1$ — это разрешение растянуть, и её отсутствие — тоже утверждение;
• растягивание не копирует вход, но результат материализуется целиком;
• самая дорогая ошибка не падает: формы сходятся, а тензор оказывается больше нужного в сотни раз.