Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization

Resumo Técnico: Diferença de Programação Convexa no Espaço de Wasserstein com Aplicações à Otimização de MMD

Enunciado do Problema
A otimização de funcionais sobre o espaço de medidas de probabilidade é central para tarefas de aprendizado de máquina que variam desde a inferência variacional até a modelagem generativa. Uma abordagem padrão envolve realizar a otimização diretamente sobre o espaço de Wasserstein $(P_2(\mathbb{R}^d), W_2)$ através da discretização de fluxos de gradiente de Wasserstein. Embora eficaz para funcionais convexos (ex: divergência KL, divergências f), muitos objetivos praticamente relevantes são não-convexos ao longo de geodésicas de Wasserstein. Exemplos notáveis incluem a distância de Wasserstein ao quadrado, a distância de Sliced-Wasserstein e a Discrepância de Média Máxima (MMD). Métodos de primeira ordem padrão, como o Gradiente de Wasserstein (WGD), têm dificuldade nesses cenários não-convexos, muitas vezes exigindo injeção delicada de ruído ou otimização em espaços alternativos (ex: espaços Wasserstein-MMD ou Wasserstein-Fisher-Rao) que alteram os pesos das distribuições ou misturam passos implícitos.

Metodologia
Os autores propõem um novo algoritmo baseado em partículas adaptado para funcionais não-convexos que admitem uma decomposição de Diferença de Convexos (DC) dentro do espaço de Wasserstein. A metodologia central envolve elevar o clássico Procedimento Convexo-Côncavo (CCCP) de espaços euclidianos para o espaço de Wasserstein, resultando no Procedimento Convexo-Côncavo de Wasserstein (WCCCP).

1. Decomposição DC no Espaço de Wasserstein: O funcional objetivo $F$ é decomposto como $F(\mu) = F_+(\mu) - F_-(\mu)$, onde tanto $F_+$ quanto $F_-$ são funcionais totalmente convexos em $P_2(\mathbb{R}^d)$. Ao contrário de trabalhos anteriores que restringiam a parte côncava a potenciais de energia (funcionais lineares), esta abordagem permite que $F_-$ seja qualquer funcional diferenciável.
2. Esquema Algorítmico: Em cada iteração $k$, dada uma medida $\mu_k$, o algoritmo resolve um problema de majorização-minimização. Ele lineariza a parte côncava $F_-$ em torno de $\mu_k$ para construir um limite superior e minimiza o resultante sobregiante convexo sobre o espaço de mapas de transporte $T \in L^2(\mu_k)$:
   $$T_{k+1} = \arg\min_{T \in L^2(\mu_k)} \left( F_+(T\#\mu_k) - \langle \nabla_{W_2}F_-(\mu_k), T - \text{Id} \rangle_{L^2(\mu_k)} \right)$$
   O próximo iterado é então definido como $\mu_{k+1} = (T_{k+1})\#\mu_k$.
3. Implementação: A minimização interna é realizada via gradiente descendente em $L^2(\mu_k)$. Para distribuições empíricas (partículas), isso se traduz na atualização das posições das partículas. Os autores também estabelecem conexões entre WCCCP e Descida de Espelho/Descida Proximal de Bregman no espaço de Wasserstein.
4. Especificidades de MMD: Para a Discrepância de Média Máxima, os autores propõem decomposições DC explícitas ao dividir o próprio kernel $k(x,y) = \psi(x-y)$ em componentes convexos $\psi = \psi_+ - \psi_-$. Eles fornecem decomposições específicas para kernels invariantes por translação, incluindo kernels Gaussianos e de Riesz, usando decomposição de Jordan ou divisão algébrica (ex: $\cosh/\sinh$).

Principais Contribuições

• Algoritmo WCCCP: Introdução do Procedimento Convexo-Côncavo de Wasserstein, estendendo a otimização DC para o espaço de Wasserstein sem restringir o componente côncavo a potenciais lineares.
• Convergência Teórica:
  • Prova de que o valor do objetivo $F(\mu_k)$ é não-crescente ao longo dos iterados.
  • Estabelecimento de resultados de quase estacionalidade. Sob suposições de suavidade e convexidade forte (convexidade total) dos componentes, os autores provam taxas de convergência sublineares para a distância ao quadrado entre iterados e a norma do gradiente de Wasserstein. Especificamente, $\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$.
  • Demonstração de que o esquema é equivalente à Descida de Bregman Proximal e Descida de Espelho sob condições específicas.
• Decomposições DC Explícitas para MMD: O artigo fornece uma maneira sistemática de decompor funcionais MMD com kernels invariantes por translação em formas DC, satisfazendo as necessárias condições de convexidade e suavidade para garantias de convergência.
• Validação Empírica: Experimentos com Distância de Energia e MMD com kernels Gaussianos demonstram que decomposições DC bem escolhidas produzem convergência significativamente mais rápida e estável do que o WGD padrão e o método de Gradiente Proximal de Wasserstein (Forward-Backward).

Resultados

• Distância de Energia: Em dados sintéticos (formas de espiral e gato) e imagens CIFAR10, o WCCCP converge significativamente mais rápido que o WGD e o esquema Forward-Backward (FB). O algoritmo minimiza com sucesso a Distância de Energia para corresponder a distribuições alvo complexas.
• MMD com Kernel Gaussiano: Minimizar o MMD com um kernel Gaussiano é notoriamente sensível à largura de banda (bandwidth) e inicialização. Os autores mostram que o WGD e o esquema FB (usando uma decomposição padrão) frequentemente ficam presos em mínimos locais onde as partículas derivam para longe do alvo. Em contraste, o WCCCP usando decomposções de kernel específicas (Jordan ou $\cosh/\sinh$) alcança uma melhor convergência. A decomposição de Jordan é notada por ser mais suave e menos propensa a instabilidade numérica, enquanto a decomposição $\cosh/\sinh$ garante que todas as partículas eventualmente converjam, embora a uma taxa mais lenta.
• Comparação: Os resultados destacam que o desempenho do WCCCP é altamente dependente da escolha da decomposição DC. Uma decomposição "bem escolhida" é crítica para superar a descida de gradiente padrão.

Significância e Alegações
O artigo afirma fornecer um framework robusto para otimizar funcionais não-convexos no espaço de Wasserstein, abordando uma lacuna onde métodos de gradiente padrão falham devido à não-convexidade. Ao elevar o CCCP para este cenário, os autores oferecem uma alternativa teoricamente fundamentada que evita os problemas de mistura de pesos de abordagens anteriores específicas de MMD. O trabalho enfatiza que, embora o algoritmo seja geral, seu sucesso prático depende fortemente da decomposição DC específica do kernel. Os autores observam modestamente que trabalhos futuros são necessários para entender mais profundamente o impacto de diferentes decomposições e para projetar estratégias automáticas e adaptativas para selecioná-las. Adicionalmente, a análise teórica atualmente assume soluções exatas para os problemas de otimização interna, e a contabilização de solvers inexatos é identificada como uma direção necessária para pesquisas futuras.