Resumen Técnico: LightGP

Enunciado del Problema

La inferencia de Procesos Gaussianos (GP) es un marco bayesiano principiado para regresión con incertidumbre calibrada, pero su costo computacional—$O(N^3)$ para inferencia exacta y memoria $O(N^2)$—ha limitado históricamente su aplicación a conjuntos de datos pequeños. Aunque aproximaciones escalables (por ejemplo, métodos de puntos de inducción, gradientes conjugados, interpolación de kernels estructurados) han extendido estos límites a cientos de miles de puntos, implementaciones dominantes como GPyTorch y GPflow dependen de marcos de aprendizaje profundo pesados (PyTorch, TensorFlow). Estos marcos imponen una huella de dependencias significativa (a menudo superior a 2 GB), una sobrecarga de despacho de Python no trivial por operación y un soporte limitado para la implementación en entornos móviles o embebidos. Además, las bibliotecas existentes a menudo carecen de implementaciones nativas en C++ libres de dependencias que puedan aprovechar aceleradores de hardware específicos como Apple Metal o NVIDIA CUDA sin la hinchazón de marcos.

Metodología

LightGP es una biblioteca C++17 libre de dependencias diseñada para abordar estas limitaciones. Proporciona enlaces para Python pero opera sobre un núcleo nativo en C++, eliminando la necesidad de PyTorch o TensorFlow. La biblioteca está estructurada en cuatro capas:

1. Tensor y BLAS: Una clase Tensor que almacena datos float32 en orden de fila. Despacha la multiplicación de matrices a Apple Accelerate (utilizando el coprocesor AMX en macOS) o OpenBLAS (en Linux). La factorización de Cholesky utiliza LAPACK spotrf.
2. Kernels: Una clase base abstracta Kernel soporta composición mediante operadores (+, *, Scale). Implementaciones concretas incluyen kernels RBF, Matérn, Periódicos y Lineales.
3. Backends de GPU:
   • Metal (Apple Silicon): Implementa la construcción de matrices de kernels en mosaico, un producto kernel-vector sin matriz fusionado, Cholesky bloqueado con GEMM de memoria de grupo de hilos en mosaico y resoluciones triangulares.
   • CUDA (NVIDIA): Refleja la funcionalidad de Metal utilizando cuBLAS, cuSOLVER y kernels personalizados para operaciones sin matriz.
   • Selección Automática: Una heurística Backend::Auto selecciona CPU o GPU basándose en el tamaño y dimensionalidad del problema.
4. Modelos de Inferencia:
   • GPExact: Soporta tres solucionadores: Cholesky Exacto ($O(N^3)$), Gradientes Conjugados (CG) con log-determinante de Lanczos estocástico (sin matriz, $O(N^2k)$), e Interpolación de Kernel Estructurado (SKI) mediante FFT ($O(N \log N)$).
   • GPSparse: Implementa la Energía Libre Variacional (VFE) de Titsias con inicialización de puntos de inducción por punto más lejano y optimización Adam. Incluye una característica de "refit en caliente" para omitir la reinicialización en ajustes repetidos.

La biblioteca se compila como una única biblioteca estática con cero dependencias externas y es instalable vía pip install lightgp.

Contribuciones Clave

1. Implementación Nativa en C++: Una implementación completa de cuatro rutas de inferencia GP (exacta, CG, VFE dispersa, SKI) con kernels composables, desplegable sin Python ni dependencias de marcos.
2. Kernels Optimizados para Hardware: Shaders de cálculo Metal ajustados a mano para Apple Silicon y kernels CUDA para GPUs NVIDIA. Cabe destacar que introduce un producto kernel-vector sin matriz fusionado que habilita inferencia con memoria $O(N)$ para CG, una característica no soportada nativamente por el backend Metal de PyTorch.
3. SKI Acelerado por FFT: Una implementación de SKI utilizando Apple Accelerate vDSP y cuFFT para lograr multiplicación matriz-vector $O(N \log N)$ a escala.
4. Benchmarks Multiplataforma: Evaluación sistemática que revela que en Apple Silicon, el coprocesor AMX basado en CPU (vía Accelerate) supera significativamente a Metal para factorización densa de Cholesky, un hallazgo que informa las heurísticas de selección de backend.
5. Enlaces Python de Cero Dependencias: Exposición completa de la API vía pybind11 con interoperabilidad NumPy, permitiendo la instalación sin la huella pesada de marcos de aprendizaje profundo.

Resultados

Los benchmarks se realizaron en un Apple M4 (Metal 3) y un escritorio Linux con un Intel i7-12700 y NVIDIA RTX 3060 (CUDA 12.0).

• Rendimiento vs. GPyTorch:
  • Apple M4 (CPU): LightGP es 2.6–8.7× más rápido que GPyTorch para inferencia GP exacta y ~1.5× más rápido para GP disperso en todas las escalas probadas. La aceleración se atribuye a la eliminación de la sobrecarga de despacho de Python y PyTorch.
  • Linux (CUDA): LightGP es 2.3–6.7× más rápido que GPyTorch para GP exacto hasta $N=2,048$. En $N=4,096$, GPyTorch cierra la brecha debido a su pipeline de autograd fusionado que amortiza la sobrecarga. Para GP disperso, LightGP es más rápido en $N=10,000$ (1.7×), pero GPyTorch gana en $N=50,000$ debido a tensores de dispositivo persistentes y cálculo de gradientes compilado.
• Operaciones sin Matriz: En Apple M4, el producto kernel-vector sin matriz fusionado de Metal logra una aceleración de 32× sobre la ruta explícita de CPU en $N=20,000$ y es la única opción viable más allá de $N=30,000$ en 16 GB de memoria. En Linux, CUDA logra hasta 259× de aceleración en $N=10,000$.
• Escalado de Memoria: CG sin matriz, VFE dispersa y SKI mantienen un uso de memoria $O(N)$. Por el contrario, Cholesky exacto ($O(N^2)$) agota 16 GB de RAM en $N \approx 42,000$–45,000 en Mac y 12 GB de VRAM en $N \approx 35,000$ en Linux.
• Precisión: Las predicciones de backends de CPU y GPU coinciden dentro de $10^{-5}$ para media y varianza en conjuntos de datos estándar (Motocicleta, Mauna Loa, kin40k).

Significado

LightGP demuestra que la inferencia GP nativa en C++ puede igualar o superar el rendimiento de bibliotecas basadas en marcos mientras requiere cero dependencias externas. Al eliminar la sobrecarga de despacho de Python y los marcos de aprendizaje profundo, LightGP ofrece ganancias de rendimiento consistentes, particularmente en Apple Silicon donde el coprocesor AMX se aprovecha directamente. La biblioteca habilita inferencia GP a escalas ($N > 100,000$) previamente imprácticas en hardware de consumo mediante métodos sin matriz de memoria $O(N)$ y SKI acelerado por FFT.

El artículo destaca una idea crítica de hardware: en Apple Silicon, la CPU (vía Accelerate/AMX) a menudo es superior a la GPU (Metal) para tareas de álgebra lineal densa como la factorización de Cholesky, un matiz que las heurísticas de LightGP gestionan automáticamente. Mientras GPyTorch sigue siendo competitivo o superior a escalas muy grandes ($N \ge 50,000$) en CUDA debido a su autograd maduro e infraestructura de tensores persistentes, LightGP llena un nicho para inferencia GP ligera, de alto rendimiento y desplegable sin la carga de una pila completa de aprendizaje profundo.