2026-08-26·無料·14分で読めます

テンソルと形状操作 — einsumが読めれば論文が読める

論文の Σ_j A_ij B_jk と、コードの x.transpose(1,2) は同じことを言っています。その橋渡しをするのが einsum 記法です。軸・ブロードキャスト・縮約という3つの道具だけで、Attentionの実装が1行で書けるところまで。

#tensor #einsum #broadcasting #shape #numpy #pytorch #math

対象textタスクmath

「形が合わない」で止まる

論文を読み始めた人が最初につまずくのは、たいてい理論ではありません。形（shape）です。

紙の上には $\sum_j A_{ij} B_{jk}$ と書いてある。実装を開くと x.transpose(1, 2).reshape(b, -1, d) が並んでいる。動かすと shapes (32,128,64) and (32,64,128) not aligned と怒られる。式は分かった気がするのに、式とコードの間に橋が架からない——この状態はごく普通です。

その橋になるのが einsum記法 です。torch.einsum('bhqd,bhkd->bhqk', q, k) という一行は、慣れると論文の添字よりも読みやすくなります。逆に言えば、einsumが読めるようになった時点で、論文の $\sum$ とコードの transpose が同じものの言い換えだと分かります。この記事はその読み書きを、前提知識ゼロから身につけるためのものです。

比喩: 軸に名札を付けた棚

テンソルという言葉は物理や微分幾何では厳密な意味を持ちますが、深層学習の実務ではずっと素朴です。多次元に並んだ数値の箱、それだけだと思って構いません。数字が1個ならスカラー、一列ならベクトル、縦横なら行列、その先は個別の名前がないので全部まとめてテンソルと呼んでいる、という程度の話です。

本当に大事なのは箱そのものではなく、それぞれの軸(axis)に意味があることです。たとえば (32, 8, 128, 64) という形を持つテンソルは、ただの数値の塊ではありません。

0番目の軸 32 = バッチ（同時に処理する文の本数）
1番目の軸 8 = ヘッド（注意機構の視点の数）
2番目の軸 128 = 位置（文の中の何番目のトークンか）
3番目の軸 64 = 特徴（1トークンを表す数値の並びの長さ）

棚に喩えるなら、「32棟の倉庫があり、各倉庫に8フロア、各フロアに128個の引き出し、各引き出しに64個の小物入れ」という構造です。倉庫とフロアを取り違えても、数値の総数は同じなのでプログラムは黙って動いてしまいます。形状のバグが静かなのは、これが理由です。

だから最初に身につけるべき習慣は、shapeを数字の羅列として読まないこと。(32, 8, 128, 64) を見たら、頭の中で (batch, head, position, feature) と名札を貼り直す。これだけで、後に出てくる操作のほとんどが「名札の付け替え」として理解できるようになります。

直感: やっていることは3つしかない

深層学習のテンソル計算は、見た目の複雑さのわりに、操作の種類はごく限られています。

1. 軸を潰す（縮約） — ある軸に沿って足し合わせ、その軸を消す。合計・平均・内積・行列積は全部これです。(32, 128, 64) の最後の軸を潰せば (32, 128) になります。

2. 軸を並べる（保持・追加） — 軸を消さずに残す、あるいは新しい軸を作る。外積や、バッチ次元をそのまま持ち回るのがこれです。

3. 軸を揃える（ブロードキャスト） — 形の違う2つのテンソルを、片方の軸を仮想的に引き伸ばして足し算・掛け算できるようにする。

あとは値を動かさず並べ替えるだけの操作（転置・reshape）があるきりです。逆に言えば、論文に出てくるどんな式も、この語彙で読み下せます。einsumは、この3つを1つの文字列で書ける記法です。

揃える: ブロードキャストの3行ルール

先にブロードキャストを片付けておきます。規則は3行です。

2つの形を右端で揃える（左が足りない側には長さ1の軸があると見なす）
各位置で、長さが等しいか、片方が1なら通る
長さ1の軸は、相手の長さぶんコピーされたかのように振る舞う（実際にはコピーせず、同じ値を読み直すだけ）

例を見ます。

x = np.zeros((32, 128, 64))   # (batch, position, feature)
b = np.zeros((64,))           # 特徴ごとのバイアス
x + b                         # OK → (32, 128, 64)

b は右端で揃えると (1, 1, 64) と見なされ、バッチと位置の全マスに同じ64個の値が足されます。「特徴ごとに1つのバイアスを持つ」という意図そのものです。

危ないのは、規則を通ってしまう意図しない組み合わせです。

a = np.zeros((1000,))      # (N,)
c = np.zeros((1000, 1))    # (N, 1)
a - c                       # 通る → (1000, 1000)

引き算をしたつもりが、1000×1000 の距離行列ができています。エラーは出ません。この「長さ1の軸が黙って増殖する」性質が、ブロードキャスト最大の落とし穴です。メモリ使用量が急に跳ねたときは、まずここを疑ってください。

