Resumo Técnico: f-Trajectory Balance

Declaração do Problema

O alinhamento e o ajuste fino de modelos generativos, incluindo Modelos de Linguagem de Grande Escala (LLMs) e Redes de Fluxo Generativo (GFlowNets), geralmente dependem da minimização da discrepância entre a distribuição de um modelo e uma distribuição-alvo. No ajuste fino por Aprendizado por Reforço (RL), isso é frequentemente enquadrado como a maximização da recompensa esperada sujeita a uma penalidade de divergência de Kullback-Leibler (KL).

Um desafio significativo neste domínio é o trade-off entre comportamentos de busca de modo e cobertura de modo. Objetivos padrão, como a divergência KL reversa (frequentemente implementada via Trajectory Balance ou perda KL quadrada), tendem a ser de busca de modo, fazendo com que os modelos colapsem em poucos modos de alta probabilidade. Isso é subótimo para tarefas que exigem exploração, como descoberta de fármacos ou geração diversificada de candidatos, onde a "cobertura de modo" (capturar a diversidade completa da distribuição-alvo) é preferida.

Além disso, métodos existentes frequentemente lutam com a validade off-policy. Embora estimadores como REINFORCE ou PPO sejam comuns, eles sofrem de alta variância. Abordagens recentes que usam a perda KL quadrada ($KL_{sq}$) oferecem gradientes de baixa variância on-policy e validade off-policy, mas estão inerentemente ligadas às propriedades da divergência KL reversa, limitando a capacidade de controlar o trade-off entre exploração e exploração.

Metodologia

Os autores propõem o f-Trajectory Balance, uma família de funções de perda de substituição que generaliza a perda KL quadrada para toda a família de f-divergências.

Fundamentação Teórica

A contribuição teórica central é estabelecer uma correspondência um para um entre funções de perda invariantes por translação em log-probabilidades e f-divergências.

1. Da Divergência para a Perda: Para qualquer função convexa $f$ definindo uma f-divergência $D_f(p_\theta \| p^*)$, os autores derivam uma função de perda $L_f$ definida na diferença de log-probabilidades $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$:
   $$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$$
   Esta perda possui duas propriedades críticas:

   • Equivalência de Gradiente On-Policy: Quando amostras são extraídas da política do modelo ($\mu = p_\theta$), os gradientes esperados auto-diferenciáveis de $L_f$ correspondem aos gradientes da f-divergência correspondente $D_f$.
   • Validade Off-Policy: Quando avaliada em dados off-policy (amostrados de uma distribuição de comportamento $\mu$), a perda permanece um objetivo válido com o mesmo minimizador global ($p_\theta = p^*$), desde que o suporte do alvo esteja contido no suporte da política de comportamento.

2. Da Perda para a Divergência: Inversamente, qualquer perda estritamente convexa e invariante por translação $\ell$ em log-probabilidades corresponde à minimização de uma f-divergência específica. O artigo prova que esses mapeamentos são inversos um do outro.

Implementações Práticas

Para abordar desafios práticos no treinamento, os autores introduzem várias extensões:

• Perda DevGrad (Normalização por Lote): Para lidar com funções de partição intratáveis (comuns em GFlowNets e alvos não normalizados), os autores generalizam o estimador VarGrad. Em vez de usar a variância do lote, eles usam o desvio generalizado por lote das diferenças de log-probabilidade. Isso envolve estimar a constante de normalização $\log Z$ minimizando a perda sobre o lote, efetivamente centralizando os coeficientes da função de pontuação e reduzindo a variância.
• Perda Temperada: Para evitar overflow numérico ao lidar com parâmetros de temperatura $\beta$ pequenos (onde a distribuição-alvo é $p^* \propto \pi_{ref} \exp(r/\beta)$), os autores propõem uma perda temperada. Isso escala a perda por $1/\beta$ e a aplica à distribuição temperada, garantindo que os gradientes permaneçam limitados e estáveis.
• f-Trajectory Balance para GFlowNets: O framework é aplicado a GFlowNets impondo restrições de equilíbrio de trajetória usando a perda $L_f$. Isso generaliza o Trajectory Balance padrão (que corresponde à KL reversa) para f-divergências arbitrárias.

