Compilation optimale des sommes d'Einstein

Question

La sommation d'Einstein est un moyen pratique d'exprimer des opérations de tenseur qui ont trouvé leur place dans les bibliothèques de tenseurs telles que numpy, torch, tensorflow, etc.

Sa flexibilité permet de représenter le produit de trois matrices,$X$,$Y$,$Z$de dimensions$(a,b)$,$(b,c)$,$(c,d)$comme

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

Cependant, ce qui précède se compile en quelque chose comme

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_]

Ce natif est quadratique dans la taille des matrices en faisant simplement

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

Serait simplement cubique.

Il existe à peu près trois niveaux d'optimisation dont nous pouvons imaginer une mise en œuvre plus intelligente einsum.

1. Décomposer un einsum de tenseurs$(x_1, \ldots, x_n)$en un einsum de paires de tenseurs$x_1, x_2$,$e(x_1, x_2), x_3$, etc pour optimiser le temps de calcul.

2. Appuyez-vous sur l'associativité (le cas échéant) pour choisir judicieusement ces paires (il s'agit d'un problème de programmation dynamique classique) et construire les tenseurs intermédiaires appropriés.

3. Découvrez des formules de type Strassen pour le calcul spécifique du tenseur

Alors que 3 semble clairement hors de portée, 1 et 2 semblent être réalisables exactement avec un algorithme raisonnablement simple. De tels algorithmes sont-ils connus pour les sommations génériques d'Einstein ? Ont-ils été étudiés ?

Answer 1

Il semble que le problème général de trouver l'ordre de contraction optimal soit NP-difficile [1]. Un article récent sur l'optimisation approximative de l'ordre de contraction, et contenant des références pertinentes, est [2].

[1] Chi-Chung, Lam, P. Sadayappan et Rephael Wenger. "Sur l'optimisation d'une classe de boucles multidimensionnelles avec réduction pour une exécution parallèle." Lettres de traitement parallèle 7.02 (1997) : 157-168.

[2] Schindler, Frank et Adam Jermyn. "Algorithmes pour l'ordre de contraction du réseau de tenseurs." Apprentissage automatique : science et technologie (2020).