Particle-based Generalised Stochastic Optimisation

Resumo Técnico: Otimização Estocástica Generalizada Baseada em Partículas

Enunciado do Problema
O artigo aborda o desafio de otimizar funções de perda $\ell: \Theta \to \mathbb{R}$ onde o gradiente é uma expectativa intratável em relação a uma distribuição dependente de um parâmetro $\pi_\theta$. Especificamente, o gradiente assume a forma:
$$\nabla_\theta \ell(\theta) = \mathbb{E}_{\pi_\theta}[F(\theta, X)] = \int_X F(\theta, x)\pi_\theta(dx)$$
Esta estrutura é prevalente no treinamento de Modelos Baseados em Energia (EBMs), modelos de variáveis latentes, Estimativa de Máxima Verossimilhança Marginal (MMLE) e no ajuste fino (fine-tuning) de modelos generativos. A distribuição $\pi_\theta$ tipicamente possui uma densidade não normalizada $e^{-E(\theta, x)}/C_\theta$, onde a constante de normalização $C_\theta$ é intratável.

Abordagens padrão frequentemente empregam esquemas de otimização-amostragem aninhados, onde um loop interno de Markov Chain Monte Carlo (MCMC) aproxima amostras de $\pi_\theta$ para cada atualização de gradiente externa. Esses métodos são computacionalmente caros e sofrem com problemas de convergência relacionados ao viés de MCMC, tempos de mistura (mixing times) e sensibilidade ao tamanho do passo. Embora a literatura recente tenha introduzido métodos de partículas interagentes (por exemplo, Particle Gradient Descent, Momentum Particle Gradient Descent) para evitar loops aninhados, estes são frequentemente adaptados a otimizadores ou processos de difusão específicos, carecendo de um arcabouço teórico unificado que acomode otimizadores modernos e sofisticados (como Adam ou métodos baseados em momentum) e dinâmicas de amostragem de ordem superior.

Metodologia
Os autores propõem um arcabouço geral e contínuo baseado em dinâmica de campo médio (mean-field dynamics) e suas aproximações de partículas interagentes. O núcleo da metodologia é um sistema acoplado de Equações Diferenciais Estocásticas (SDEs) envolvendo um estado do otimizador $\vartheta_t$ e um estado do amostrador $Z_t$:

1. Dinâmica do Otimizador:
   $$d\vartheta_t = \Psi_t(\vartheta_t, G_{\nu_t}(\vartheta_t))dt$$
   Aqui, $\vartheta_t \in \Theta \times M$ inclui parâmetros e variáveis auxiliares (ex: momentum). $\Psi_t$ é uma função de atualização geral que permite vários otimizadores (gradiente descendente, heavy-ball, tipos de Adam). $G_{\nu_t}$ é a estimativa de gradiente empírica derivada da distribuição de partículas $\nu_t$.

2. Dinâmica do Amostrador:
   $$dZ_t = b_{\Pi_\theta(\vartheta_t)}(Z_t)dt + \sigma_{\Pi_\theta(\vartheta_t)}(Z_t)dB_t$$
   $Z_t \in X \times V$ representa o estado das partículas. O drift $b$ e a difusão $\sigma$ são projetados de tal forma que, para um parâmetro congelado $\theta$, a distribuição marginal do componente $X$ converge para $\pi_\theta$. O arcabouço permite dinâmicas de Langevin superamortadas (overdamped), subamortadas (underdamped) e de ordem superior ao escolher apropriadamente $b$ e $\sigma$ (por exemplo, via uma construção de Langevin generalizada envolvendo matrizes simétricas e antissimétricas).

3. Aproximação por Partículas:
   O sistema de campo médio é aproximado por um sistema de $N$ partículas interagentes $\{Z^i_t\}_{i=1}^N$ impulsionadas por movimentos Brownianos independentes. As partículas interagem através do estado do otimizador compartilhado $\vartheta^N_t$, que atualiza com base na média empírica dos gradientes computados a partir das partículas.

