アインシュタインの合計を最適にコンパイルする

Aug 18 2020

アインシュタインの縮約は、numpy、torch、tensorflowなどのテンソルライブラリで使用されているテンソル演算を表現するための便利な方法です。

その柔軟性により、3つの行列の積を表すことができます。 $X$$Y$$Z$ 寸法の $(a,b)$$(b,c)$$(c,d)$ なので

X.Y.Z = einsum('ab,bc,cd->ad',X,Y,Z)

ただし、上記は次のようにコンパイルされます

for a_ in range(a):
  for d_ in range(d):
     res[a_,d_] = 0
     for b_ in range(b):
        for c_ in range(c):
           res[a_,d_] += X[a_,b_] * Y[b_,c_] * Z[c_, d_]

このネイティブは、単純に実行すると、行列のサイズが2次式になります。

einsum('ac,cd->ad',einsum('ab,bc'->'ac', X, Y), Z)

単に立方体になります。

実行するためのよりスマートな実装を想像できる最適化には、およそ3つのレベルがありeinsumます。

  1. テンソルの縮約記を分解します $(x_1, \ldots, x_n)$ テンソルのペアのeinsumに $x_1, x_2$$e(x_1, x_2), x_3$、などを使用して計算時間を最適化します。

  2. 結合性(該当する場合)に依存して、これらのペアを慎重に選択し(これは古典的な動的計画問題です)、適切な中間テンソルを構築します。

  3. 特定のテンソル計算のためのStrassenのような式を発見する

3は明らかに手の届かないところにあるように見えますが、1と2は、かなり単純なアルゴリズムで正確に達成できるようです。そのようなアルゴリズムは、一般的なアインシュタインの総和で知られていますか?彼らは研究されましたか?

回答

2 smapers Aug 19 2020 at 16:33

最適な収縮順序を見つけるための一般的な問題はNP困難であるようです[1]。収縮順序をほぼ最適化し、関連する参考文献を含む最近の論文は[2]です。

[1] Chi-Chung、Lam、P。Sadayappan、およびRephaelWenger。「並列実行のための削減を伴う多次元ループのクラスの最適化について。」並列処理レター7.02(1997):157-168。

[2]シンドラー、フランク、およびアダムジャーミン。「テンソルネットワーク収縮順序付けのアルゴリズム。」機械学習:科学技術(2020)。