Resumo Técnico: FlashSinkhorn: Transporte Ótimo Entópico Consciente de IO em GPU

Declaração do Problema

O Transporte Ótimo Entópico (EOT), resolvido via iterações de Sinkhorn, é uma ferramenta fundamental no aprendizado de máquina moderno para comparar medidas de probabilidade, variando desde modelagem generativa até alinhamento distribucional. No entanto, escalar o EOT para grandes nuvens de pontos (por exemplo, $n \approx 70.000$) em GPUs permanece um gargalo persistente.

Os solucionadores de GPU existentes enfrentam duas ineficiências primárias:

1. Implementações Tensorizadas: Estas materializam a matriz de custos densa $n \times m$, levando a um tráfego quadrático de Memória de Alta Largura de Banda (HBM) ($O(nm)$ leituras/escritas) e problemas de capacidade de memória que causam erros de Esgotamento de Memória (OOM) à medida que $n$ cresce.
2. Backends Online: Embora estes evitem armazenar a matriz densa ao avaliar interações sob demanda (por exemplo, via KeOps ou OTT-JAX), eles dependem de kernels genéricos de redução map-reduce em blocos. Esses kernels sofrem de fusão limitada, impedindo a exploração de estruturas GEMM de tensor-core e resultando em utilização de computação subótima e tempo de execução elevado, apesar da redução no uso de memória.

O gargalo central dos sistemas é que os métodos atuais são ou limitados por memória (devido ao tráfego excessivo de HBM) ou limitados por computação com baixa eficiência de instruções, falhando em escalar para as resoluções repetidas em grande escala exigidas em pipelines downstream.

Metodologia

Os autores introduzem o FlashSinkhorn, um solucionador de GPU consciente de IO para EOT com custo euclidiano quadrado. A metodologia traça um paralelo direto entre as atualizações de Sinkhorn no domínio log estabilizado e o mecanismo de atenção em Transformers, aproveitando os princípios do FlashAttention para otimizar padrões de acesso à memória.

1. Reformulação na Forma de Atenção

A principal insight teórica (Proposição 1) é que, para custo euclidiano quadrado, as atualizações de Sinkhorn no domínio log estabilizado podem ser reescritas como reduções LogSumExp (LSE) linha a linha de uma matriz de pontuação de produto escalar enviesada.

• Especificamente, a atualização $f_i \leftarrow -\varepsilon \text{LSE}_j [(g_j - C_{ij})/\varepsilon + \log b_j]$ é equivalente a calcular o LSE de uma matriz $S_X = (QK^\top + \mathbf{1}(\hat{g} + \delta))/\varepsilon$, onde $Q=\sqrt{2}X$, $K=\sqrt{2}Y$, e $\delta, \gamma$ são vieses de peso log pré-calculados.
• Esta reformulação transforma a iteração de Sinkhorn na mesma estatística de normalização exata usada na atenção de produto escalar escalonada, permitindo a aplicação de técnicas de streaming conscientes de IO.

2. Kernels de Streaming e Tiling

O FlashSinkhorn implementa esta reformulação usando kernels Triton fundidos que empregam um modelo de memória de dois níveis (HBM lenta e SRAM on-chip rápida):

• Sem Materialização: A matriz de custos $n \times m$ nunca é materializada na HBM. Em vez disso, blocos (tiles) de $Q$ e $K$ são transmitidos em fluxo.
• LSE Online: Os kernels mantêm estatísticas de soma de exp e máximo em execução linha a linha on-chip (SRAM) enquanto processam blocos da matriz de pontuação. Isso permite o cálculo de reduções LSE exatas com significativamente menos acessos à HBM.
• Fusão: Toda a atualização (cálculo de pontuação, adição de viés, redução LSE e atualização potencial) é fundida em uma única passagem de kernel. Isso elimina escritas intermediárias na HBM e reduz a sobrecarga de lançamento de kernel.
• Traversão Estacionária por Linha: Ao contrário do FlashAttention, que frequentemente itera sobre chaves/valores no loop externo, o FlashSinkhorn adota uma travessia estacionária por linha (fixando um bloco de linha de consulta na SRAM e transmitindo todos os blocos de chave) para melhor corresponder à natureza de redução linha a linha das atualizações de Sinkhorn.

3. Diferenciação e Otimização de Segunda Ordem

O artigo estende esta estratégia consciente de IO para diferenciação e métodos de segunda ordem:

• Aplicação da Matriz de Transporte: A projeção baricêntrica e o cálculo do gradiente são expressos como operações semelhantes à atenção (médias ponderadas), permitindo que sejam computados via os mesmos kernels de streaming sem materializar a matriz de transporte $P^*$.
• Produtos Hessiano-Vetor (HVP): Os autores derivam um oráculo de HVP em streaming que decompõe o produto Hessiano-vetor em uma sequência de produtos transporte-vetor e transporte-matriz. Isso evita a materialização do tensor Hessiano $O(n^2 d^2)$ ou da matriz de transporte $O(nm)$, reduzindo a memória de trabalho para $O((n+m)d)$.

Contribuições Principais

1. Atualização de Sinkhorn na Forma de Atenção: Os autores expressam as atualizações de Sinkhorn estabilizadas como reduções LSE de produto escalar enviesado, permitindo streaming exato no estilo FlashAttention para EOT.
2. Solucionador EOT Consciente de IO: Eles fornecem kernels Triton fundidos para iterações de Sinkhorn em fluxo e operadores de streaming sem matriz de transporte para diferenciação, reduzindo substancialmente o volume de IO da HBM.
3. Otimização de Segunda Ordem Escalável: Eles introduzem HVPs em streaming que habilitam métodos Newton-CG e de região de confiança escaláveis para perdas baseadas em EOT, anteriormente inviáveis em grande escala devido a restrições de memória.
4. Avaliação em Grande Escala: O artigo demonstra acelerações substanciais e eficiência de memória em GPUs A100 em benchmarks sintéticos e tarefas downstream como Distância de Conjunto de Dados de Transporte Ótimo (OTDD) e regressão embaralhada.

Resultados Experimentais

Avaliado em GPUs NVIDIA A100-80GB, o FlashSinkhorn demonstra ganhos de desempenho significativos sobre baselines de última geração (GeomLoss com backends KeOps/Tensorizados e OTT-JAX):

• Velocidade:
  • Passada Forward: Alcança até 32× de aceleração sobre o KeOps e 5,1× sobre o OTT-JAX.
  • End-to-End (Forward + Backward): Alcança até 161× de aceleração sobre o KeOps em altas dimensões ($d=512$).
  • HVP: Alcança 3–25× de aceleração dependendo da dimensionalidade, tornando o Newton-CG prático em escalas onde outros métodos esgotam o tempo.
• Eficiência de Memória:
  • O FlashSinkhorn mantém complexidade de memória linear $O((n+m)d)$, enquanto as bases tensorizadas escalam como $O(nm)$ ou $O(n^2)$.
  • Ele processa com sucesso nuvens de pontos com $n=50.000$ e dimensões até $d=1024$, onde métodos tensorizados esgotam a memória e bases online (KeOps) excedem os limites de tempo.
• Tarefas Downstream:
  • OTDD: O FlashSinkhorn escala o cálculo da distância OTDD e o fluxo de gradiente para $n=60.000$ (MNIST $\leftrightarrow$ Fashion-MNIST) com pico de memória abaixo de 1 GB, enquanto as bases tensorizadas sofrem OOM em $n > 20.000$.
  • Escape de Sela: Na regressão linear embaralhada, o HVP em streaming permite monitorar o menor autovalor do Hessiano para detectar pontos de sela, permitindo que um otimizador híbrido Adam-Newton escape de selas e convirja rapidamente.

Significado e Alegações

O artigo afirma que o FlashSinkhorn torna o EOT em grande escala prático no regime limitado por memória. Ao reescrever as atualizações de Sinkhorn para corresponder à estrutura consciente de IO do FlashAttention, os autores fecham a lacuna entre o transporte ótimo e as técnicas modernas de aceleração de transformers.

O significado reside em:

• Otimização em Nível de Sistema: Os ganhos são atribuídos à especialização em nível de kernel (fusão e tiling) e não a mudanças algorítmicas na própria iteração de Sinkhorn.
• Habilitação de Novas Cargas de Trabalho: A redução na pegada de memória e no tempo de execução permite o uso de EOT em tarefas downstream (como distância de conjunto de dados e diferenciação implícita) em escalas anteriormente consideradas computacionalmente proibitivas.
• Reprodutibilidade: Os autores lançam uma implementação de código aberto para facilitar pesquisas futuras em transporte ótimo escalável.

Os autores notam modestamente que esta abordagem atualmente se aplica a custos com uma estrutura específica (euclidiano quadrado e distância cosseno após normalização). Custos sem essa estrutura (por exemplo, euclidiano bruto com raízes quadradas ou custos neurais aprendidos) são identificados como trabalho futuro.