Resumo Técnico: Alinhamento de Modelos Generativos de Poucos Passos por Amortização de Inferência Variacional Baseada em Amostras

Declaração do Problema

Alinhar modelos generativos de poucos passos (por exemplo, GANs, modelos de consistência, mapas de fluxo, modelos de deriva) com objetivos específicos, como funções de recompensa ou preferências humanas, apresenta um desafio significativo. As estruturas de alinhamento existentes geralmente dependem de suposições estruturais restritivas que geradores de poucos passos frequentemente violam:

1. Verossimilhanças Tratáveis: Muitos métodos exigem avaliação explícita de verossimilhança sobre trajetórias de remoção de ruído, o que não está disponível para modelos implícitos como GANs ou mapeamentos de um único passo.
2. Solucionadores Específicos: Abordagens baseadas em Controle Ótimo Estocástico (SOC) ou Aprendizado por Reforço (RL) frequentemente assumem solucionadores específicos de Equações Diferenciais Ordinárias (ODE) ou Equações Diferenciais Estocásticas (SDE) e processos sequenciais de remoção de ruído.
3. Especificidade do Modelo: Formulações são frequentemente adaptadas para famílias específicas de modelos (por exemplo, modelos de difusão), falhando em generalizar para arquiteturas diversas de poucos passos, como VAEs ou mapas de fluxo.

Consequentemente, há uma falta de uma estrutura de alinhamento geral que possa ajustar finamente geradores diversos de poucos passos, preservando sua vantagem central: inferência rápida e de passo único.

Metodologia: FAV

Os autores propõem o FAV (Alinhamento de Modelos Generativos de Poucos Passos via Inferência Variacional Baseada em Amostras), uma estrutura geral que desacopla o alinhamento da família do gerador e da dinâmica de amostragem. O FAV formula o alinhamento como amostragem de uma distribuição inclinada por recompensa ancorada a uma distribuição de referência:
$$q^*(x) \propto p_{\text{ref}}(x) \exp(\beta \cdot r(x))$$
onde $p_{\text{ref}}$ é uma distribuição de referência (por exemplo, um gerador pré-treinado ou dados empíricos) e $r(x)$ é a função objetivo.

O método procede em três etapas principais:

1. Aproximação Baseada em Amostras via SVGD

O FAV aproxima a distribuição alvo $q^*$ usando Descida de Gradiente Variacional de Stein (SVGD). Em vez de exigir avaliação de densidade, o SVGD representa a distribuição com partículas e as transporta iterativamente em direção ao alvo. O campo de velocidade de Stein ótimo $\phi^*$ é derivado para maximizar a diminuição na divergência KL:
$$\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \left[ k(x', x) \nabla_{x'} \log p_{\text{ref}}(x') + \beta \cdot k(x', x) \nabla_{x'} r(x') + \nabla_{x'} k(x', x) \right]$$
Este campo compreende três componentes:

• Alinhamento de Prioridade: Impulsiona amostras em direção à distribuição de referência $p_{\text{ref}}$.
• Orientação por Recompensa: Empurra amostras para regiões de alta recompensa.
• Aplicação de Diversidade: Previne o colapso de modos por meio do termo de gradiente do kernel.

2. Estimativa de Pontuação Não Paramétrica (KDE)

Um desafio crítico é que a pontuação $\nabla \log p_{\text{ref}}$ é frequentemente intratável para dados empíricos ou geradores implícitos. O FAV aborda isso estimando a pontuação de forma não paramétrica usando Estimação de Densidade Kernel (KDE). Ao convolver a distribuição de referência com um kernel gaussiano, a pontuação é aproximada como uma média ponderada das diferenças entre amostras de referência e partículas atuais, contornando efetivamente a necessidade de verossimilhanças explícitas.

3. Amortização via Regressão de Ponto Fixo

Para preservar a velocidade de inferência rápida de modelos de poucos passos, o FAV não realiza atualizações iterativas de SVGD no momento da inferência. Em vez disso, ele amortiza o transporte nos parâmetros do gerador $\theta$ via regressão de ponto fixo. O gerador é treinado para minimizar a distância entre sua saída e o alvo "transportado":
$$\mathcal{L}(\theta) = \mathbb{E}_{\epsilon} \left[ \| \text{stopgrad}(f_\theta(\epsilon) + \hat{\phi}^*_{q_\theta, q^*_\sigma}(f_\theta(\epsilon))) - f_\theta(\epsilon) \|_2^2 \right]$$
Aqui, a amostra transportada atua como um alvo de regressão estável. Isso permite que o gerador aprenda a produzir amostras alinhadas em uma única passagem direta.