Principais Contribuições

1. Família de Perdas Generalizada: Derivação de uma família geral de perdas de substituição invariantes por translação ($L_f$) cujos gradientes on-policy correspondem aos da f-divergência correspondente, mantendo a validade off-policy.
2. Equivalência Teórica: Prova da relação inversa entre perdas de log-probabilidade invariantes por translação e f-divergências, permitindo o design sistemático de novas funções de perda com propriedades desejadas (por exemplo, cobertura de modo).
3. Extensão para Alvos Não Normalizados: Generalização da técnica VarGrad para f-divergências arbitrárias via perda DevGrad, abordando o problema da função de partição intratável no ajuste fino de RL.
4. Controle da Exploração: Demonstração de que a escolha da f-divergência (parametrizada por $\alpha$ em $\alpha$-divergências) controla diretamente o comportamento de busca de modo versus cobertura de modo. Valores de $\alpha$ mais baixos (por exemplo, KL Direta, Hellinger) promovem cobertura de modo, enquanto valores mais altos (por exemplo, $\chi^2$ de Pearson) promovem busca de modo.

Resultados Experimentais

Os autores avaliam a família de perdas em quatro tarefas distintas:

1. Grid Sintético (GFlowNets): Em uma tarefa de grade 2D projetada para testar a cobertura de modo, perdas com $\alpha$ mais baixo (KL Direta, Hellinger) descobriram com sucesso todos os quatro modos da distribuição-alvo significativamente mais rápido e com mais confiabilidade do que o Trajectory Balance padrão (KL Reversa). Por outro lado, o $\chi^2$ de Pearson ($\alpha=2$) exibiu forte comportamento de busca de modo, ficando preso no primeiro modo encontrado.
2. Descoberta de Moléculas (SynFlowNet): Aplicado ao framework SynFlowNet para geração de moléculas sinteticamente acessíveis. Os autores descobriram que ajustar $\alpha$ permitiu controlar a diversidade das moléculas geradas. Um cronograma de recozimento (começando com $\alpha$ mais baixo para exploração e aumentando para $\alpha=1$ para exploração) produziu o melhor trade-off, gerando moléculas de maior recompensa com maior diversidade em comparação ao Trajectory Balance padrão.
3. Modelos de Difusão: Na amostragem condicional de dígitos MNIST (ímpares vs. pares), o Trajectory Balance padrão falhou em amostrar uniformemente entre os modos-alvo (superamostrando dígitos específicos como 0 e 6). As perdas f-Trajectory Balance, particularmente aquelas com $\alpha$ mais baixo, alcançaram uma amostragem mais uniforme entre os modos-alvo.
4. Ajuste Fino Assíncrono de LLMs: As perdas foram aplicadas para ajustar finamente LLMs (famílias Qwen e OLMo) em problemas de matemática (GSM8k, MATH) com um atraso assíncrono de 50 passos. As perdas de f-divergência demonstraram treinamento estável e trade-offs consistentes entre entropia e recompensa, enquanto implementações padrão de PPO exibiram instabilidade sob as mesmas condições assíncronas.

Significado e Alegações

O artigo afirma que a eficácia da perda KL quadrada não é um fenômeno isolado, mas uma instância específica de uma família mais ampla de perdas. Ao estender a propriedade de "correspondência de gradiente on-policy" e "validade off-policy" para toda a família de f-divergências, os autores fornecem um framework unificado para o ajuste fino de modelos generativos.

O significado principal reside na capacidade de controlar explicitamente o trade-off entre exploração e exploração simplesmente selecionando uma função de perda de substituição diferente (ou ajustando um parâmetro $\alpha$), sem alterar o algoritmo de treinamento subjacente. Isso permite que praticantes escolham entre comportamentos de busca de modo (para maximização de recompensa) e cobertura de modo (para diversidade e exploração), mantendo a estabilidade e a compatibilidade off-policy necessárias para o treinamento assíncrono em grande escala moderno de LLMs e GFlowNets. O trabalho sugere que este framework é particularmente valioso para aplicações como descoberta de fármacos e agentes exploratórios, onde cobrir o espaço de candidatos de alta recompensa é crítico.