2020.08.20
Python

アインシュタインの縮約記法

はじめに

　今回は、Pythonの数値計算モジュールNumPyが提供している関数einsumの使い方を解説する。einsumはアインシュタイン（Einstein）の縮約記法を実装した関数である。テンソル間の複雑な計算を1行で書くことができる。前回実装した変分推論においても用いたので、この機会に紹介したい。

アインシュタインの縮約記法とは

　いま、ベクトル $x$ と $y$ の内積を考える。

(1) $\begin{equation*} x\cdot y=\sum_{i}x_i\;y_i \end{equation*}$

ここで、上の式の和の記号 $\sum$ を省略し、同じ文字の添え字が2つ続く場合はそれらについて和を取ると約束する。

(2) $\begin{equation*} x_i\;y_i\equiv\sum_{i}x_i\;y_i \end{equation*}$

左辺の簡略化した書き方をアインシュタインの縮約記法と呼ぶ。この記法を最初に考案したのはアインシュタインである。厳密な話をすると、アインシュタインが一般相対論を定式化する際に用いた数学（リーマン幾何学）では2種類の添え字を扱う。上付き添え字 $x^i$ と下付き添え字 $x_i$ である。前者を反変ベクトル、後者を共変ベクトルと呼び、これらは区別される。そして、同じ上付き添え字と下付き添え字が続く場合の和を簡略化する記法として、上の縮約記法が導入された。

(3) $\begin{equation*} x_i\;y^i\equiv\sum_{i}x_i\;y^i \end{equation*}$

本解説では簡単のため、最初に説明した方（下付き添え字だけの場合）の記法もアインシュタインの縮約記法と呼ぶことにして、話を進める。

NumPy.einsumの使い方

　最初に、上で見た2つのベクトルの内積の場合を考える。コードは以下の通り。

    # 2つのベクトルを定義する。
    x = np.array([1, 2, 3])
    y = np.array([3, 2, 1])

    # 内積を計算する。dotを使用した場合。
    z = x.dot(y)
    assert(z == 10)

    # einsumの場合。
    u = np.einsum("i,i", x, y)
    assert(u == z)

# 2つのベクトルを定義する。

x = np.array([1, 2, 3])

y = np.array([3, 2, 1])

# 内積を計算する。dotを使用した場合。

z = x.dot(y)

assert(z == 10)

# einsumの場合。

u = np.einsum("i,i", x, y)

assert(u == z)

10行目でeinsumに渡している第1引数の文字列"i,i"は、式(2)の添え字に一致している。すなわち、式(2)の $x_{i}$ の添え字 $i$ と $y_{i}$ の添え字 $i$ とを対応させ和を取ることをeinsumに教えている。

　次に、2つの行列 $A,B$ の積を考える。縮約記法で書くと以下のようになる。

(4) $\begin{equation*} (AB)_{ij}=A_{in}B_{nj} \end{equation*}$

ここで、 $A_{ij}$ は行列 $A$ の $(i,j)$ 成分を表す。コードは以下の通りである。

   # 2つの行列を定義する。
    A = np.array([[1, 2], [2, 1]])
    B = np.array([[0, 3], [3, 0]])

    # 積を計算する。matmulを使用した場合。
    C = np.matmul(A, B)
    assert(np.all(C == np.array([[6, 3], [3, 6]])))

    # einsumの場合。
    D = np.einsum("in,nj->ij", A, B)
    assert(np.all(C == D))

# 2つの行列を定義する。

A = np.array([[1, 2], [2, 1]])

B = np.array([[0, 3], [3, 0]])

# 積を計算する。matmulを使用した場合。

C = np.matmul(A, B)

assert(np.all(C == np.array([[6, 3], [3, 6]])))

# einsumの場合。

D = np.einsum("in,nj->ij", A, B)

assert(np.all(C == D))

10行目でeinsumの第1引数に渡している文字列"in,nj->ij"は式(4)の添え字と一致している。すなわち、式(4)の右辺にある $A_{in}$ の添え字 $in$ と $B_{nj}$ の添え字 $nj$ から、左辺の $(AB)_{ij}$ の添え字 $ij$ を作る操作であることをeinsumに教えている。

　今度は、2つのベクトル $x,y$ から行列を作る演算を考える。

(5) $\begin{equation*} A=x\;y^T \end{equation*}$

