Fixed-Point Neural Optimal Transport without Implicit Differentiation

Resumo Técnico: Transporte Ótimo Neural de Ponto Fixo sem Diferenciação Implícita

Declaração do Problema
O Transporte Ótimo (TO) busca um mapeamento entre distribuições de probabilidade que minimize um custo de transporte prescrito. Embora abordagens recentes de aprendizado profundo tenham habilitado solucionadores de TO escaláveis, elas frequentemente dependem de otimização adversarial min–max (por exemplo, WGANs) ou arquiteturas multi-rede envolvendo mapas de transporte e potenciais duais. Esses métodos frequentemente sofrem de instabilidade no treinamento, sensibilidade a hiperparâmetros e dificuldades de convergência, particularmente em configurações de alta dimensão. Abordagens alternativas baseadas em fluxos contínuos exigem a resolução de equações diferenciais ordinárias (EDOs), incorrendo em sobrecarga computacional significativa, enquanto formulações de Hamilton–Jacobi–Bellman (HJB) lutam com a escalabilidade em altas dimensões. Além disso, muitos métodos existentes exigem redes auxiliares ou procedimentos de otimização alternada para impor a viabilidade dual.

Metodologia
Os autores propõem um framework de rede única para transporte ótimo quadrático que elimina o treinamento adversarial e a diferenciação implícita. O núcleo do método é uma formulação neural implícita derivada do problema dual de Kantorovich.

1. Parametrização de Potencial Único: O potencial dual $\phi$ é parametrizado como o negativo de uma função convexa $g_\theta$ representada por uma única rede neural ($\phi_\theta(x) = -g_\theta(x)$). Isso aproveita o teorema de Brenier, que afirma que o mapa de transporte ótimo é o gradiente de um potencial convexo.
2. Reformulação Proximal de Ponto Fixo: Sob essa parametrização, a transformada $c$ (um componente chave do objetivo dual) reduz-se a um problema de otimização proximal fortemente convexo:
   $$(-g_\theta)^c(z) = \inf_{y \in \mathbb{R}^d} \left\{ \frac{1}{2}\|y - z\|^2 + g_\theta(y) \right\}$$
   O minimizador $y^*_\theta(z)$ deste problema define o mapa de transporte reverso ($T_{\nu \to \mu}(z) = y^*_\theta(z)$), enquanto o mapa direto é dado por $T_{\mu \to \nu}(x) = x + \nabla g_\theta(x)$.
3. Treinamento via Iteração de Ponto Fixo: O objetivo de treinamento envolve minimizar a perda dual negativa. Crucialmente, a minimização interna (o operador proximal) é resolvida usando iterações padrão de ponto fixo (por exemplo, etapas de descida de gradiente) em vez de otimização desenrolada ou diferenciação implícita.
4. Cálculo de Gradiente sem Diferenciação Implícita: Uma percepção teórica chave é que o gradiente da perda em relação aos parâmetros da rede $\theta$ não requer retropropagação através das iterações de ponto fixo. Devido às condições de otimalidade de primeira ordem do operador proximal, os termos envolvendo a derivada do ponto fixo em relação a $\theta$ ($dy^*/d\theta$) cancelam-se exatamente. Consequentemente, o gradiente depende apenas das derivadas explícitas de $g_\theta$ avaliadas nas amostras de entrada e nas soluções aproximadas de ponto fixo. Isso permite um treinamento eficiente usando retropropagação padrão.
5. Convergência: Os autores estabelecem que a Descida de Gradiente Estocástica (SGD) converge para um ponto estacionário mesmo quando o problema interno de ponto fixo é resolvido aproximadamente (dentro de uma tolerância especificada), desde que o erro de aproximação seja limitado.
6. Extensão Condicional por Classe: O framework estende-se naturalmente ao TO condicional por classe condicionando o potencial $g_\theta(x, k)$ aos rótulos de classe $k$, permitindo que uma única rede lide com múltiplas tarefas de transporte simultaneamente.

Contribuições Principais

• Reformulação de Rede Única: Os autores derivam uma formulação de TO quadrático que requer apenas uma rede neural para recuperar tanto os mapas de transporte direto quanto reverso, removendo a necessidade de configurações adversariais min–max ou redes conjugadas auxiliares.
• Treinamento sem Diferenciação Implícita: Eles demonstram que o objetivo pode ser diferenciado sem diferenciação implícita, confiando no cancelamento exato de termos induzido pelas condições de otimalidade do operador proximal.
• Garantias de Convergência: O artigo fornece provas teóricas para a convergência da SGD a um ponto estacionário sob avaliações de ponto fixo inexatas, abordando a realidade prática de que o loop interno não é resolvido com precisão de máquina.
• Framework Unificado: O método recupera simultaneamente os mapas direto e reverso e estende-se naturalmente a configurações condicionais por classe sem alterações arquitetônicas.

Resultados Experimentais
O método proposto foi avaliado em benchmarks gaussianos de alta dimensão, conjuntos de dados de física do mundo real (UCI) e tarefas de tradução de imagens (MNIST para Fashion-MNIST).

• Gaussianos de Alta Dimensão: Em benchmarks gaussianos com dimensões de até $d=64$, o método alcançou uma Porcentagem de Variância Não Explicada (UVP) significativamente menor em comparação com baselines como NOT, MM-v1, MM:R, NCF e WGAN-QC. Mantive precisão estável conforme a dimensionalidade aumentava, enquanto as baselines mostraram degradação.
• Eficiência Computacional: O método demonstrou eficiência computacional e de memória favorável, exigindo menos parâmetros e menos tempo de treinamento do que métodos baseados em fluxo contínuo (FFJORD, RNODE) e outras abordagens neurais de TO.
• Dados do Mundo Real: Em diversos conjuntos de dados UCI (Power, Gas, HEPMASS, MINIBOONE) e patches de imagem BSDS300, o método alcançou desempenho de correspondência de distribuição competitivo ou superior (medido pela Discrepância Média Máxima) com tamanhos de modelo substancialmente menores.
• Tradução de Imagem: Em tarefas de tradução de imagem condicional por classe, o método alcançou a menor Distância Fréchet Inception (FID) e a maior precisão de classificação entre as baselines, indicando preservação superior da estrutura de classe e conteúdo semântico.
• Estudos de Ablação: Experimentos mostraram que o método é robusto à tolerância de parada do solucionador de ponto fixo; uma inexatidão moderada no loop interno produz um compromisso favorável entre precisão e velocidade. Adicionalmente, embora impor convexidade estrita via projeção de pesos tenha estabilizado as iterações de ponto fixo, isso degradou a precisão geral de transporte, sugerindo que uma parametrização não convexa (sem projeção rígida) oferece um melhor equilíbrio entre expressividade e estabilidade.

Significado e Alegações
O artigo alega que este framework oferece uma alternativa estável, eficiente e escalável aos métodos existentes de TO neural. Ao eliminar o treinamento adversarial e a sobrecarga computacional de solucionadores de EDOs ou resolução de equações HJB sobre todo o domínio, o método fornece um mecanismo "livre de causalidade" para impor condições de otimalidade. Os autores afirmam que sua abordagem é amplamente aplicável, capaz de lidar com problemas de alta dimensão e condicionais por classe enquanto recupera mapas de transporte precisos em ambas as direções. O trabalho posiciona-se como um passo em direção à simplificação da paisagem de otimização do TO, afastando-se de formulações min–max complexas em direção a um problema de minimização de objetivo único que é teoricamente fundamentado e praticamente eficiente.