Resumo Técnico: Calibração de Modelos Generativos a Restrições Distribucionais

Declaração do Problema

Modelos generativos frequentemente sofrem de má calibração, onde as estatísticas da distribuição de amostragem (por exemplo, frequências de classes, proporções de estruturas secundárias ou representação de gênero) desviam-se sistematicamente de valores desejados. Essa questão manifesta-se em diversos domínios: modelos de imagem exibem colapso de modos; modelos de linguagem reforçam vieses sociais; e modelos de proteínas/ADN geram estruturas com frequências de subestruturas atípicas. Esses erros frequentemente decorrem de desequilíbrios nos conjuntos de dados, dinâmicas de treinamento subótimas ou ajustes post-hoc, como amostragem com temperatura baixa.

Os autores enquadram a calibração como um problema de otimização com restrições. Dado um modelo generativo base $p_{\theta_{base}}(x)$, uma estatística $h(x)$ e uma expectativa alvo $h^*$, o objetivo é encontrar uma nova distribuição $p_\theta$ que satisfaça $E_{p_\theta}[h(x)] = h^*$ enquanto minimiza a divergência de Kullback-Leibler (KL) em relação ao modelo base:
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{s.t.} \quad E_{p_\theta}[h(x)] = h^*$$
Esta formulação busca o modelo mais próximo da distribuição base que satisfaça a restrição de calibração. No entanto, resolver isso exatamente é intratável, pois tanto o objetivo quanto a restrição envolvem expectativas difíceis de computar para modelos generativos de alta dimensão.

Metodologia

Para abordar a intratabilidade, o artigo introduz dois objetivos substitutos para ajuste fino, denominados CGM-relax e CGM-reward. Ambos os métodos dependem de otimização estocástica utilizando estimativas de gradiente não tendenciosas derivadas de amostras.

1. CGM-relax (Perda de Relaxamento)

Esta abordagem substitui a restrição rígida por um termo de penalidade. A função objetivo é:
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• Violação de Restrição ($L_{viol}$): A norma $\ell_2$ ao quadrado da diferença entre a estatística empírica e o alvo. Os autores derivam um estimador não tendencioso para este termo usando uma fórmula de variância amostral corrigida por viés.
• Divergência KL ($L_{KL}$): Medida via amostragem de Monte Carlo.
• Otimização: O método usa amostragem por importância para calcular gradientes não tendenciosos. Especificamente, amostra-se do modelo atual $p_\theta$ (com stop-gradient) e reponderam-se as amostras para estimar gradientes do termo KL e da violação de restrição. Um hiperparâmetro $\lambda$ equilibra a satisfação da restrição e a proximidade com o modelo base.

2. CGM-reward (Perda de Recompensa)

Esta abordagem conecta o problema de calibração ao princípio de Máxima Entropia. A solução para o problema de máxima entropia (minimizar KL em relação ao modelo base sujeito a restrições de momentos) é um tilt exponencial do modelo base: $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• Estimativa de $\alpha^*$: Como a verdadeira distribuição de máxima entropia é intratável, o método estima primeiro os parâmetros ótimos de recompensa $\alpha^*$ resolvendo um problema de otimização dual usando $N$ amostras do modelo base.
• Função de Perda: O objetivo de ajuste fino é a divergência KL entre o modelo atual $p_\theta$ e a distribuição de máxima entropia estimada $p_{\hat{\alpha}_N}$:
  $$L_{reward}(\theta) = D_{KL}(p_\theta \parallel p_{\hat{\alpha}_N}) = D_{KL}(p_\theta \parallel p_{\theta_{base}}) + E_{p_\theta}[-r_{\hat{\alpha}_N}(x)] + C$$
  Isso é equivalente a um problema de ajuste fino por recompensa onde a "recompensa" é derivada das restrições.
• Otimização: Similar ao CGM-relax, utiliza estimadores de gradiente de função de pontuação (score-function) com bases leave-one-out para reduzir a variância.

Estimativa de Gradiente

Ambos os métodos empregam estimadores de gradiente de função de pontuação (estilo REINFORCE) combinados com amostragem por importância. Para lidar com a dependência da expectativa em relação a $\theta$, os autores usam um truque de "stop-gradient" onde amostras são retiradas de uma versão fixa do modelo atual, e os gradientes são computados via razão de densidades (função de pontuação). A redução de variância é alcançada usando bases leave-one-out.

