Оптимальное составление сумм Эйнштейна

Aug 18 2020

Суммирование Эйнштейна - удобный способ выразить тензорные операции, нашедшие свое применение в тензорных библиотеках, таких как numpy, torch, tensorflow и т. Д.

Его гибкость позволяет нам представить произведение трех матриц, $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_]

Этот родной код квадратичен по размеру матриц при простом выполнении

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

Было бы просто кубическим.

Мы можем представить себе более разумную реализацию примерно трех уровней оптимизации einsum.

  1. Разложите эйнсум тензоров $(x_1, \ldots, x_n)$ в эйнсум пар тензоров $x_1, x_2$, $e(x_1, x_2), x_3$и т. д. для оптимизации времени вычислений.

  2. Положитесь на ассоциативность (где это применимо), чтобы разумно выбрать эти пары (это классическая проблема динамического программирования) и построить соответствующие промежуточные тензоры.

  3. Откройте для себя формулы Штрассена для вычисления конкретного тензора

В то время как 3 кажется явно недостижимым, 1 и 2 кажутся достижимыми с помощью достаточно простого алгоритма. Известны ли такие алгоритмы для общих суммирований Эйнштейна? Были ли они изучены?

Ответы

2 smapers Aug 19 2020 at 16:33

Кажется, что общая задача поиска оптимального порядка сжатия NP-трудна [1]. Недавняя статья, посвященная приблизительной оптимизации порядка сжатия и содержащая соответствующие ссылки, - [2].

[1] Чи-Чунг, Лам, П. Садаяппан и Рефаэль Венгер. «Об оптимизации класса многомерных циклов с редукцией для параллельного выполнения». Письма о параллельной обработке 7.02 (1997): 157-168.

[2] Шиндлер, Франк и Адам Джермин. «Алгоритмы для упорядочения сжатия тензорной сети». Машинное обучение: наука и технологии (2020).