Resumo Técnico: torch-sla

Declaração do Problema

Álgebra linear esparsa diferenciável é um requisito fundamental para aprendizado de máquina científico, incluindo redes neurais em grafos (GNNs), operadores neurais em malhas não estruturadas e pipelines de simulação diferenciável. No entanto, o ecossistema PyTorch carece de uma biblioteca unificada para suportar essas cargas de trabalho. O suporte nativo do PyTorch está fragmentado: torch.sparse fornece kernels de baixo nível e um spsolve não diferenciável e exclusivo para CPU (um wrapper SuperLU), enquanto torch.linalg é restrito a matrizes densas.

As soluções alternativas existentes são insuficientes para cenários de treinamento em grande escala ou complexos:

1. Quebra do Grafo: Chamar solvers externos (por exemplo, cuDSS, PETSc) quebra o grafo de autograd, forçando os usuários a descartar gradientes ou escrever wrappers torch.autograd.Function personalizados para cada problema.
2. Diferenciação Ineficiente: A retropropagação ingênua através de $k$ iterações de um solver iterativo cria um grafo computacional de $O(k)$ nós. Para um problema com $10^6$ graus de liberdade (DOF) e 1.000 iterações, armazenar vetores intermediários requer aproximadamente 80 GB de memória, tornando o treinamento em grande escala inviável.
3. Limites de Escalabilidade: Problemas com $10^7$–$10^9$ incógnitas exigem decomposição de domínio distribuída. As primitivas distribuídas existentes do PyTorch não suportam nativamente a troca de halo transposta necessária para adjuntos de operadores esparsos na passagem de retropropagação.
4. Fragmentação de Backend: Solvers diretos (por exemplo, cuDSS) são rápidos para problemas pequenos ($<10^5$ DOF), mas sofrem com preenchimento $O(n^{1.5})$ e exaustão de memória para sistemas maiores. Solvers iterativos escalam melhor, mas incorrem em sobrecarga de lançamento de kernel em problemas pequenos. Não existe uma interface unificada para distribuir entre esses backends heterogêneos com base no dispositivo e no tamanho do problema.

Metodologia

Os autores apresentam torch-sla, uma biblioteca de código aberto que preenche essa lacuna por meio de três componentes principais:

1. Abstração Unificada de Backend

O torch-sla expõe uma única API consciente de autograd (.solve, .matvec, .eigsh, .det) que distribui entre cinco backends intercambiáveis com base no dispositivo e no tamanho do problema:

• CPU: SciPy (SuperLU, UMFPACK, CG, BiCGStab) e Eigen (CG, BiCGStab).
• GPU: NVIDIA cuDSS (LU, Cholesky, LDLT), CuPy (iterativo e direto) e um solver iterativo nativo do PyTorch.
• Auto-Distribuição: O sistema seleciona automaticamente o backend: solvers diretos para problemas pequenos a médios (por exemplo, $<2 \times 10^6$ DOF em GPU) e solvers iterativos para problemas maiores para evitar estouro de memória. Simetria e definição positiva são detectadas para otimizar a fatoração (por exemplo, atualizar LU para Cholesky).

A biblioteca suporta quatro tipos de tensores esparsos organizados por localidade e agrupamento: SparseTensor (matriz única/padrão compartilhado), SparseTensorList (padrões distintos por lote) e suas variantes distribuídas DSparseTensor e DSparseTensorList.

2. Framework de Diferenciação Adjoint

Para evitar a complexidade de memória e grafo $O(k)$ da retropropagação ingênua, o torch-sla emprega uma abordagem baseada no teorema da função implícita.

• Sistemas Lineares: Para $Ax=b$, a passagem de retropropagação resolve um único sistema adjunto $A^\top \lambda = \partial L / \partial x$. O gradiente em relação à matriz é computado como $-\lambda_i x_j$ apenas no padrão de esparsidade.
• Sistemas Não Lineares: Para $F(u, \theta)=0$, a passagem de retropropagação requer uma resolução linear adjunta ($J^\top \lambda = \partial L / \partial u$) e um produto vetor-Jacobiano, independentemente do número de iterações não lineares.
• Problemas de Autovalor: Gradientes para autovalores utilizam o teorema de Hellmann–Feynman ($\partial \lambda / \partial A_{ij} = v_i v_j$), exigindo apenas um produto externo no padrão de esparsidade.
• Complexidade: Este framework reduz o grafo de autograd para $O(1)$ nós e a memória para $O(n + \text{nnz})$, independentemente do número de iterações do solver ($k$).