仕組み: 総和記号を1つずつ消していく

行列積の定義から始めます。

C_{ik} = \sum_{j=1}^{n} A_{ij} B_{jk}

(1)

読み下すと「出力の $i$ 行 $k$ 列の値は、 $A$ の $i$ 行目と $B$ の $k$ 列目を、共通の添字 $j$ に沿って掛けて足し合わせたもの」です。 $i$ と $k$ は出力に残る添字、 $j$ は足し合わせて消える添字です。

ここでアインシュタインが気づいたのは、足し合わせる添字は、見ればわかるということでした。 $j$ は右辺に2回出てくるのに左辺には出てこない。ならば $\sum_j$ を書く必要はない。こうして総和記号を省いたのが、式(1)と同じ意味を持つ次の書き方です。

C_{ik} = A_{ij} B_{jk}

(2)

これがアインシュタインの総和規約です。einsum（Einstein summation）という関数名はここから来ています。そしてこの式をそのまま文字列にしたものが、

C = np.einsum('ij,jk->ik', A, B)

です。カンマが入力の区切り、-> の右が出力。添字の文字は自分で選べます。

読み書きのルールは3つだけです。

-> の右に出てくる文字は、出力に残る軸（自由添字）
入力にあって出力にない文字は、その軸に沿って合計されて消える（縮約）
複数の入力に同じ文字が出てきたら、その軸同士を対応付けて掛け合わせる

つまり「どの軸を残すか」を書けば、残りは自動的に潰れる。何をするかではなく、何が残るかを書く記法だと捉えると腑に落ちます。

読み書きの練習

同じルールで、よく使う操作がすべて書けます。声に出して読み下してみてください。

einsum	意味	読み方
`'i,i->'`	内積	同じ軸を掛けて、その軸を潰す。何も残らないのでスカラー
`'i,j->ij'`	外積	共通の文字がないので潰れない。2軸が並ぶ
`'ij->ji'`	転置	出力で文字の順を入れ替えただけ
`'ij->i'`	行ごとの合計	`j` が消えるので列方向に潰れる
`'ii->'`	トレース	同じ文字を1つの入力内で2回使うと対角線を取る
`'ii->i'`	対角成分	対角線を取り、潰さずに残す
`'ij,jk->ik'`	行列積	式(2)そのもの
`'bij,bjk->bik'`	バッチ行列積	`b` は両方にあり出力にも残るので、潰れずに並走する
`'...ij,...jk->...ik'`	先頭の軸は何本でもよい	`...` は「残りの軸を右端で揃える」＝ブロードキャスト

一番覚えておく価値があるのは b の振る舞いです。両方の入力にあり、かつ出力にも残る文字は、潰れずに「その軸ごとに独立して同じ計算をする」という意味になります。バッチ・ヘッド・時刻など、並列に持ち回りたい軸はすべてこの形で書けます。

そしてAttentionのスコア計算は、この語彙でこう書けます。

# q, k: (batch, head, position, dim)
scores = torch.einsum('bhqd,bhkd->bhqk', q, k)   # (batch, head, q位置, k位置)

読み下すと「バッチとヘッドはそのまま並走させ、特徴軸 d を潰し、クエリ位置 q とキー位置 k を残す」。これは $QK^\top$ と全く同じもので、しかも「どの軸で内積を取るのか」が文字列に書いてあるぶん、transpose(-2, -1) を挟んだ書き方より読み間違えにくい。Attention自体の意味はAttention機構を1から理解するで扱っています。

最小単位である内積の感覚は、実際に動かすのが早いです。

FIG 1einsumの `'i,i->'` が測っているもの。2本のベクトルを回すと内積が変わる。向きが揃うほど大きく、直交すると0、逆向きで負。「軸を潰す」とはこの1つの数に畳むこと

コードで書く

Attentionの本体を、einsumだけで組み立ててみます。

import torch

def attention(q, k, v, mask=None):
    # q, k, v: (b, h, n, d)
    d = q.shape[-1]
    scores = torch.einsum('bhqd,bhkd->bhqk', q, k) / d**0.5
    if mask is not None:
        scores = scores.masked_fill(mask, float('-inf'))  # mask: (b, 1, q, k) 等
    w = scores.softmax(dim=-1)
    return torch.einsum('bhqk,bhkd->bhqd', w, v)          # 重み付き平均

2つのeinsumが対になっているのが見えます。1つ目は特徴軸 d を潰してキー位置 k を作り、2つ目はキー位置 k を潰して特徴軸 d を戻す。行って帰ってくる構造がそのまま文字列に現れています。

軸そのものを組み替えたいときは、einsumより einops の rearrange が読みやすくなります。

from einops import rearrange
# (b, n, h*d) を (b, h, n, d) に分解する
q = rearrange(q, 'b n (h d) -> b h n d', h=8)

