Resumen Técnico: torch-sla

Declaración del Problema

El álgebra lineal dispersa diferenciable es un requisito fundamental para el aprendizaje automático científico, incluidas las redes neuronales de grafos (GNN), los operadores neuronales en mallas no estructuradas y las tuberías de simulación diferenciables. Sin embargo, el ecosistema de PyTorch carece de una biblioteca unificada para soportar estas cargas de trabajo. El soporte nativo de PyTorch está fragmentado: torch.sparse proporciona kernels de bajo nivel y un spsolve no diferenciable y solo para CPU (un envoltorio de SuperLU), mientras que torch.linalg está restringido a matrices densas.

Las soluciones alternativas existentes son insuficientes para escenarios de entrenamiento a gran escala o complejos:

1. Ruptura del Grafo: Llamar a solucionadores externos (por ejemplo, cuDSS, PETSc) rompe el grafo de autograd, obligando a los usuarios a descartar gradientes o escribir envoltorios personalizados torch.autograd.Function para cada problema.
2. Diferenciación Ineficiente: Retropropagar ingenuamente a través de $k$ iteraciones de un solucionador iterativo crea un grafo computacional de $O(k)$ nodos. Para un problema con $10^6$ grados de libertad (DOF) y 1.000 iteraciones, almacenar vectores intermedios requiere aproximadamente 80 GB de memoria, haciendo inviable el entrenamiento a gran escala.
3. Límites de Escalabilidad: Los problemas con $10^7$–$10^9$ incógnitas requieren descomposición de dominio distribuida. Los primitivos distribuidos existentes de PyTorch no soportan nativamente el intercambio de halo transpuesto requerido para los adjuntos de operadores dispersos en la pasada inversa.
4. Fragmentación del Backend: Los solucionadores directos (por ejemplo, cuDSS) son rápidos para problemas pequeños ($<10^5$ DOF) pero sufren de llenado $O(n^{1.5})$ y agotamiento de memoria para sistemas más grandes. Los solucionadores iterativos escalan mejor pero incurren en sobrecarga de lanzamiento de kernels en problemas pequeños. No existe una interfaz unificada para despachar entre estos backends heterogéneos basándose en el dispositivo y el tamaño del problema.

Metodología

Los autores presentan torch-sla, una biblioteca de código abierto que llena este vacío mediante tres componentes principales:

1. Abstracción Unificada de Backend

torch-sla expone una única API consciente de autograd (.solve, .matvec, .eigsh, .det) que despacha a través de cinco backends intercambiables basándose en el dispositivo y el tamaño del problema:

• CPU: SciPy (SuperLU, UMFPACK, CG, BiCGStab) y Eigen (CG, BiCGStab).
• GPU: NVIDIA cuDSS (LU, Cholesky, LDLT), CuPy (iterativo y directo) y un solucionador iterativo nativo de PyTorch.
• Despacho Automático: El sistema selecciona automáticamente el backend: solucionadores directos para problemas pequeños a medianos (por ejemplo, $<2 \times 10^6$ DOF en GPU) y solucionadores iterativos para problemas más grandes para evitar el desbordamiento de memoria. Se detecta la simetría y la definida positiva para optimizar la factorización (por ejemplo, actualizar LU a Cholesky).

La biblioteca admite cuatro tipos de tensores dispersos organizados por localidad y agrupación: SparseTensor (matriz única/patrón compartido), SparseTensorList (patrones distintos por lote) y sus variantes distribuidas DSparseTensor y DSparseTensorList.

2. Marco de Diferenciación Adjoint

Para evitar la complejidad de memoria y grafo de $O(k)$ de la retropropagación ingenua, torch-sla emplea un enfoque basado en el teorema de la función implícita.

• Sistemas Lineales: Para $Ax=b$, la pasada inversa resuelve un único sistema adjunto $A^\top \lambda = \partial L / \partial x$. El gradiente con respecto a la matriz se calcula como $-\lambda_i x_j$ solo en el patrón de dispersión.
• Sistemas No Lineales: Para $F(u, \theta)=0$, la pasada inversa requiere una resolución lineal adjunta ($J^\top \lambda = \partial L / \partial u$) y un producto vector-Jacobiano, independientemente del número de iteraciones no lineales.
• Problemas de Valores Propios: Los gradientes para los valores propios utilizan el teorema de Hellmann–Feynman ($\partial \lambda / \partial A_{ij} = v_i v_j$), requiriendo solo un producto externo en el patrón de dispersión.
• Complejidad: Este marco reduce el grafo de autograd a $O(1)$ nodos y la memoria a $O(n + \text{nnz})$, independientemente del número de iteraciones del solucionador ($k$).

