Resumo Técnico: Além do Proxy: Guia Destilado por Trajetórias para Treinamento Offline de GFlowNet

Declaração do Problema

Redes de Fluxo Generativo (GFlowNets) são modelos poderosos para amostrar objetos composicionais proporcionais às suas recompensas. No entanto, em muitas aplicações práticas — como descoberta de moléculas, design de sequências biológicas e sistemas de recomendação — a interação ativa com o ambiente para consultar recompensas é inviável devido aos altos custos experimentais ou à necessidade de avaliação humana. Consequentemente, as GFlowNets devem ser treinadas usando conjuntos de dados estáticos offline.

Os paradigmas atuais de treinamento offline geralmente dependem do treinamento de um modelo de recompensa proxy para aproximar sinais de recompensa a partir do conjunto de dados, que a GFlowNet consulta então para otimização da política. Essa abordagem enfrenta desafios significativos:

1. Escassez de Dados e Custo: Construir um modelo proxy confiável frequentemente requer grandes conjuntos de dados diversos e expertise significativa no domínio.
2. Propagação de Erros: Consultas de recompensa para amostras ou estados fora da distribuição (OOD) estimados com imprecisão pelo proxy podem propagar erros através dos gradientes, comprometendo a política aprendida.
3. Limitações de Métodos sem Proxy: Abordagens existentes sem proxy (por exemplo, RO-GFlowNets, COFlowNet) frequentemente impõem restrições grosseiras para alinhar a política ao conjunto de dados. Essas restrições podem limitar a generalização, inibir a exploração eficaz e aumentar a sensibilidade à qualidade dos dados.

A questão central abordada é: Como podemos aproveitar efetivamente trajetórias offline para treinar GFlowNets quando as recompensas terminais são insuficientes para construir um modelo proxy confiável?

Metodologia: GFlowNet Destilado por Trajetórias (TD-GFN)

Os autores propõem o TD-GFN, um novo paradigma sem proxy que extrai informações estruturais ricas do Grafo Acíclico Direcionado (DAG) do ambiente para guiar o treinamento da política sem depender de um modelo de recompensa proxy. A metodologia consiste em três fases principais:

1. Extração de Recompensa de Aresta via Aprendizado por Reforço Inverso (IRL)

Em vez de prever recompensas terminais, o TD-GFN utiliza IRL para destilar recompensas de aresta no nível de transição ($R_E$) a partir de trajetórias offline.

