Resumo Técnico: Modelagem Generativa de Estruturas Latentes Discretas via Gradientes de Política Dinâmicos

Definição do Problema

O artigo aborda o desafio de inferir estados latentes mecanísticos não observados ($S^*$) a partir de observações indiretas ($X$) em domínios científicos onde o espaço latente é combinatorialmente grande. Diferente do aprendizado não supervisionado geral (ex: agrupamento ou modelagem de tópicos), que aprende representações latentes artificiais, este trabalho foca na reconstrução de estados mecanísticos de verdade fundamental, como vias de reações químicas, redes de transporte ou isoformas de genes.

O problema central é formulado como uma tarefa de aprendizado (Problema 2.1) para encontrar parâmetros de modelo $\theta$ para uma distribuição generativa $Pr(S|\theta)$ que maximize a verossimilhança dos dados observados $X_{1:N}$, onde $Pr(X|\theta) = \sum_{S} Pr(X|S)Pr(S|\theta)$. Uma tarefa secundária de inferência (Problema 2.2) envolve estimar o estado latente específico $\hat{S}_i$ para cada observação $X_i$.

Abordagens clássicas como o Expectation Maximization (EM) enfrentam dificuldades aqui porque o passo E (E-step) exige o cálculo de expectativas sobre espaços de estados exponencialmente grandes, o que é frequentemente intratável. Embora o Aprendizado por Reforço (RL) seja naturalmente adequado para a geração sequencial de estados estruturados, os gradientes de política padrão otimizam recompensas esperadas sob funções de recompensa estacionárias, não a verossimilhança marginal de observações indiretas.

Metodologia: GReinSS

Os autores introduzem o GReinSS (Generative Reinforcement Learning of Structured States), um framework que adapta os gradientes de política para otimizar a verossimilhança marginal de observações indiretas.

Redimensionamento Dinâmico de Recompensa

A inovação central é uma função de recompensa dinamicamente redimensionada. No RL padrão, o gradiente de política é $\mathbb{E}_{\tau}[r(\tau) \nabla_\theta \log Pr(\tau|\theta)]$. O GReinSS define a recompensa $r(\tau)$ para uma trajetória $\tau$ (que gera um estado latente $S(\tau)$) como:
$$r(\tau) = \sum_{i=1}^{N} \frac{Pr(X_i | \tau)}{Pr(X_i | \theta)}$$
onde $Pr(X_i | \theta) = \mathbb{E}_{\tau}[Pr(X_i | \tau)]$.

O Teorema 3.1 estabelece que esta função de recompensa específica torna o gradiente de política um estimador não viesado do gradiente da log-verossimilhança dos dados $\nabla_\theta \log Pr(X_{1:N} | \theta)$.

• Mecanismo: O denominador $Pr(X_i | \theta)$ redimensiona a contribuição de cada observação. Isso evita que a política colapse em uma única trajetória de alta recompensa (o que maximizaria a verossimilhança de uma observação, mas falharia nas demais) e, em vez disso, incentiva a política a aprender uma distribuição $Pr(S|\theta)$ que equilibre a verossimilhança entre todas as observações.
• Ciclo de Treinamento: O algoritmo alterna entre amostrar trajetórias, calcular recompensas dinâmicas baseadas no $\theta$ atual e atualizar $\theta$ via ascensão de gradiente.

Aprendizado Off-Policy

Para lidar com a esparsidade de trajetórias que resultam em verossimilhança não nula para observações específicas, o artigo utiliza amostragem off-policy. O Teorema 3.3 identifica a distribuição de proposta off-policy ótima para minimização de variância como $q(\tau | X_{1:N}, \theta) = \frac{1}{N} \sum_{i=1}^N Pr(\tau | X_i, \theta)$. Embora a amostragem exata desta distribuição seja frequentemente intratável, os autores utilizam heurísticas (ex: enviesar a amostragem com base em características das observações) combinadas com correções de amostragem por importância para manter estimativas de gradiente não viesadas.

Mini-batching

O framework suporta mini-batching (Corolário 3.2), onde as recompensas são calculadas sobre um subconjunto de observações $B \subset [N]$, permitindo que o método escale para grandes conjuntos de dados.

Principais Contribuições