3. Camada Distribuída com Troca de Halo

Para execução multi-GPU, o torch-sla implementa decomposição de domínio onde cada processo possui um bloco de linha e metadados de halo.

• Passagem Direta: Os processos trocam valores de halo antes das operações locais de Matriz Esparsa-Vetor (SpMV).
• Passagem de Retropropagação: A biblioteca implementa uma troca de halo transposta. Enquanto a passagem direta dispersa valores possuídos para halos vizinhos, a passagem de retropropagação soma gradientes de halo recebidos de volta às entradas de fronteira do processo proprietário.
• Implementação: Tanto as comunicações diretas quanto as de retropropagação são envoltas como primitivas torch.autograd.Function usando torch.distributed, garantindo diferenciabilidade ponta a ponta em resoluções distribuídas.

Contribuições Principais

1. API Unificada: Uma única biblioteca fornecendo interfaces conscientes de autograd para solvers diretos, iterativos, não lineares e de autovalor em cinco backends (SciPy, Eigen, cuDSS, CuPy, PyTorch-nativo).
2. Framework Adjoint Escalável: Um método de diferenciação de grafo $O(1)$ que elimina o gargalo de memória da retropropagação de solver iterativo, permitindo treinamento em problemas com milhões de incógnitas.
3. Autograd Distribuído: A primeira implementação de solvers esparsos distribuídos no PyTorch que suporta nativamente diferenciação ponta a ponta via troca de halo transposta.
4. Extensibilidade: Uma abstração de backend que permite integrar bibliotecas externas (por exemplo, PETSc, Trilinos, hypre) ou pré-condicionadores aprendidos implementando uma interface comum.

Resultados Experimentais

Os autores avaliaram o torch-sla em benchmarks de Poisson 2D usando GPUs NVIDIA H200:

• Escalabilidade: A biblioteca escala para 169M DOF em uma única GPU e 400M DOF em três GPUs.
• Desempenho: Para um problema de 1M DOF, o backend CG do torch-sla (190 ms) supera o cuDSS (7,3 s) em 38$\times$. Acima de 2M DOF, solvers diretos esgotam a memória, enquanto solvers iterativos mantêm escalabilidade quase linear.
• Eficiência de Memória: Em uma comparação entre retropropagação ingênua e o método adjunto para 1.000 iterações CG, a abordagem ingênua exigiu 64 GB (OOM em 2.000 iterações), enquanto o método adjunto manteve um uso constante de **328 MB**, alcançando uma redução de memória de 195$\times$.
• Verificação de Gradiente: Gradientes foram verificados contra diferenças finitas para problemas não lineares e de autovalor, mostrando erros relativos abaixo de $10^{-5}$.
• Aplicação Ponta a Ponta: Uma tarefa de aprendizado de coeficiente inverso (recuperar condutividade espacialmente variável em uma equação de Poisson) convergiu com sucesso para um erro L2 relativo de 0,23% em menos de um minuto usando um loop Adam padrão do PyTorch.

Significância

O artigo afirma que o torch-sla fecha uma lacuna de longa data no ecossistema PyTorch, fornecendo capacidades análogas a jax.lax.custom_linear_solve e jax.scipy.sparse.linalg do JAX, mas dentro de uma biblioteca unificada do PyTorch. Isso permite que pesquisadores:

• Treinem modelos envolvendo resoluções lineares esparsas sem descartar gradientes ou portar código para JAX.
• Escalem cargas de trabalho de aprendizado de máquina científico para centenas de milhões de graus de liberdade em clusters distribuídos de GPU.
• Tratem solvers esparsos como camadas diferenciáveis, facilitando a pesquisa de pré-condicionadores aprendidos e otimização ponta a ponta de simulações físicas.

Os autores observam limitações, incluindo a dependência atual de pré-condicionamento Jacobi para o backend nativo do PyTorch (o que limita a velocidade de convergência para sistemas muito grandes) e um foco em benchmarks de Poisson 2D, com EDPs 3D e sistemas indefinidos identificados como trabalho futuro.