テンソル縮約と縮約順序最適化

テンソル縮約とは何か、なぜ順序が問題になるか

テンソル縮約（tensor contraction）は、複数の多次元配列を共有する添字について和を取り、1つの配列にまとめる演算です。行列積 C[i,k] = Σ_j A[i,j] B[j,k] はその最も単純な例で、共有添字 j について和を取り、j を消して i,k を残しています。ベクトル内積、行列とベクトルの積、トレース（Σ_i A[i,i]）もすべて、この「共有添字について和を取る」という一つの枠組みで表せます。

実務で扱うのは2テンソルの積だけではありません。量子多体系、機械学習、量子化学では、数十個のテンソルを一括で縮約する式が現れます。ここで決定的に効いてくるのが縮約する順番です。3つ以上のテンソルを縮約するとき、どのペアから縮約するかで生成される中間テンソルの大きさが変わり、総演算量（FLOP数）が桁違いに変化します。順序は結果の値には影響しませんが、計算コストには決定的に影響します。

einsum記法 ── 縮約を添字の並びだけで書く

einsum（Einstein summation、アインシュタインの縮約記法）は、各テンソルの添字を文字列で並べ、繰り返し現れる添字を自動的に和の対象とみなす記法です。数式を書かずに縮約の構造だけを宣言できます。

einsum の読み方（規則は2つだけ）

  入力側にだけ現れ、出力側に無い添字  → 縮約する（その添字について和を取る）
  入力側にも出力側にも現れる添字      → 残す（和を取らない）

  例:
  "ij,jk->ik"    行列積      j を縮約、i,k を残す
  "ij,ij->"      全要素の内積  i,j を縮約、出力はスカラー
  "ii->"         トレース      対角の和
  "ij->ji"       転置         縮約なし、添字の並べ替えのみ
  "ijk,jkl->il"  2添字を一括縮約  j,k を縮約

einsum の強みは、和・転置・対角抽出・ブロードキャストといった別々に見える操作を、添字文字列という単一の宣言的インターフェースに統合できる点です。ただし naive な einsum 実装は、全添字にわたる多重ループを素朴に回すため計算量が爆発します。実用的な実装は、この宣言から後述の「良い縮約順序」を導き、各ステップを行列積へ帰着させて高速化します。

縮約順序で計算量が桁違いに変わる

3テンソルの縮約 "ij,jk,kl->il" を例に、順序による差を見ます。各添字の次元をNとします。

式: A[i,j] B[j,k] C[k,l] を縮約して結果[i,l] を得る（各次元 = N）

順序1: (A・B) を先に縮約
  中間 T1[i,k] = Σ_j A[i,j] B[j,k]   コスト ≈ N³（i,k,j の三重ループ）
  結果[i,l]   = Σ_k T1[i,k] C[k,l]   コスト ≈ N³
  合計 ≈ 2 N³

順序2: 素朴に3テンソルを一度に回す（全添字ループ）
  結果[i,l] = Σ_j Σ_k A[i,j] B[j,k] C[k,l]
  i,j,k,l の四重ループ            コスト ≈ N⁴
  合計 ≈ N⁴

N=1000 なら順序1は約20億回、素朴な四重ループは約1兆回で、約500倍の差です。ペア単位で縮約して中間結果を再利用するだけで、これだけ変わります。さらに、どのペアから始めるか（左結合か右結合か）でも中間テンソルの形状が変わるため、支配的なのは常にその順序で生じる最大の中間テンソルと、その生成にかかる演算量です。

最適順序を探す ── 行列連鎖積と同じ構造

最適な縮約順序を求める問題は、行列連鎖積（matrix-chain multiplication）の最適括弧付け問題を、一般のテンソルへ拡張したものです。ペア縮約のコストは「関与する全添字の次元の積」で見積もれます。たとえば T1[i,k] = Σ_j A[i,j]B[j,k] のコストは、内側で回る添字 i,j,k すべての積、すなわちおよそ dim(i)×dim(j)×dim(k) です。

テンソルの個数が十数個までなら、部分集合ごとの最適コストを表に持つ動的計画法で厳密な最適順序を求められます。しかし対象数が増えると探索空間が組合せ的に爆発するため、実用ライブラリは「各段で生成される中間テンソルが最も小さくなるペアを選ぶ」貪欲法や、コストを見積もりながら枝刈りする探索を併用します。ここで重要なのは、順序探索そのものの計算量は縮約本体に比べれば桁違いに小さいため、多少探索に時間をかけても、選んだ順序で本体のFLOPが大きく減れば十分に元が取れるという点です。演算強度と律速要因の見極めについては Roofline性能モデル の枠組みが有効です。

GEMMへの帰着 ── なぜ速くできるのか