1. Formulação: O artigo formula o problema de aprender distribuições sobre estados latentes discretos a partir de observações indiretas como um problema de otimização de política, unindo o aprendizado de probabilidade e o aprendizado por reforço.
2. Conexão Teórica: Fornece uma prova teórica (Teorema 3.1) de que gradientes de política com recompensas redimensionadas dinamicamente são equivalentes à maximização da verossimilhança marginal de observações indiretas, um resultado não alcançado pelo RL padrão ou pelas formulações de Inferência Variacional (VI).
3. Generalização: Diferente de aplicações anteriores de técnicas similares (ex: filogenia de câncer ou variação de número de cópias), este trabalho generaliza o procedimento para quaisquer estruturas discretas (grafos, conjuntos, sequências) e estabelece conexões com métodos existentes como Generalized Expectation Maximization (GEM) e GFlowNets sob condições específicas.

Resultados Experimentais

Os autores avaliam o GReinSS em conjuntos de dados simulados e do mundo real, comparando-o com busca local, gradientes de política padrão, GFlowNets e implementações baseadas em GEM usando VAEs, modelos autorregressivos e difusão discreta.

1. Inferência de Grafo de Processo Simulado

• Tarefa: Reconstruir grafos direcionados a partir de endpoints de caminhadas aleatórias.
• Resultados: O GReinSS alcançou consistentemente os maiores escores medianos de F1. Notavelmente, com baixa informação por observação ($k=10$ caminhadas aleatórias), o GReinSS atingiu um F1 de 0,891, enquanto todos os baselines (incluindo gradientes de política ingênuos e GFlowNets) pontuaram abaixo de 0,55. O baseline de gradiente de política ingênuo frequentemente colapsava para prever grafos vazios, destacando a necessidade do redimensionamento dinâmico de recompensa.

2. Inferência de Subconjunto Simulado

• Tarefa: Reconstruir subconjuntos a partir de medições de elementos ruidosos.
• Resultados: O GReinSS superou todos os baselines em vários níveis de ruído ($\sigma$) e tamanhos de universo ($|U|$).
  • Para universos pequenos ($|U|=10$), a maioria dos métodos performou bem.
  • Para universos grandes ($|U|=1000$), os baselines (VAE, difusão, autorregressão) colapsaram (mediana de F1 < 0,4), enquanto o GReinSS manteve uma mediana de F1 de 0,938.
  • Os resultados demonstram que o GReinSS escala efetivamente para grandes espaços combinatórios onde outros modelos generativos falham.

3. Aplicação no Mundo Real: Reconstrução de Isoformas de RNA

• Tarefa: Inferir estruturas e proporções de isoformas de RNA a partir de dados de sequenciamento de leitura curta (dataset GTEx), validados contra o ground truth de sequenciamento de leitura longa.
• Comparação: Comparado ao RSEM, o algoritmo padrão baseado em Expectation Maximization usado no projeto GTEx.
• Resultados: O GReinSS superou significativamente o RSEM.
  • Para o gene MBD2, o RSEM previu uma isoforma principal com proporção de 18% (vs. 96% do ground truth) e isoformas menores incorretas, resultando em um erro de predição de 0,5374. O GReinSS previu 92% para a isoforma principal com um erro de 0,0067.
  • Em 14.390 genes, o GReinSS superou o RSEM por uma margem de $\ge 0,05$ em 46,6% dos casos, enquanto o RSEM superou o GReinSS em apenas 9,4% dos casos. A diferença mediana de erro foi de $-0,0405$ a favor do GReinSS.

Significância e Alegações

O artigo afirma que o GReinSS é uma "abordagem fundamentada e praticamente eficaz" para modelagem generativa em espaços combinatórios. Sua principal significância reside em:

• Superação da Intratabilidade: Fornece uma solução para a inferência de estados latentes em espaços onde o EM exato é inviável devido à complexidade exponencial.
• Superioridade sobre o RL Padrão: Demonstra que os gradientes de política padrão falham neste contexto sem o redimensionamento dinâmico de recompensa, e que os GFlowNets (que equilibram probabilidades de trajetória) são insuficientes quando a estrutura de recompensa não se alinha perfeitamente com as verossimilhanças das observações.
• Impacto no Mundo Real: A aplicação bem-sucedida à splicing de RNA, superando um padrão amplamente utilizado (RSEM) em dados biológicos reais, valida a utilidade do método além de benchmarks sintéticos.

Os autores observam que, embora o método atualmente dependa de heurísticas para amostragem off-policy e assuma um $Pr(X|S)$ conhecido, o framework oferece uma base para extensões futuras, incluindo Q-learning, frameworks de actor-critic e o aprendizado do próprio modelo de observação.