Contribuições Principais

1. Formulação: Enquadrar a calibração de modelos generativos como um problema de minimização de KL com restrições.
2. Algoritmos: Propor dois algoritmos práticos de ajuste fino (CGM-relax e CGM-reward) que aproximam a solução do problema restrito usando otimização estocástica.
3. Escalabilidade: Demonstrar que esses métodos podem lidar com centenas de restrições simultâneas e modelos com até 9 bilhões de parâmetros.
4. Conexão Teórica: Estabelecer o vínculo entre calibração, máxima entropia e ajuste fino por recompensa, mostrando que o CGM-reward visa a solução de máxima entropia, enquanto o CGM-relax visa uma versão regularizada.

Resultados

Os autores avaliam os métodos em três aplicações diversas:

1. Design de Proteínas (Genie2 e ESM3):

   • Objetivo: Calibrar o conteúdo de estrutura secundária (hélices alfa e folhas beta) para corresponder às estatísticas de proteínas naturais (conjunto de dados CATH).
   • Restrições: Até 99 restrições simultâneas (CDFs bivariadas).
   • Resultado: O CGM-relax reduziu a distância às distribuições de proteínas naturais em mais de 4x para o Genie2 e mais de 2x para o ESM3. O CGM-reward alcançou melhorias mais modestas, parcialmente devido à dificuldade de estimar $\alpha^*$ com restrições de alta dimensão. Ambos os métodos mantiveram a plausibilidade biofísica, embora com um ligeiro aumento em "falhas de design".

2. Geração de Imagens (TarFlow):

   • Objetivo: Equilibrar as proporções de classes dentro da categoria "vida selvagem" do conjunto de dados AFHQ (por exemplo, equalizando leões, tigres, raposas e lobos).
   • Restrições: 5 restrições simultâneas.
   • Resultado: O CGM-relax reduziu a distância de variação total das proporções de classes de 0,306 para 0,101. Embora o realismo da amostra (FID) tenha diminuído ligeiramente, o modelo calibrado produziu gerações mais diversas e equilibradas. O CGM-reward não conseguiu reduzir significativamente a violação de restrição.

3. Modelagem de Linguagem (Gemma-2-9B-IT):

   • Objetivo: Reduzir o viés de gênero em histórias infantis geradas para várias profissões (por exemplo, garantindo probabilidade igual de advogados homens/femininos).
   • Restrições: 8 restrições condicionais à profissão.
   • Resultado: Ambos os métodos reduziram o desequilíbrio de gênero. O CGM-relax reduziu o desequilíbrio em mais de três vezes em relação à linha de base, enquanto o CGM-reward mostrou uma melhoria estatisticamente significativa, mas menor. Os modelos ajustados mantiveram a qualidade da história (sem degradação na log-probabilidade).

Experimentos Sintéticos:

• Eventos Raros: O CGM conseguiu aumentar o peso de modos raros em misturas gaussianas, mesmo quando a probabilidade base era tão baixa quanto $10^{-3}$.
• Altas Dimensões: O CGM-relax lidou efetivamente com até $10^3$ restrições. Em contraste, o CGM-reward tornou-se inviável para restrições $>30$ devido à dificuldade de resolver o problema empírico de máxima entropia (Equação 7).

Significado e Alegações

O artigo alega que o CGM fornece um framework prático para calibrar modelos generativos de alta dimensão para atender a restrições distribucionais sem exigir mudanças na arquitetura do modelo ou treinamento do zero.

• Robustez: O CGM-relax mostrou-se mais robusto que o CGM-reward, particularmente em configurações de restrições de alta dimensão onde estimar os parâmetros de máxima entropia se torna intratável.
• Compensações: Os métodos oferecem uma compensação navegável entre a precisão da calibração e a fidelidade ao modelo base (controlada por $\lambda$).
• Limitações: Os autores reconhecem que erros residuais permanecem, especialmente para eventos muito raros. Além disso, o framework atualmente depende de modelos com verossimilhanças tratáveis (modelos de difusão, fluxos normalizadores, modelos autoregressivos), limitando sua aplicação direta a modelos implícitos como GANs.

O trabalho posiciona a calibração como uma direção de pesquisa distinta e frutífera, separada do ajuste fino padrão por recompensa, ao usar as próprias restrições como o sinal impulsor para o alinhamento distribucional.