順序が決まっても、各ペア縮約を素朴な多重ループで実装したのでは、キャッシュ効率が悪く演算器を使い切れません。実用実装の鍵は、任意のペア縮約を行列積（GEMM、一般行列乗算）へ帰着させ、高度に最適化されたBLASライブラリに任せることです。

横にスクロール

ペア縮約を GEMM に落とす3ステップ（例: "ijk,klm->ijlm"）

  1. 添字を「残す群」と「縮約する群」に並べ替え（permute / transpose）
     A: i,j を「行側」、k を「列側」へ  → (ij) × (k)
     B: k を「行側」、l,m を「列側」へ  → (k)  × (lm)

  2. 各群をまとめて2次元に変形（reshape）
     A → 行列 [ (i·j) , (k) ]
     B → 行列 [ (k) , (l·m) ]

  3. 通常の行列積を呼ぶ
     C_mat[(i·j),(l·m)] = A_mat × B_mat   ← BLAS の GEMM
     C_mat を [i,j,l,m] に戻す（reshape）

この「並べ替え → 2次元化 → GEMM → 形状復元」という定石は TTGT（Transpose-Transpose-GEMM-Transpose）と呼ばれます。縮約する添字をまとめて1つの内積次元に、残す添字を行・列にたたむことで、あらゆるペア縮約を単一のGEMMに写せます。GEMMは演算強度が高く（問題サイズを上げるほどデータ再利用が増える）、演算バウンド側で動くため、CPUのSIMD・FMAやGPUのTensorコアといったピーク性能に近い速度を引き出せます。

演算バウンドな密テンソル縮約に対し、対象が疎（ゼロが多い）な場合は事情が変わります。ゼロ要素まで律儀に掛けるのは無駄なので、非ゼロだけを扱う格納形式と専用カーネルが要ります。この違いは 疎行列の格納形式と演算 と同じ構図で、密は演算律速・疎はメモリアクセスと間接参照が律速になります。

テンソルネットワークでの応用

縮約順序最適化が単なる高速化を超えて「計算できるかどうか」を分けるのが、テンソルネットワーク（tensor network）です。量子多体系の波動関数や量子回路の振幅は、多数の小さなテンソルを共有添字（ボンド）でつないだネットワークとして表され、その振幅を求めることはネットワーク全体を1つのスカラーへ縮約することに等しくなります。

テンソルネットワーク縮約のコストを決めるもの

  ボンド次元 χ : テンソル同士をつなぐ内部添字の大きさ
  縮約木の形   : どの順でネットワークを畳むか（= 縮約順序）

  中間テンソルの階数が r なら、その要素数は χ^r のオーダー
  → 縮約木の中で最大の中間階数が、時間・空間計算量を支配する
  この「最大中間階数」はネットワークのグラフ構造の treewidth に対応する

ここでは中間テンソルの階数が指数の肩に乗るため、順序を誤ると中間テンソルの要素数が χ の高いべき乗となり、現実的なメモリに収まらなくなります。逆に、ネットワークのグラフ構造をうまく畳む順序（グラフ理論でいう木幅、treewidth の小さい分解に対応）を見つけられれば、同じネットワークが桁違いに小さいコストで縮約できます。量子回路シミュレーションで「どの順で縮約するか」の探索自体が主要な研究テーマになっているのはこのためで、最適順序の発見がシミュレーション可能な回路規模の限界を直接押し上げます。

こうした大規模テンソルネットワークの縮約を実際に走らせる段では、巨大な中間テンソルを複数ノード・複数GPUへ分散し、ペアGEMMをまたぐデータ移動をネットワーク越しに行います。ここでの通信設計は 並列FFTアルゴリズム の分散転置と同型の課題であり、相互接続の帯域とトポロジについては GPUクラスタの相互接続（NVLink・InfiniBand） が律速を左右します。

まとめ

• テンソル縮約は「共有添字について和を取り階数を下げる」演算で、行列積・内積・トレースを含む一般的枠組み。einsum記法は縮約添字と出力添字を並べるだけで任意の縮約を宣言でき、繰り返す添字が消えるという規則で意味が決まる。
• 3テンソル以上では縮約順序で総FLOP数と中間テンソルのメモリ量が桁違いに変わる。最適順序を求める問題は行列連鎖積と同じ構造で、支配的なのは各段で生じる最大の中間テンソルの大きさ。厳密解は動的計画法、大規模では貪欲法や枝刈り探索が使われる。
• 各ペア縮約は「並べ替え → 2次元化 → GEMM → 形状復元」（TTGT）で行列積に帰着でき、演算バウンドなGEMMとしてBLASやTensorコアのピーク性能を引き出せる。転置コストを避けるGETT方式もある。
• テンソルネットワークでは中間テンソルの階数が指数の肩に乗るため、順序最適化は速度だけでなく計算可能性そのものを決める。良い縮約順序（小さいtreewidthの分解）の発見が、シミュレーションできる系や量子回路の規模の限界を押し上げる。