• Reequilíbrio: Para abordar o fato de que dados offline podem não representar comportamento de especialista, os autores introduzem uma estratégia de reequilíbrio. Trajetórias são amostradas com probabilidade proporcional às suas recompensas terminais ($P(\tau) \propto R(s_T)$), criando um conjunto de dados pseudo-experto.
• Treinamento Adversarial: Usando um framework modificado de GAIL (Aprendizado de Imitação Adversarial), um discriminador $D_\phi$ distingue entre transições no conjunto de dados reequilibrado e aquelas geradas por uma política $\pi_\psi$.
• Derivação de Recompensa: Uma função de recompensa de aresta no nível de aresta não enviesada é extraída do discriminador:
  $$R_E(s, s') = \log D_\phi(s, s') - \log(1 - D_\phi(s, s'))$$
  Teoricamente, essa recompensa aproxima a pontuação de fluxo de entrada condicional em log ($\log P_B(s|s')$), quantificando a importância de transições específicas para aprender uma política eficaz. Crucialmente, essas recompensas fornecem orientação intermediária densa, em vez de previsões terminais.

2. Poda do DAG Guiada por Recompensa

Para garantir robustez e prevenir a propagação de erros de estimativas de IRL potencialmente imprecisas, o TD-GFN não usa recompensas de aresta para moldar diretamente os gradientes da política. Em vez disso, ele as usa para podar o DAG do ambiente.

• Limiarização: Arestas com recompensas abaixo de um limiar estatístico (média menos $K$ vezes o desvio padrão da distribuição de recompensas) são removidas.
• Garantia Teórica: Os autores provam que essa estratégia de poda preserva todos os modos de "alto gargalo" (estados terminais alcançáveis via caminhos com altas pontuações de aresta) enquanto remove transições de baixa utilidade. Isso cria um espaço de ação compacto e expressivo que prioriza trajetórias de alto valor, ao mesmo tempo que isola a política de erros de aproximação do IRL.

3. Amostragem Reversa Priorizada

O treinamento é conduzido no subgrafo podado usando um mecanismo de amostragem reversa priorizada.

• Estratégia de Amostragem: Estados terminais são amostrados a partir da interseção do conjunto de dados e do grafo podado com probabilidade proporcional às suas recompensas verdadeiras (ground-truth).
• Construção de Trajetória Reversa: A partir de cada terminal amostrado, trajetórias são construídas em reverso em direção à raiz usando uma política reversa aprendida $P_B$, definida pelas recompensas de aresta aprendidas:
  $$P_B(s_t | s_{t+1}) = \frac{\exp\{R_E(s_t, s_{t+1})\}}{\sum \exp\{R_E(s, s_{t+1})\}}$$
• Atualizações de Gradiente: A política da GFlowNet é otimizada usando essas trajetórias amostradas e suas recompensas terminais verdadeiras. Ao depender exclusivamente de recompensas terminais registradas para atualizações de gradiente, o método evita os problemas de propagação de erros inerentes a métodos baseados em proxy.

Contribuições Principais

1. Um Novo Paradigma sem Proxy: O TD-GFN elimina a dependência de modelos proxy e consultas de recompensa OOD. Ele emprega IRL para destilar recompensas de aresta no nível de transição a partir de conjuntos de dados reequilibrados, fornecendo orientação estruturada para treinamento e generalização da política.
2. Mecanismo de Orientação Robusto: O framework introduz dois mecanismos de orientação indireta: poda de DAG e amostragem reversa priorizada. Esse design garante que as atualizações de gradiente dependam exclusivamente de recompensas verdadeiras, prevenindo a propagação de erros enquanto ainda aproveita orientação estrutural densa para exploração eficiente.
3. Desempenho de Estado da Arte: Experimentos extensivos demonstram que o TD-GFN alcança velocidade de convergência superior e qualidade de amostra em comparação com uma ampla gama de baselines, incluindo métodos baseados em proxy (Proxy-GFN), algoritmos de RL offline (CQL, IQL) e outras variantes de GFlowNet sem proxy (COFlowNet, Dataset-GFN).

Resultados Experimentais

Os autores avaliaram o TD-GFN em três benchmarks:

• Hypergrid (Sintético): Em uma grade de 84 dimensões com recompensas esparsas e multimodais, o TD-GFN identificou todos os 16 modos de recompensa com menos de 5.000 visitas a estados (uma melhoria de seis vezes sobre as baselines) e convergiu para o menor Erro L1 Empírico duas vezes mais rápido que o melhor concorrente (COFlowNet). Demonstrou forte robustez sob dados ruidosos (conjunto de dados Mixed) e escassez de dados (10% dos dados).
• Design de Biosequências (Peptídeos Antimicrobianos): O TD-GFN gerou sequências com as maiores recompensas e maior diversidade entre as 100 sequências principais, superando tanto métodos baseados em proxy quanto outros métodos offline.
• Design de Moléculas: O TD-GFN descobriu mais moléculas de alta recompensa e modos diversos do que as baselines, exigindo significativamente menos trajetórias para convergir (2,75 × 10^4 vs. 43,7 × 10^4 para Proxy-GFN). Ele igualou o desempenho de um Oracle-GFN (que usa recompensas verdadeiras diretamente) na identificação de moléculas de alta recompensa.

Significado e Alegações

O artigo afirma que o TD-GFN estabelece um paradigma mais robusto, estável e eficiente para treinamento offline de GFlowNet. Ao aproveitar as informações estruturais inerentes às trajetórias offline através do IRL, o método preenche a lacuna entre conjuntos de dados offline limitados e os requisitos complexos de exploração da modelagem generativa.

Os autores enfatizam que o TD-GFN:

• Desbloqueia Potencial: Permite a implantação de GFlowNets em domínios complexos onde consultas ativas de recompensa são proibitivas.
• Previne Propagação de Erros: Ao desacoplar o mecanismo de orientação (recompensas de aresta) do mecanismo de atualização de gradiente (recompensas terminais), evita as armadilhas dos modelos proxy.
• Melhora a Exploração: A orientação estrutural densa permite que o modelo explore regiões de alta recompensa efetivamente, mesmo quando elas são representadas esparsamente ou não observadas nos dados de treinamento.

O trabalho sugere que informações no nível de trajetória, frequentemente negligenciadas em métodos baseados em proxy, são uma pedra angular para treinamento eficaz de GFlowNet offline, oferecendo uma solução escalável para aplicações do mundo real na descoberta de medicamentos, design biológico e otimização combinatória.