Compilare somme di einstein in modo ottimale

Aug 18 2020

La sommatoria di Einstein è un modo conveniente per esprimere operazioni tensoriali che ha trovato la sua strada nelle librerie tensoriali come numpy, torch, tensorflow, ecc.

La sua flessibilità ci permette di rappresentare il prodotto di tre matrici,$X$,$Y$,$Z$di dimensioni$(a,b)$,$(b,c)$,$(c,d)$come

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

Tuttavia, quanto sopra si compila in qualcosa di simile

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

Questo nativo è quadratico nella dimensione delle matrici quando si fa semplicemente

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

Sarebbe semplicemente cubico.

Ci sono circa tre livelli di ottimizzazione di cui possiamo immaginare un'implementazione più intelligente einsumda eseguire.

  1. Scomporre un einsum di tensori$(x_1, \ldots, x_n)$in un einsum di coppie di tensori$x_1, x_2$,$e(x_1, x_2), x_3$, ecc. per ottimizzare il tempo di calcolo.

  2. Affidati all'associatività (ove applicabile) per scegliere quelle coppie con giudizio (questo è un classico problema di programmazione dinamica) e costruisci i tensori intermedi appropriati.

  3. Scopri le formule simili a Strassen per il calcolo del tensore specifico

Mentre 3 sembra chiaramente fuori portata, 1 e 2 sembrano realizzabili esattamente con un algoritmo ragionevolmente semplice. Tali algoritmi sono noti per le sommatorie di Einstein generiche? Sono stati studiati?

Risposte

2 smapers Aug 19 2020 at 16:33

Sembra che il problema generale di trovare l'ordine di contrazione ottimale sia NP-difficile [1]. Un documento recente sull'ottimizzazione approssimativa dell'ordine di contrazione, contenente riferimenti rilevanti, è [2].

[1] Chi-Chung, Lam, P. Sadayappan e Rephael Wenger. "Sull'ottimizzazione di una classe di cicli multidimensionali con riduzione per l'esecuzione parallela." Parallel Processing Letters 7.02 (1997): 157-168.

[2] Schindler, Frank e Adam Jermyn. "Algoritmi per l'ordinamento della contrazione della rete tensoriale". Apprendimento automatico: scienza e tecnologia (2020).