TopoGPU acelera o cálculo de homologia persistente em média 53 vezes

A homologia persistente é cada vez mais usada em deep learning, mas calculá-la dentro do loop de treinamento é muito lento. Os autores de um novo artigo no arXiv propõem o TopoGPU, um pipeline GPU em streaming que constrói diagramas de persistência para complexos cúbicos gerados por imagens 2D e 3D. O concorrente, Cubical Ripser, é considerado uma das melhores implementações existentes. O TopoGPU superou o rival em todos os benchmarks testados: 53,24 vezes mais rápido em média no tempo total de cálculo (end-to-end) e até 198,01 vezes no pico. O código é aberto, e há binários prontos, notebooks para Google Colab e imagens Docker.

Onde estava o gargalo

A homologia persistente acompanha como componentes conexos, laços e cavidades surgem e desaparecem conforme o limiar de brilho da imagem muda. O resultado é um diagrama de persistência, uma descrição compacta da forma. Em tarefas como segmentação, ele entra como parte da função de perda: a rede é penalizada por vasos rompidos ou buracos a mais na máscara. Mas, nesse caso, o diagrama precisa ser recalculado a cada passo do treinamento. O Cubical Ripser roda na CPU, e, em volumes 3D grandes, o cálculo consome uma fatia considerável do tempo de treinamento.

Como funciona o pipeline

O TopoGPU divide a imagem de entrada em blocos e os envia à GPU um de cada vez. Cada bloco é processado por kernels massivamente paralelos, em uma grade de blocos de GPU. As relações de borda encontradas são acumuladas na memória do host, e é ali que a CPU executa a redução da matriz de bordas. A parte mais paralelizável do trabalho vai para a placa de vídeo, enquanto a etapa final, sequencial, fica no processador. Pela descrição, o esquema em streaming não exige manter o volume inteiro na GPU ao mesmo tempo, embora o resumo do artigo não informe os limites exatos de memória.

A principal ideia técnica envolve a teoria de Morse discreta. Ela permite "colapsar" o complexo, removendo pares de células sem alterar a homologia, de modo que a matriz de bordas fica bem menor. O problema é que, no processamento em streaming, cada bloco enxerga só parte dos dados, e um emparelhamento de Morse ingênuo pode distorcer o resultado. Os autores introduzem o stratification-aware discrete Morse matching e demonstram que ele preserva a homologia persistente no esquema em streaming. Para rodar na GPU, foram adicionados uma ordenação topológica paralela e um algoritmo paralelo de paridade de V-paths, que deriva as bordas de Morse.

• Aceleração média end-to-end contra o Cubical Ripser: 53,24x
• Aceleração máxima: 198,01x
• Dados de entrada: imagens 2D e 3D, complexos cúbicos
• Divisão de trabalho: a GPU constrói as relações de borda, a CPU reduz a matriz
• Licença e distribuição: código aberto, binários, Colab, Docker

Para quem isso serve

Os autores testaram o método em uma rede com consciência topológica e mostraram que o custo de calcular a homologia persistente durante o treinamento cai de forma notável. Essa é a principal conclusão prática. Se o seu modelo usa perdas topológicas, por exemplo para segmentar estruturas finas em exames 3D na área médica, o TopoGPU pode eliminar o gargalo sem mudar a abordagem de treinamento. Com as imagens Docker e os notebooks Colab prontos, dá para testá-lo em uma noite.

Há ressalvas. A comparação é feita com o Cubical Ripser, e os números de 53x e 198x dependem do conjunto de benchmarks e do hardware, que o resumo do artigo não descreve. A redução da matriz continua na CPU, então, em dados complexos, ela pode virar o novo limitante. Antes de mexer no seu pipeline, vale rodar o TopoGPU em suas próprias imagens e comparar o tempo de cálculo nos tamanhos de volume que você usa, em vez de se basear apenas no ganho médio do artigo.