Principais Contribuições

1. Arcabouço Unificado: O artigo estabelece um arcabouço de difusão geral que unifica métodos de partículas existentes (como PGD, MPGD e Implicit Diffusion) como casos especiais. Crucialmente, permite a integração de quaisquer otimizadores modernos e dinâmicas de amostragem complexas dentro de uma única estrutura teórica.
2. Garantias de Convergência Teórica:
   • Contração Exponencial: Sob suposições de bem-postura e uma condição de monotonia forte conjunta (uma condição de contratividade torcida no sistema acoplado otimizador-amostrador), os autores provam que o sistema de campo médio converge exponencialmente para uma distribuição estacionária.
   • Otimalidade Estacionária: Sob estruturas adicionais de conservatividade e convexidade forte conjunta, mostra-se que o parâmetro estacionário é o minimizador global único da perda.
   • Limites de Erro de Partícula Finita: O artigo deriva limites de erro não assintóticos para a aproximação de partículas. Especificamente, demonstra um erro de interação de $O(N^{-1/2})$ e fornece um limite para o erro do parâmetro que decai exponencialmente no tempo, desde que o número de partículas $N$ seja suficientemente grande.
3. Design de Novos Algoritmos: O arcabouço é instanciado para criar dois novos algoritmos:
   • HOL-MYPGD: Um algoritmo de dinâmica de Langevin de ordem superior para MMLE, utilizando técnicas proximais para lidar com priors não diferenciáveis (ex: Variação Total).
   • M-MFPO: Um otimizador de partículas baseado em momentum para treinamento de EBMs, generalizando o MPGD para uma classe mais ampla de problemas com dinâmicas de amostragem subamortadas.

Resultos
Os autores avaliam os métodos propostos em duas tarefas distintas:

1. Desfoque de Imagem Bayesiano (MMLE): Em uma configuração com priors de Variação Total anisotrópica, o proposto HOL-MYPGD (Langevin de ordem superior) é comparado contra um baseline de Moreau-Yosida Particle Gradient Descent (MYPGD). Os resultados indicam que o esquema de ordem superior alcança menor Erro Quadrático Médio (MSE) e maior Similaridade Estrutural (SSIM), produzindo reconstruções de imagem mais nítidas.
2. Treinamento de Modelo Baseado em Energia: O algoritmo M-MFPO é testado em conjuntos de dados sintéticos 2D (anéis, contas, rede) contra o Persistent Contrastive Divergence (PCD). O método baseado em momentum demonstra convergência mais rápida em termos de Divergência de Fisher em uma escala logarítmica comparado ao PCD, mantendo qualidade visual comparável nos exemplos gerados.

Significância e Alegações
O artigo alega fornecer uma abordagem "principiada e de propósito geral" para otimização estocástica de funções de perda com gradientes intratáveis. Sua principal significância reside em:

• Generalização: Ir além de construções específicas de problemas ou otimizadores para um arcabouço flexível que pode incorporar otimizadores de ponta (como Adam ou heavy-ball) e dinâmicas de amostragem avançadas (ordem superior, não reversíveis).
• Rigor Teórico: Fornecer provas de convergência rigorosas para esta classe geral de algoritmos, estabelecendo especificamente contração exponencial e limites de erro de amostra finita sob suposições de contratividade conjunta.
• Utilidade Prática: Demonstrar que a instanciação deste arcabouço com escolhas específicas (dinâmicas de ordem superior para MMLE, momentum para EBMs) resulta em melhorias de desempenho sobre métodos especializados existentes.

Os autores reconhecem limitações, observando que as garantias teóricas dependem de contratividade global e suposições de momentos uniformes, que podem ser restritivas para objetivos altamente não convexos em grandes redes neurais. Além disso, os resultados atuais aplicam-se a sistemas de tempo contínuo, e estabelecer limites de erro de passos finitos para esquemas de discretização específicos permanece como uma tarefa para trabalhos futuros.