Precision autotuning for linear solvers via contextual bandit-based RL

Resumo Técnico: Ajuste Automático de Precisão para Solucionadores Lineares via RL Baseado em Bandit Contextual

1. O Problema

A computação científica e o aprendizado de máquina modernos dependem intensivamente de aritmética de ponto flutuante. Com o avanço de hardware como GPUs (ex: NVIDIA A100 e H100), o uso de precisão reduzida (como FP16, BF16, FP8) oferece ganhos significativos de velocidade (até 64x) e eficiência energética em comparação com a precisão dupla (FP64).

No entanto, a redução de precisão introduz erros de arredondamento e instabilidade numérica, o que pode comprometer a convergência e a precisão de algoritmos iterativos, como solucionadores de sistemas lineares ($Ax = b$).

• Desafio Atual: Métodos de mixed-precision (precisão mista) tradicionais exigem análise teórica complexa ou ferramentas de ajuste (tuning) baseadas em tentativa e erro (trial-and-error). Essas ferramentas existentes frequentemente falham em generalizar para dados não vistos, exigem múltiplas execuções do programa para calibração e não escalam bem para grandes problemas.
• Objetivo: Desenvolver um método que aprenda automaticamente a selecionar a configuração de precisão ideal para cada etapa de um algoritmo numérico, equilibrando eficiência computacional e estabilidade numérica, sem depender de conhecimento prévio detalhado sobre o problema específico.

2. Metodologia

Os autores propõem um Framework de Aprendizado por Reforço (RL) formulado como um problema de Bandit Contextual. Diferente de RL tradicional que busca recompensas cumulativas de longo prazo, o Bandit Contextual foca na recompensa imediata baseada no contexto atual (características do problema), o que é ideal para instâncias independentes de problemas lineares.

Componentes Principais do Framework:

• Formulação do Problema:

  • Contexto (Estado $s$): Um vetor de características extraído da matriz $A$ e do sistema, discretizado para gerenciar a complexidade. As características principais são o número de condição ($\kappa(A)$) e a norma da matriz ($\|A\|_\infty$).
  • Ação ($a$): Uma configuração de precisão para as etapas computacionais do algoritmo. Para o caso de estudo (Refinamento Iterativo baseado em GMRES), as ações são tuplas de precisões para quatro etapas: Fatoração LU, Cálculo do Resíduo, Solução GMRES e Atualização da Solução.
  • Espaço de Ações Reduzido: Para evitar a explosão combinatória ($m^k$), o espaço de ações é restrito a combinações onde a precisão não diminui ao longo das etapas críticas (ordenação $u_1 \le u_2 \le \dots \le u_k$), reduzindo o espaço de 256 para 35 combinações válidas.
  • Recompensa ($r$): Uma função multi-objetivo que pondera:
    1. Precisão/Erro: Penaliza erros de forward e backward relative.
    2. Custo Computacional: Incentiva o uso de formatos de menor precisão (BF16, TF32).
    3. Penalidade de Iteração: Penaliza o número excessivo de iterações (convergência lenta) ou falhas (estagnação).

• Algoritmo de Aprendizado:

  • Utiliza uma Tabela Q (Q-Table) tabular para estimar o valor esperado das ações em estados discretizados.
  • Estratégia de seleção: $\epsilon$-greedy com decaimento linear, equilibrando exploração (testar novas configurações) e exploração (usar o melhor conhecimento atual).
  • Fase de Inferência: Após o treinamento em um conjunto limitado de dados, o modelo seleciona a melhor configuração de precisão para novos sistemas lineares (dados out-of-sample) usando a política $\pi(a|s) = \arg\max Q(s, a)$.

3. Estudo de Caso: Refinamento Iterativo GMRES (GMRES-IR)

O framework foi aplicado ao método GMRES-IR, um algoritmo bem estabelecido para resolver sistemas lineares. O objetivo foi determinar dinamicamente a precisão para:

1. Fatoração LU e solução inicial.
2. Cálculo do resíduo.
3. Resolução do sistema corrigido via GMRES.
4. Atualização da solução.

4. Resultados Experimentais

Os experimentos foram realizados em matrizes densas e esparsas com variados números de condição ($10^1$ a $10^9$), comparando o agente RL com uma linha de base de precisão dupla (FP64).

• Sistemas Densos (Bons Condicionados):

  • O agente RL aprendeu a utilizar formatos de baixa precisão (BF16, TF32) para sistemas com baixo número de condição, reduzindo o custo computacional.
  • Para sistemas mal condicionados, o agente transicionou automaticamente para FP64 para garantir a estabilidade.
  • Política Conservadora (W1): Mantém precisão comparável ao FP64, usando apenas FP32/FP64.
  • Política Agressiva (W2): Utiliza BF16/TF32 mais frequentemente, aceitando um leve aumento no erro ou no número de iterações internas do GMRES para ganhar eficiência, mantendo a convergência na maioria dos casos.

• Sistemas Esparsos (Mal Condicionados):

  • Em matrizes com números de condição muito altos ($10^8$ a $10^9$), o agente RL adotou uma estratégia altamente conservadora, utilizando quase exclusivamente FP64.
  • Isso demonstra que o modelo aprendeu a "fronteira de sobrevivência" numérica: quando a instabilidade é alta, qualquer redução de precisão leva à divergência, e o agente prioriza a convergência sobre a eficiência.

• Generalização:

  • O modelo treinado em um conjunto de dados demonstrou capacidade de generalização para dados não vistos, ajustando a precisão corretamente sem necessidade de re-treinamento.
  • A ablação da penalidade de iteração mostrou que, sem ela, o agente tende a escolher precisões muito baixas, resultando em um número excessivo de iterações para compensar a perda de precisão, invalidando o ganho de eficiência.

5. Contribuições Chave

1. Primeira Aplicação de RL para Ajuste de Precisão com Generalização Verificada: Diferente de ferramentas anteriores que exigem re-treinamento ou busca exaustiva, este trabalho demonstra que um modelo de RL pode aprender características salientes (como número de condição) e inferir configurações ótimas para dados não vistos.
2. Formulação como Bandit Contextual: Simplifica o problema de decisão, focando na recompensa imediata baseada no estado do sistema, o que é computacionalmente mais eficiente que RL de longo prazo para este tipo de problema.
3. Redução Eficiente do Espaço de Ações: A imposição de restrições de ordem nas precisões reduz drasticamente o espaço de busca, permitindo treinamento rápido e convergência estável.
4. Adaptabilidade Dinâmica: O sistema consegue alternar entre modos de alta eficiência (baixa precisão) e alta estabilidade (alta precisão) dependendo apenas das características numéricas da matriz de entrada.

6. Significado e Impacto

Este trabalho avança o estado da arte em computação de precisão mista e IA para métodos numéricos. Ele oferece uma solução prática para o dilema entre velocidade e precisão em solvers lineares, que são fundamentais em simulações científicas e aprendizado de máquina.

• Automação: Reduz a dependência de especialistas humanos para configurar parâmetros de precisão.
• Eficiência: Permite o uso seguro de hardware de baixa precisão (como FP8/FP16) em cenários onde antes era arriscado, potencialmente acelerando grandes cargas de trabalho científicas.
• Generalidade: O framework é projetado para ser estendido a outros algoritmos numéricos além do GMRES-IR, desde que características relevantes do problema possam ser extraídas como estado.

Em suma, o artigo demonstra que o Aprendizado por Reforço, quando bem formulado e limitado a características físicas do problema, pode superar métodos heurísticos tradicionais, oferecendo um ajuste de precisão "inteligente" e adaptativo.