Resumo Técnico: cuRegOT

Formulação do Problema

O Transporte Ótimo (OT) é uma ferramenta fundamental no aprendizado de máquina moderno para comparar medidas de probabilidade, amplamente utilizada em adaptação de domínio, modelagem generativa e visão computacional. No entanto, resolver o problema de OT discreto em escala é computacionalmente proibitivo devido à complexidade supercúbica dos solucionadores padrão de programação linear. Embora a regularização entrópica (Cuturi, 2013) transforme o problema em um problema de otimização dual suave e sem restrições, solucionável via o algoritmo Sinkhorn, este último frequentemente sofre de convergência lenta em problemas desafiadores.

Avanços recentes introduziram métodos de segunda ordem e quase-Newton (por exemplo, SNS, SSNS e o algoritmo Sparse-Plus-Low-Rank ou SPLR) que oferecem taxas de convergência superiores explorando a estrutura da matriz Hessiana. Contudo, esses métodos enfrentam um gargalo crítico ao serem implantados em GPUs: a necessidade de resolver grandes sistemas lineares esparsos. Especificamente, as etapas de análise simbólica e reordenamento necessárias para a decomposição de Cholesky esparsa são inerentemente sequenciais, difíceis de paralelizar em arquiteturas SIMT e tipicamente executadas na CPU. Isso cria um gargalo de sincronização onde a GPU permanece ociosa, prejudicando severamente a eficiência de solucionadores quase-Newton, que de outra forma seriam poderosos.

Metodologia

O artigo apresenta o cuRegOT, um solucionador de alta performance para GPU adaptado para OT regularizado entrópicamente, construído sobre o algoritmo SPLR. Para superar o gargalo de sincronização CPU-GPU e as limitações de largura de banda de memória, os autores propõem três otimizações específicas algorítmicas e arquitetônicas:

1. Análise Simbólica Amortizada

O SPLR padrão reavalia o padrão de esparsidade da aproximação da Hessiana ($H_\Omega$) a cada iteração, acionando uma análise simbólica baseada em CPU dispendiosa. Os autores observam que o padrão de esparsidade (determinado pelas entradas top-$k$ do plano de transporte $T$) evolui lentamente ao longo das iterações.

• Estratégia: O solucionador realiza a análise simbólica baseada em CPU e o reordenamento da matriz apenas uma vez a cada $S$ iterações (por exemplo, $S=10$). Para as iterações intermediárias, o padrão de esparsidade estabelecido é reutilizado, atualizando apenas os valores numéricos dos elementos não nulos.
• Garantia Teórica: Os autores provam que reutilizar o padrão de esparsidade não expande a faixa de autovalores da aproximação da Hessiana. Consequentemente, o número de condição permanece limitado, garantindo que as propriedades de convergência do método quase-Newton sejam preservadas.

2. Computação Colaborativa CPU-GPU (Geração Assíncrona de Candidatos)

Mesmo com a amortização, a GPU permanece ociosa durante as fases periódicas de análise simbólica baseadas em CPU.

• Estratégia: Os autores introduzem um mecanismo assíncrono onde, enquanto a CPU realiza a análise simbólica, a GPU executa simultaneamente iterações leves do Sinkhorn sobre as variáveis duais atuais. Isso gera uma iteração "candidata" ($x^s_k$).
• Seleção: Uma vez que a CPU completa a análise e a GPU termina a etapa principal do passo quase-Newton (gerando $x^q_k$), o solucionador avalia a função objetivo e as normas do gradiente para ambos os candidatos. A iteração que produz maior melhoria é selecionada.
• Benefício: Essa abordagem mascara a latência da CPU e fornece exploração adicional do espaço de soluções sem aumentar o tempo de relógio. A análise teórica confirma que a inserção dessas etapas de seleção de candidatos preserva a convergência global do algoritmo original.

3. Kernel CUDA Fundido para Avaliação de Gradiente

O cálculo ingênuo do gradiente envolve múltiplas passagens sobre a matriz de custo $M$ e o plano de transporte $T$, levando a gargalos severos de largura de banda de memória em GPUs.

• Estratégia: Um kernel CUDA fundido é projetado para realizar computação elemento a elemento, redução de linha e redução de coluna em uma única passagem sobre os dados.
• Implementação: O kernel utiliza um loop de passo de grade e maximiza a reutilização de dados em registradores e memória compartilhada. Emprega intrínsecos de shuffle em nível de warp para reduções de linha (uso zero de memória compartilhada) e buffers de memória compartilhada para reduções de coluna, minimizando escritas na memória global e contenção atômica.

Contribuições Principais

1. Biblioteca cuRegOT: Um solucionador de alta performance para GPU para OT regularizado entrópicamente que adapta com sucesso o método quase-Newton SPLR ao hardware de GPU.
2. Inovações Algorítmicas:
   • Uma estratégia de análise simbólica amortizada que reduz a sobrecarga da CPU reutilizando padrões de esparsidade.
   • Um mecanismo de geração assíncrona de candidatos que utiliza o tempo ocioso da GPU durante a análise da CPU para acelerar a convergência.
   • Um kernel de gradiente fundido que reduz significativamente a sobrecarga de I/O de memória.
3. Garantias Teóricas: Provas rigorosas demonstrando que as modificações propostas (reutilização de padrão e seleção de candidatos) mantêm a convergência global e a taxa de convergência linear do algoritmo SPLR original.

Resultados Experimentais

Os autores avaliaram o cuRegOT contra solucionadores baseados em GPU de última geração, incluindo POT (Sinkhorn, Greenkhorn), OTT-JAX (Sinkhorn, aceleração de Anderson) e AccSinkhorn, usando dados sintéticos e o conjunto de dados CIFAR-10.

• Desempenho: O cuRegOT consistentemente alcança convergência significativamente mais rápida do que todas as linhas de base, particularmente quando alta precisão é exigida.
• Escalabilidade: A vantagem de desempenho do cuRegOT aumenta com o tamanho do problema (testado até $6400 \times 4800$ e $5000 \times 5000$), validando sua escalabilidade para aplicações em grande escala.
• Estudo de Ablação: Experimentos removendo individualmente a análise amortizada ("-A") ou a geração de candidatos ("-S") mostraram convergência mais lenta, confirmando que ambos os componentes são essenciais para o desempenho do solucionador.
• Linhas de Base: Embora a aceleração de Anderson tenha proporcionado acelerações para o Sinkhorn padrão em alguns casos, ela ainda enfrentou convergência lenta para alta precisão em comparação com o cuRegOT. O Greenkhorn foi considerado muito lento para comparação prática.

Significância

O artigo aborda uma lacuna crítica na implantação de métodos de otimização avançados para Transporte Ótimo em hardware moderno. Ao preencher a lacuna de eficiência entre a velocidade de convergência teórica dos métodos quase-Newton e as limitações práticas da execução em GPU (especificamente no que diz respeito a solucionadores lineares esparsos), o cuRegOT permite a solução eficiente de problemas de OT em grande escala e de alta precisão. O trabalho demonstra que, com um design algorítmico cuidadoso — especificamente amortizando tarefas sequenciais da CPU e fundindo operações limitadas por memória —, métodos de segunda ordem podem ser tornados viáveis e superiores às linhas de base de primeira ordem em aceleradores de GPU.