(h d) という括弧が「この軸は h と d が結合したもの」を表します。reshape(b, n, 8, -1).transpose(1, 2) と同じ処理ですが、何を何に分けたのかが文字列に書いてあるぶん、レビューで間違いに気づけます。

縮約の順序で、計算量が桁で変わる

einsumは3つ以上の入力も取れます。

np.einsum('ij,jk,kl->il', A, B, C)

数学的には $(AB)C = A(BC)$ で結果は同じです。ところが計算量は同じではありません。 $A$ が $(100, 1000)$ 、 $B$ が $(1000, 1000)$ 、 $C$ が $(1000, 5)$ のとき、

$(AB)C$ : $100{\times}1000{\times}1000 + 100{\times}1000{\times}5 \approx 1.005$ 億回
$A(BC)$ : $1000{\times}1000{\times}5 + 100{\times}1000{\times}5 \approx 550$ 万回

18倍以上の差です。掛け算の順序を変えただけで、です。

FIG 2縮約の順序を間違えると、同じ答えを出すのに1段上の曲線を登ることになる。nが小さいうちは誤差でも、nが伸びると差が桁になる

NumPyの np.einsum は既定では順序を最適化しません。3つ以上を一度に縮約するなら optimize=True を付けるか、np.einsum_path で選ばれた順序を確認してください。

print(np.einsum_path('ij,jk,kl->il', A, B, C, optimize='optimal')[1])

PyTorchの torch.einsum は opt_einsum が入っていれば順序探索を使い、torch.backends.opt_einsum で挙動を切り替えられます。とはいえ最適化が効くのは3項以上のときで、2項なら順序の余地はありません。行列積そのもののコストがどこから来るかは行列積のコストで扱っています。

現場ではこう使う

誰がいつ使うか。 論文を読んで実装に落とすリサーチエンジニア、他人のモデルコードを読んで推論を最適化する担当、そして学習が回らない原因を追うMLOps。共通しているのは「他人が書いたテンソル操作を、短時間で正しく読む」場面です。einsumとeinopsは、そのための共通語になっています。

触るもの。 torch.einsum / np.einsum / np.einsum_path / optimize=True、einops の rearrange reduce repeat、そしてデバッグ時の tensor.shape tensor.stride() tensor.is_contiguous()。

事故になる落とし穴。

1. reshapeで軸を並べ替えようとする。 (b, n, h, d) を (b, h, n, d) にしたいときに reshape を使うと、エラーは出ないのに値が別の場所へ混ざります。並べ替えは transpose / permute、結合と分解だけが reshape の仕事です。学習が「なんとなく収束しない」形でしか現れないので、発見が非常に遅れます。

2. view が落ちる。 transpose の後のテンソルはメモリ上で連続していないため view は例外を投げます。reshape は通りますが、裏で黙ってコピーが走ります。

3. マスクの形。 Attentionのマスクは (b, 1, 1, k)（パディング）と (b, 1, q, k)（因果）で意味が違います。ブロードキャストが通るので間違った形でも動く——動くけれど未来のトークンが見えている、という最悪の壊れ方をします。

4. 中間テンソルの実体化。 'bhqd,bhkd->bhqk' は $b{\times}h{\times}q{\times}k$ の配列をメモリ上に作ります。バッチ8、ヘッド32、文長4096、fp16なら $8{\times}32{\times}4096{\times}4096{\times}2$ バイト ≈ 8GiB。出力より中間結果のほうが大きいという逆転が起き、これがOOMの典型的な原因です。FlashAttentionのような手法は、この表を作らずに済ませる工夫だと理解できます。

5. 添字の文字を使い回す。 'bhqd,bhkd->bhqk' の q と k を両方 n と書くと、クエリとキーが同じ軸だと宣言したことになります。文字は必ず意味ごとに変える。

面接で問われる形。 「torch.einsum('bhqd,bhkd->bhqk', q, k) は何をしていますか」は定番です。①潰れる文字（d）を指し、②残る文字（b, h, q, k）を指し、③「バッチとヘッドは並走、特徴軸で内積、位置×位置の表ができる」と言い切る。ここまで言えれば、続く「その表は何バイトですか」にも上の掛け算で答えられます。

まとめ

テンソルは「軸に意味がある箱」。shapeは数字ではなく名札として読む
操作は3種類しかない: 潰す（縮約）・並べる（保持）・揃える（ブロードキャスト）
einsumは「何をするか」ではなく「何が残るか」を書く記法。出力に無い文字は消える
3項以上では縮約の順序で計算量が桁で変わる。optimize=True を忘れない
ブロードキャストと reshape は、エラーを出さずに間違える。だから危ない

ベクトルと行列そのものが何をしているのかに戻りたくなったら、AIのための線形代数から読み直すと、einsumの文字列が急に自然な言葉に見えてきます。