Parallelised Differentiable Straightest Geodesics for 3D Meshes

1. O Problema

O aprendizado de máquina em domínios não euclidianos (como superfícies 3D representadas por malhas) tem avançado, mas métodos geometricamente precisos ainda estão atrasados. Os principais obstáculos são:

• Falta de operadores riemannianos de forma fechada: Muitas superfícies complexas não possuem soluções analíticas para operadores como o mapa exponencial.
• Não diferenciabilidade: As implementações discretas existentes de geodésicas (caminhos mais curtos/retos em superfícies) não são diferenciáveis, impedindo o uso de retropropagação (backpropagation) em redes neurais e otimizadores.
• Ineficiência computacional: Os métodos atuais são sequenciais (executados na CPU) e não paralelizáveis, tornando-os inviáveis para grandes lotes de dados ou aplicações em tempo real.
• Limitações de métodos existentes: Soluções baseadas em Eikonal ou integração numérica de ODEs são computacionalmente caras, muitas vezes não diferenciam corretamente ou produzem resultados fora da malha.

2. Metodologia

Os autores propõem um framework completo para computar o Mapa Exponencial (Exponential Map) em malhas triangulares 3D de forma diferenciável e paralelizada.

• Geodésicas Mais Retas (Straightest Geodesics): Em vez de resolver problemas de valor de contorno (BVP) para encontrar o caminho mais curto, o método utiliza geodésicas "mais retas", que preservam o ângulo entre segmentos da curva ao cruzar arestas e vértices da malha. Isso é feito iterativamente, mantendo a condição de que os ângulos de entrada e saída ($\theta_l = \theta_r$) sejam iguais em relação ao ângulo total do vértice.
• Paralelização em GPU: O algoritmo de traçado de geodésicas foi reescrito para rodar em GPUs usando CUDA. Cada thread é responsável por traçar um caminho geodésico completo independentemente, permitindo o processamento de dezenas de milhares de pontos simultaneamente em múltiplas malhas.
• Esquemas de Diferenciação: Como o traçado de geodésicas envolve escolhas combinatórias discretas (não diferenciáveis), os autores propõem dois métodos para calcular os gradientes:
  1. Proxy Externo (EP - Extrinsic Proxy): Utiliza uma função proxy analítica baseada em rotações e translações no espaço euclidiano para aproximar a derivada em relação ao vetor inicial ($v$). É extremamente rápido, mas menos preciso para variações na posição inicial ($p$).
  2. Diferenças Finitas Geodésicas (GFD - Geodesic Finite Differences): Um esquema baseado em diferenças finitas que perturba a posição inicial e o vetor, traçando novas geodésicas para estimar o Jacobiano. É mais lento que o EP, mas oferece maior precisão para ambas as variáveis ($p$ e $v$).
• Transporte Paralelo: O método também calcula o transporte paralelo de vetores ao longo da geodésica, essencial para manter a consistência dos campos vetoriais na superfície.

3. Principais Contribuições

1. Primeiro Mapa Exponencial Diferenciável e Paralelizado: Uma implementação eficiente em GPU que permite o cálculo de mapas exponenciais, traçado de geodésicas e transporte paralelo em malhas 3D.
2. Biblioteca digeo: Uma biblioteca PyTorch compatível (com kernels C++/CUDA) que integra geometria riemanniana discreta em pipelines de aprendizado modernos.
3. Convoluções Geodésicas Adaptativas (AGC): Uma nova camada de convolução que aprende dinamicamente o tamanho do "patch" (receptivo) durante o treinamento, adaptando-se à geometria local da superfície, superando métodos de patch fixo.
4. MeshFlow: Um método de Flow Matching (correspondência de fluxo) para malhas arbitrárias. Utiliza o mapa exponencial diferenciável para mover amostras de uma distribuição base para uma alvo sem necessidade de projetar pontos de volta para a superfície a cada passo (evitando erros de projeção).
5. Mesh-LBFGS: Um otimizador de segunda ordem (baseado em LBFGS) adaptado para variedades riemannianas em malhas, aplicado com sucesso ao problema de Tesselação de Voronoi Centroidal Geodésica (GCVT).

4. Resultados

• Desempenho e Precisão: A implementação em GPU é 3 ordens de magnitude mais rápida que as melhores implementações existentes em CPU (como pp3d e geometry-central) para grandes lotes, mantendo precisão comparável.
• Validação de Diferenciação: Os esquemas EP e GFD foram validados contra mapas exponenciais de forma fechada (esfera) e soluções de ODEs (toro). O GFD mostrou alta correlação cosinual com os gradientes reais, enquanto o EP foi eficaz para gradientes de vetores.
• Aplicações:
  • Segmentação Corporal: O modelo AGC alcançou 92.3% de precisão na segmentação de partes do corpo humano, superando métodos anteriores como GCNN e MDGCNN.
  • Geração de Modelos (MeshFlow): O MeshFlow foi 16.000 vezes mais rápido na inferência e usou 97% menos memória de GPU que o método de Riemannian Flow Matching (RFM) anterior, com métricas de divergência KL e distância de Chamfer superiores.
  • Otimização (Mesh-LBFGS): Convergiu mais rápido e alcançou soluções finais melhores para a tesselação de Voronoi em comparação com o algoritmo clássico de Lloyd, especialmente quando as sementes iniciais estavam agrupadas.

5. Significado e Impacto

Este trabalho representa um marco fundamental para a geometria computacional e o aprendizado de máquina em superfícies. Ao resolver o problema da não diferenciabilidade e da ineficiência dos operadores riemannianos em malhas discretas, os autores permitem que:

• Redes neurais aprendam diretamente sobre a geometria intrínseca de superfícies complexas.
• Otimizadores de segunda ordem sejam aplicados a problemas geométricos complexos.
• Métodos generativos modernos (como Flow Matching) sejam aplicados a dados 3D de forma eficiente e precisa.

A liberação da biblioteca digeo e dos modelos elimina a barreira de entrada para pesquisadores que desejam integrar geometria diferencial rigorosa em seus pipelines de deep learning, acelerando o avanço em áreas como visão computacional, gráficos computacionais e bioinformática.