$x$ が $N$ 次元、 $y$ が $M$ 次元のとき、行列 $A$ のサイズは $N\times M$ となる。成分で書くと

(6) $\begin{equation*} A_{ij}=x_i\;y_j \end{equation*}$

となる。この式には和は出現しないが、以下のようにeinsumを用いて書くことができる。

    # 2つのベクトルを定義する。
    x = np.array([1, 2, 3])
    y = np.array([1, 3])

    z = np.einsum("i,j->ij", x, y)
    assert(np.all(z == np.array([[1, 3], [2, 6], [3, 9]])))

# 2つのベクトルを定義する。

x = np.array([1, 2, 3])

y = np.array([1, 3])

z = np.einsum("i,j->ij", x, y)

assert(np.all(z == np.array([[1, 3], [2, 6], [3, 9]])))

5行目のeinsumに渡している引数"i,j->ij"は式(6)の添え字と対応している。すなわち、右辺にある $x_{i}$ の添え字 $i$ と $y_{j}$ の添え字 $j$ から、左辺の $A_{ij}$ の添え字 $ij$ を作る操作であることをeinsumに教えている。本来の縮約記法は和記号を省略することを意味したが、NumPy.einsumの適用範囲は拡張されている。さらに、次の例を考えてみたい。

(7) $\begin{equation*} A=\sum_n x_n\;y_n^T \end{equation*}$

ここで、 $x_n,y_n$ はベクトルである。成分表示すると

(8) $\begin{equation*} A_{ij}=x_{ni}\;y_{nj} \end{equation*}$

となる。ここで、添え字 $n$ について縮約記法を用いた。einsumを用いると次のよう書ける。

   # 2つの行列を定義する。
    x = np.array([[1, 2, 3], [1, 2, 3]])
    y = np.array([[1, 3], [1, 3]])

    z = np.einsum("ni,nj->ij", x, y)
    assert(np.all(z == np.array([[2, 6], [4, 12], [6, 18]])))

# 2つの行列を定義する。

x = np.array([[1, 2, 3], [1, 2, 3]])

y = np.array([[1, 3], [1, 3]])

z = np.einsum("ni,nj->ij", x, y)

assert(np.all(z == np.array([[2, 6], [4, 12], [6, 18]])))

もう、説明は必要ないであろう。

　最後に、複雑な例を挙げて終わりにしたい。

(9) $\begin{equation*} A_{ij}=x_{iabc}\;y_{abc}\;z_{abcj} \end{equation*}$

上の式は次式の縮約記法である。

(10) $\begin{equation*} A_{ij}=\sum_{abc}x_{iabc}\;y_{abc}\;z_{abcj} \end{equation*}$

コードは以下の通り。

    x = np.arange(16).reshape(2, 2, 2, 2)
    y = np.arange(8).reshape(2, 2, 2)
    z = np.arange(16).reshape(2, 2, 2, 2)
    w = np.einsum("iabc,abc,abcj->ij", x, y, z)
    assert((2, 2) == w.shape)

x = np.arange(16).reshape(2, 2, 2, 2)

y = np.arange(8).reshape(2, 2, 2)

z = np.arange(16).reshape(2, 2, 2, 2)

w = np.einsum("iabc,abc,abcj->ij", x, y, z)

assert((2, 2) == w.shape)

まとめ

　今回は、PyNum.einsumの使い方を説明した。通常であればループの入れ子を用いて書くことになる計算を1行で書くことができる。しかし、上のサンプルコードでも示したが、内積にはdot関数が、行列同士の積にはmatmulが用意されている。同じことをする関数がすでに存在するのであれば、既存関数を使った方が高速であるようだ。速度比較についてはこちらのサイトが詳しいので参照してほしい。

Kumada Seiya

仕事であろうとなかろうと勉強し続ける、その結果”中身”を知ったエンジニアになれる

プロダクティビティ向上はじめました

見積もりと心理的安全性

アインシュタインの縮約記法

はじめに

アインシュタインの縮約記法とは

NumPy.einsumの使い方

まとめ

Kumada Seiya

最近の記事

LangExtract

LLMとMCPの連携

LLMの出力の構造化データへの変換

MCP Python SDKによるMCPサーバの構築

OpenAI Agents SDK

自動翻訳ライブラリdeep-translator

PythonのProgressBar（プログレスバー）のデフォルトの挙動の比較をしてみた

Google Vision APIでOCR

アーカイブ

カテゴリー