3. Capa Distribuida con Intercambio de Halo

Para la ejecución multi-GPU, torch-sla implementa descomposición de dominio donde cada proceso posee un bloque de filas y metadatos de halo.

• Pasada Directa: Los procesos intercambian valores de halo antes de las operaciones locales de Matriz Dispersa-Vector (SpMV).
• Pasada Inversa: La biblioteca implementa un intercambio de halo transpuesto. Mientras que la pasada directa dispersa valores propios a halos vecinos, la pasada inversa suma los gradientes de halo entrantes de nuevo a las entradas de frontera del proceso propietario.
• Implementación: Tanto las comunicaciones directas como inversas se envuelven como primitivas torch.autograd.Function usando torch.distributed, asegurando diferenciabilidad de extremo a extremo en resoluciones distribuidas.

Contribuciones Clave

1. API Unificada: Una única biblioteca que proporciona interfaces conscientes de autograd para solucionadores directos, iterativos, no lineales y de valores propios en cinco backends (SciPy, Eigen, cuDSS, CuPy, PyTorch-nativo).
2. Marco de Adjoint Escalable: Un método de diferenciación de grafo $O(1)$ que elimina el cuello de botella de memoria de la retropropagación de solucionadores iterativos, permitiendo el entrenamiento en problemas con millones de incógnitas.
3. Autograd Distribuido: La primera implementación de solucionadores dispersos distribuidos en PyTorch que soporta nativamente la diferenciación de extremo a extremo mediante intercambio de halo transpuesto.
4. Extensibilidad: Una abstracción de backend que permite integrar bibliotecas externas (por ejemplo, PETSc, Trilinos, hypre) o precondicionadores aprendidos implementando una interfaz común.

Resultados Experimentales

Los autores evaluaron torch-sla en benchmarks de Poisson 2D utilizando GPUs NVIDIA H200:

• Escalabilidad: La biblioteca escala a 169M DOF en una sola GPU y 400M DOF en tres GPUs.
• Rendimiento: Para un problema de 1M DOF, el backend CG de torch-sla (190 ms) supera a cuDSS (7.3 s) en 38$\times$. Por encima de 2M DOF, los solucionadores directos se quedan sin memoria, mientras que los solucionadores iterativos mantienen una escalabilidad casi lineal.
• Eficiencia de Memoria: En una comparación de retropropagación ingenua frente al método adjunto para 1.000 iteraciones de CG, el enfoque ingenuo requirió 64 GB (OOM a 2.000 iteraciones), mientras que el método adjunto mantuvo un uso constante de **328 MB**, logrando una reducción de memoria de 195$\times$.
• Verificación de Gradientes: Los gradientes se verificaron frente a diferencias finitas para problemas no lineales y de valores propios, mostrando errores relativos por debajo de $10^{-5}$.
• Aplicación de Extremo a Extremo: Una tarea de aprendizaje de coeficientes inversos (recuperar conductividad espacialmente variable en una ecuación de Poisson) convergió exitosamente a un error L2 relativo del 0,23% en menos de un minuto utilizando un bucle Adam estándar de PyTorch.

Significado

El artículo afirma que torch-sla cierra una brecha de larga data en el ecosistema de PyTorch, proporcionando capacidades análogas a jax.lax.custom_linear_solve y jax.scipy.sparse.linalg de JAX, pero dentro de una biblioteca unificada de PyTorch. Permite a los investigadores:

• Entrenar modelos que involucran resoluciones lineales dispersas sin descartar gradientes ni portar código a JAX.
• Escalar cargas de trabajo de aprendizaje automático científico a cientos de millones de grados de libertad en clústeres de GPU distribuidos.
• Tratar a los solucionadores dispersos como capas diferenciables, facilitando la investigación de precondicionadores aprendidos y la optimización de extremo a extremo de simulaciones físicas.

Los autores señalan limitaciones, incluida la dependencia actual del precondicionamiento de Jacobi para el backend nativo de PyTorch (lo que limita la velocidad de convergencia para sistemas muy grandes) y un enfoque en benchmarks de Poisson 2D, identificando las EDP 3D y los sistemas indefinidos como trabajo futuro.