Resumo Técnico: LightGP

Declaração do Problema

A inferência de Processos Gaussianos (GP) é um framework bayesiano principiado para regressão com incerteza calibrada, mas seu custo computacional — $O(N^3)$ para inferência exata e memória $O(N^2)$ — historicamente limitou sua aplicação a conjuntos de dados pequenos. Embora aproximações escaláveis (por exemplo, métodos de pontos de indução, gradientes conjugados, interpolação de kernel estruturada) tenham expandido esses limites para centenas de milhares de pontos, implementações dominantes como GPyTorch e GPflow dependem de frameworks pesados de aprendizado profundo (PyTorch, TensorFlow). Esses frameworks impõem dependências significativas (frequentemente excedendo 2 GB), sobrecarga de despacho Python não trivial por operação e suporte limitado para implantação em ambientes móveis ou embarcados. Além disso, bibliotecas existentes frequentemente carecem de implementações nativas em C++ livres de dependências que possam aproveitar aceleradores de hardware específicos como Apple Metal ou NVIDIA CUDA sem o inchaço de frameworks.

Metodologia

LightGP é uma biblioteca C++17 livre de dependências projetada para abordar essas limitações. Ela fornece bindings para Python, mas opera sobre um núcleo nativo em C++, eliminando a necessidade de PyTorch ou TensorFlow. A biblioteca é estruturada em quatro camadas:

1. Tensor e BLAS: Uma classe Tensor armazenando dados float32 em ordem de linha. Ela despacha multiplicação de matrizes para Apple Accelerate (usando o coprocessor AMX no macOS) ou OpenBLAS (no Linux). A fatoração de Cholesky utiliza spotrf do LAPACK.
2. Kernels: Uma classe base abstrata Kernel suporta composição via operadores (+, *, Scale). Implementações concretas incluem kernels RBF, Matérn, Periódico e Linear.
3. Backends de GPU:
   • Metal (Apple Silicon): Implementa construção de matriz de kernel em blocos, um produto kernel-vetor fundido sem matriz, Cholesky bloqueado com GEMM de memória de threadgroup em blocos e soluções triangulares.
   • CUDA (NVIDIA): Espelha a funcionalidade do Metal usando cuBLAS, cuSOLVER e kernels personalizados para operações sem matriz.
   • Seleção Automática: Uma heurística Backend::Auto seleciona CPU ou GPU com base no tamanho e dimensionalidade do problema.
4. Modelos de Inferência:
   • GPExact: Suporta três solucionadores: Cholesky Exato ($O(N^3)$), Gradientes Conjugados (CG) com log-determinante de Lanczos estocástico (sem matriz, $O(N^2k)$) e Interpolação de Kernel Estruturada (SKI) via FFT ($O(N \log N)$).
   • GPSparse: Implementa a Energia Livre Variacional (VFE) de Titsias com inicialização de pontos de indução de ponto mais distante e otimização Adam. Inclui um recurso de "reajuste a quente" para pular re-inicialização em ajustes repetidos.

A biblioteca compila como uma única biblioteca estática com zero dependências externas e é instalável via pip install lightgp.

Principais Contribuições

1. Implementação Nativa em C++: Uma implementação completa de quatro caminhos de inferência GP (exato, CG, VFE esparsa, SKI) com kernels composáveis, implantável sem Python ou dependências de frameworks.
2. Kernels Otimizados para Hardware: Shaders de computação Metal ajustados manualmente para Apple Silicon e kernels CUDA para GPUs NVIDIA. Notavelmente, introduz um produto kernel-vetor fundido sem matriz que permite inferência com memória $O(N)$ para CG, um recurso não suportado nativamente pelo backend Metal do PyTorch.
3. SKI Acelerado por FFT: Uma implementação de SKI usando vDSP do Apple Accelerate e cuFFT para alcançar multiplicação matriz-vetor $O(N \log N)$ em escala.
4. Benchmarks Cross-Plataforma: Avaliação sistemática revelando que no Apple Silicon, o coprocessor AMX baseado em CPU (via Accelerate) supera significativamente o Metal para fatoração densa de Cholesky, uma descoberta que informa heurísticas de seleção de backend.
5. Bindings Python sem Dependências: Exposição completa da API via pybind11 com interoperabilidade NumPy, permitindo instalação sem o peso pesado de frameworks de aprendizado profundo.

Resultados

Os benchmarks foram conduzidos em um Apple M4 (Metal 3) e um desktop Linux com Intel i7-12700 e NVIDIA RTX 3060 (CUDA 12.0).

• Desempenho vs. GPyTorch:
  • Apple M4 (CPU): LightGP é 2,6–8,7× mais rápido que GPyTorch para inferência GP exata e ~1,5× mais rápido para GP esparsa em todas as escalas testadas. O aumento de velocidade é atribuído à eliminação da sobrecarga de despacho Python e PyTorch.
  • Linux (CUDA): LightGP é 2,3–6,7× mais rápido que GPyTorch para GP exata até $N=2.048$. Em $N=4.096$, o GPyTorch reduz a diferença devido ao seu pipeline de autograd fundido amortizando a sobrecarga. Para GP esparsa, LightGP é mais rápido em $N=10.000$ (1,7×), mas GPyTorch vence em $N=50.000$ devido a tensores de dispositivo persistentes e computação de gradiente compilada.
• Operações Sem Matriz: No Apple M4, o produto kernel-vetor fundido Metal sem matriz alcança um aceleramento de 32× sobre o caminho explícito de CPU em $N=20.000$ e é a única opção viável além de $N=30.000$ em 16 GB de memória. No Linux, CUDA alcança até 259× de aceleração em $N=10.000$.
• Escalabilidade de Memória: CG sem matriz, VFE esparsa e SKI mantêm uso de memória $O(N)$. Em contraste, Cholesky exato ($O(N^2)$) esgota 16 GB de RAM em $N \approx 42.000$–45.000 no Mac e 12 GB de VRAM em $N \approx 35.000$ no Linux.
• Precisão: Previsões dos backends de CPU e GPU concordam dentro de $10^{-5}$ para média e variância em conjuntos de dados padrão (Motorcycle, Mauna Loa, kin40k).

Significado

LightGP demonstra que a inferência GP nativa em C++ pode igualar ou exceder o desempenho de bibliotecas baseadas em frameworks, exigindo zero dependências externas. Ao remover a sobrecarga de despacho de Python e frameworks de aprendizado profundo, LightGP oferece ganhos de desempenho consistentes, particularmente no Apple Silicon onde o coprocessor AMX é aproveitado diretamente. A biblioteca permite inferência GP em escalas ($N > 100.000$) anteriormente impraticáveis em hardware de consumidor através de métodos sem matriz de memória $O(N)$ e SKI acelerado por FFT.

O artigo destaca uma percepção crítica de hardware: no Apple Silicon, a CPU (via Accelerate/AMX) é frequentemente superior à GPU (Metal) para tarefas de álgebra linear densa como fatoração de Cholesky, uma nuance que as heurísticas do LightGP gerenciam automaticamente. Embora o GPyTorch permaneça competitivo ou superior em escalas muito grandes ($N \ge 50.000$) em CUDA devido à sua autograd madura e infraestrutura de tensor persistente, o LightGP preenche um nicho para inferência GP leve, de alto desempenho e implantável sem o fardo de uma pilha completa de aprendizado profundo.