Extensão para Recompensas de Caixa-Preta: Para recompensas não diferenciáveis, o FAV emprega estimativa de gradiente de ordem zero (suavização gaussiana) para aproximar $\nabla r(x)$, permitindo alinhamento mesmo quando a função de recompensa é uma caixa-preta.

Contribuições Principais

1. Estrutura de Alinhamento Geral: O FAV é a primeira estrutura a alinhar geradores diversos de poucos passos (incluindo VAEs, GANs, modelos de consistência e mapas de fluxo) sem exigir verossimilhanças tratáveis, solucionadores específicos de ODE/SDE ou trajetórias de múltiplos passos.
2. Formulação de VI Baseada em Amostras: Ao lançar o alinhamento como amostragem de uma distribuição inclinada por recompensa usando SVGD e KDE, o FAV desacopla o procedimento de alinhamento da dinâmica interna do gerador.
3. Preservação de Eficiência: Por meio da amortização via regressão de ponto fixo, o FAV mantém a eficiência de inferência de passo único do gerador base, evitando a sobrecarga computacional de amostragem iterativa ou processos de destilação.
4. Compatibilidade com Caixa-Preta: A integração da estimativa de gradiente de ordem zero permite que o FAV otimize recompensas arbitrárias de caixa-preta.

Resultados Experimentais

Os autores avaliaram o FAV em dois domínios principais: manipulação robótica e alinhamento de geradores de imagens.

1. Manipulação Robótica (RL Offline e Offline-to-Online)

• Benchmarks: 56 tarefas de RL offline (OGBench e D4RL) e 30 tarefas de RL de offline para online.
• Desempenho: O FAV alcançou desempenho de última geração em todas as tarefas, superando 10 baselines, incluindo gradientes de política reparametrizados, Correspondência Adjoint e políticas generativas destiladas (por exemplo, FQL, IFQL).
• Descoberta Chave: O FAV demonstrou capacidades superiores de extração de política e desempenho de transição do aprendizado offline para online, mesmo sem modelos generativos auxiliares ou remoção de ruído de múltiplos passos.

2. Alinhamento de Gerador de Imagens

• Escala: Experimentos variaram de ImageNet-256 (256x256) até síntese texto-para-imagem de alta resolução (1024x1024).
• Modelos: Testado em StyleGAN-XL, modelos de deriva, iMeanFlow e SANA-Sprint.
• Métricas: Avaliado contra pontuações estéticas, pontuações de preferência humana (HPS) e detectores de segurança (NSFW).
• Desempenho:
  • O FAV consistentemente alcançou a melhor fronteira de Pareto entre otimização de recompensa e diversidade/qualidade de amostras, mitigando a superotimização de recompensa melhor do que baselines como DRaFT e Correspondência Adjoint.
  • Eficiência de Inferência: O FAV gerou amostras alinhadas aproximadamente 180x a 280x mais rápido do que métodos de alinhamento no momento da inferência (por exemplo, Best-of-N, ReNO), mantendo qualidade de imagem competitiva.
  • Escalabilidade: O método escalou com sucesso de resoluções 256x256 para 1024x1024 e generalizou-se através de diferentes arquiteturas de modelos.

Significado e Afirmações

O artigo afirma que o FAV fornece uma solução unificada e agnóstica ao modelo para alinhar modelos generativos de poucos passos. Seu significado reside em:

• Remoção de Barreiras Estruturais: Elimina a necessidade de verossimilhanças tratáveis ou trajetórias de amostragem específicas, tornando o alinhamento acessível a uma classe mais ampla de modelos generativos implícitos.
• Equilíbrio entre Eficiência e Qualidade: Reconcilia com sucesso o trade-off entre geração de alta recompensa e a eficiência computacional da inferência de poucos passos, um desafio onde métodos anteriores frequentemente degradavam o desempenho ou exigiam destilação complexa.
• Robustez: A estrutura demonstra robustez em tarefas diversas, desde controle de baixa dimensão até síntese de imagens de alta resolução, e lida com funções de recompensa diferenciáveis e de caixa-preta.

Os autores concluem que o FAV representa uma abordagem forte e de propósito geral para extração de política em RL e alinhamento de recompensa em geração de imagens, oferecendo uma alternativa prática a métodos mais complexos e carregados de suposições.