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

はじめに

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

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

 いま、ベクトルxyの内積を考える。

(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つのベクトルの内積の場合を考える。コードは以下の通り。

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

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

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

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

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

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

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

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

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

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

5行目のeinsumに渡している引数"i,j->ij"は式(6)の添え字と対応している。すなわち、右辺にあるx_{i}の添え字iy_{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を用いると次のよう書ける。

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

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

(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*}

コードは以下の通り。

まとめ

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

Kumada Seiya

Kumada Seiya

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

最近の記事

  • 関連記事
  • おすすめ記事
  • 特集記事

アーカイブ

カテゴリー

PAGE TOP