Gradient-Free Noise Optimization for Reward Alignment in Generative Models

Resumo Técnico: ZeNO (Otimização de Ruído de Ordem Zero)

Declaração do Problema
Os métodos existentes de alinhamento de recompensa para modelos generativos (por exemplo, modelos de difusão e fluxo) geralmente dependem da otimização de trajetórias estocásticas de múltiplos passos. Embora eficazes nessas configurações, tais abordagens são difíceis de estender para geradores determinísticos ou de um único passo (por exemplo, modelos de consistência, modelos de difusão destilados), que não expõem naturalmente a mesma estrutura de trajetória. Uma alternativa natural é otimizar diretamente o ruído de entrada. No entanto, a otimização de ruído baseada em gradientes padrão requer retropropagação através tanto do gerador quanto da função de recompensa. Isso limita a aplicabilidade a pipelines diferenciáveis e pode ser pouco confiável quando os gradientes refletem sensibilidades locais espúrias em vez de melhorias semânticas significativas. Além disso, muitas funções de recompensa práticas (por exemplo, projetabilidade de proteínas, alinhamento composicional texto-imagem) são não diferenciáveis ou caixas-pretas, tornando os métodos baseados em gradientes inviáveis.

Metodologia: ZeNO
Os autores propõem o ZeNO (Otimização de Ruído de Ordem Zero), um framework sem gradiente que formula a otimização de ruído como um problema de controle de integral de caminho (PIC). A ideia central é tratar a evolução do ruído de entrada como um processo estocástico controlável, otimizando-o usando apenas avaliações de recompensa de ordem zero.

1. Formulação de Controle de Integral de Caminho:
   O método modela a otimização de ruído como uma Equação Diferencial Estocástica (EDE) controlada. O objetivo é direcionar a trajetória do ruído para regiões que produzem estados terminais de alta recompensa, enquanto regulariza o esforço de controle. O controle ótimo $u^*(z, t)$ é derivado em uma forma de integral de caminho:
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   Crucialmente, este estimador depende apenas de avaliações de custo terminal (recompensas) e não requer diferenciação através do gerador $G_\theta$ ou da função de recompensa $r$.

2. Dinâmica de Referência (Ornstein–Uhlenbeck):
   Para garantir que o ruído atualizado permaneça compatível com geradores pré-treinados, o ZeNO adota o processo de Ornstein–Uhlenbeck (OU) como dinâmica de referência. Ao contrário das EDEs de explosão de variância, o processo OU preserva a distribuição estacionária Gaussiana $N(0, I)$, que corresponde ao prior da maioria dos geradores pré-treinados. Essa escolha conecta a atualização do ruído à dinâmica de Langevin, visando implicitamente uma distribuição inclinada pela recompensa $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$.

3. Estimador Linearizado e Redução de Variância:
   O estimador padrão de integral de caminho sofre de alta variância devido ao ponderamento exponencial. Os autores introduzem uma aproximação linearizada (Proposição 1) subtraindo a recompensa média $\bar{r}$ como linha de base e linearizando o termo exponencial sob a suposição de baixa variância de recompensa:
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   Este estimador coincide com o estimador da função de pontuação usado em Estratégias de Evolução Natural e supera empiricamente as variantes exponenciais.

4. Horizonte Recessivo de Um Passo:
   Para evitar o custo computacional de simulações de longo horizonte, o ZeNO emprega uma aproximação de horizonte recessivo de um passo ($H=1$). Em cada iteração, $N$ candidatos de ruído perturbados são amostrados a partir da dinâmica de referência, avaliados pelo modelo de recompensa e a estimativa de controle linearizada é usada para atualizar o estado atual do ruído. Isso permite a escalabilidade no momento da inferência aumentando o número de partículas ($N$) ou iterações ($M$) sem modificar o gerador.

Principais Contribuições

• Alinhamento sem Gradiente: O ZeNO permite o alinhamento de recompensa para geradores determinísticos e de um único passo sem exigir retropropagação através do gerador ou do modelo de recompensa.
• Compatibilidade com Caixa-Preta: O método é aplicável a funções de recompensa não diferenciáveis e de caixa-preta, um cenário onde os métodos existentes de otimização de ruído baseados em gradientes (por exemplo, ORIGEN, ReNO) falham.
• Conexão Teórica: O framework estabelece um link teórico entre a otimização de ruído de ordem zero e a amostragem de Langevin visando uma distribuição inclinada pela recompensa, fornecendo uma justificação principiante para a regra de atualização.
• Escalabilidade no Momento da Inferência: O ZeNO introduz um mecanismo flexível para trocar o orçamento computacional (via $N$ e $M$) por desempenho de recompensa aprimorado.

Resultados Experimentais
Os autores avaliam o ZeNO em diversos geradores (SDXL-Turbo, DMD2, LCM, DFGAN) e funções de recompensa (pontuação estética, PickScore, GenEval e projetabilidade de proteínas).

• Geração de Imagens: O ZeNO melhora consistentemente as pontuações de recompensa em todos os geradores e tipos de recompensa testados. Ele alcança desempenho competitivo ou superior em comparação com baselines baseadas em gradientes (ORIGEN, $\Psi$-Sampler) e amostragem Best-of-N. Notavelmente, no gerador LCM, os métodos baseados em gradientes degradaram-se abaixo da baseline não controlada, enquanto o ZeNO manteve a melhoria, destacando a robustez dos métodos de ordem zero contra sensibilidades locais de gradiente.
• Recompensas Não Diferenciáveis: Na tarefa GenEval (contagem e posicionamento de objetos), onde as recompensas são não diferenciáveis, o ZeNO superou o Best-of-N em tarefas de contagem, demonstrando que a otimização iterativa de ruído extrai mais valor do mesmo orçamento computacional do que a amostragem independente.
• Geração de Estrutura de Proteínas: Aplicado a um gerador de esqueleto de proteína Riemannian MeanFlow (RMF) de um único passo, o ZeNO melhorou significativamente a projetabilidade (medida pelo RMSD de autoconsistência). Ele alcançou níveis de projetabilidade comparáveis a uma baseline RMF de 5 passos usando apenas uma única avaliação de fluxo, demonstrando sua capacidade de compensar passos de amostragem reduzidos via otimização de ruído.
• Diversidade: A análise usando a Pontuação Vendi indica que o ZeNO melhora a recompensa enquanto preserva a diversidade das amostras, evitando o colapso para um único modo frequentemente observado em abordagens baseadas em gradientes.

Significado e Alegações
O artigo afirma que o ZeNO fornece uma solução prática e teoricamente fundamentada para alinhar a classe crescente de modelos generativos eficientes, determinísticos e de um único passo com funções de recompensa. Seu significado principal reside em sua capacidade de operar em cenários onde gradientes estão indisponíveis ou são pouco confiáveis, como em modelos de recompensa de caixa-preta ou tarefas complexas de design biológico. Ao formular a otimização de ruído como um problema de controle de integral de caminho, o ZeNO oferece um framework escalável e agnóstico ao gerador para alinhamento no momento da inferência que não requer retreinamento ou modificação do modelo generativo subjacente. Os autores notam uma limitação: o método requer múltiplas avaliações de recompensa por passo de atualização, o que pode ser computacionalmente exigente se o próprio modelo de recompensa for caro.