Compilando somas de einstein de forma otimizada

Question

A soma de Einstein é uma maneira conveniente de expressar operações de tensores que encontraram seu caminho em bibliotecas de tensores como numpy, archote, tensorflow, etc.

Sua flexibilidade nos permite representar o produto de três matrizes,$X$,$Y$,$Z$de dimensões$(a,b)$,$(b,c)$,$(c,d)$Como

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

No entanto, o acima compila para algo como

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

Este nativo é quadrático no tamanho das matrizes ao simplesmente fazer

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

Seria meramente cúbico.

Existem aproximadamente três níveis de otimização que podemos imaginar uma implementação mais inteligente einsumpara executar.

1. Decompor uma einsum de tensores$(x_1, \ldots, x_n)$em uma einsum de pares de tensores$x_1, x_2$,$e(x_1, x_2), x_3$, etc para otimizar o tempo de computação.

2. Confie na associatividade (quando aplicável) para escolher esses pares criteriosamente (este é um problema clássico de programação dinâmica) e construir os tensores intermediários apropriados.

3. Descubra fórmulas do tipo Strassen para o cálculo específico do tensor

Enquanto 3 parece claramente fora de alcance, 1 e 2 parecem que poderiam ser alcançados exatamente com um algoritmo razoavelmente direto. Esses algoritmos são conhecidos por somas genéricas de einstein? Eles foram estudados?

Answer 1

Parece que o problema geral de encontrar a ordem de contração ótima é NP-difícil [1]. Um artigo recente sobre como otimizar aproximadamente a ordem de contração e contendo referências relevantes é [2].

[1] Chi-Chung, Lam, P. Sadayappan e Rephael Wenger. "Sobre a otimização de uma classe de loops multidimensionais com redução para execução paralela." Parallel Processing Letters 7.02 (1997): 157-168.

[2] Schindler, Frank e Adam Jermyn. "Algoritmos para ordenação de contração de rede tensorial." Machine Learning: Ciência e Tecnologia (2020).