Interpretable epistemic uncertainty decomposition in sequential generative models via polynomial chaos surrogates

Resumo Técnico: Decomposição Interpretável da Incerteza Epistêmica em Modelos Generativos Sequenciais via Surrogados de Caos Polinomial

Declaração do Problema

Modelos generativos sequenciais, como Redes de Fluxo Generativo (GFlowNets), são cada vez mais centrais para a descoberta científica impulsionada por IA, otimização combinatória e design biológico. Esses modelos frequentemente operam sob condições de incerteza epistêmica, especificamente quando as funções de recompensa que orientam sua geração são, elas mesmas, incertas (por exemplo, surrogados de redes neurais treinados com dados experimentais limitados). Embora métodos existentes, como ensembles profundos e redes neurais bayesianas, possam quantificar a magnitude da incerteza de previsão, eles falham em fornecer uma decomposição interpretável dessa incerteza. Consequentemente, os praticantes não podem determinar quais componentes específicos da incerteza da recompensa de entrada impulsionam quais decisões generativas, nem podem distinguir entre sinais químicos robustos e artefatos estatísticos decorrentes de vieses no conjunto de treinamento.

Metodologia

Os autores propõem uma estrutura de modelagem de surrogado que propaga a incerteza epistêmica através de GFlowNets usando Expansões de Caos Polinomial (PCE). A abordagem procede em três etapas:

1. Treinamento de Ensemble e Parametrização: Um ensemble modesto de GFlowNets é treinado sob condições de recompensa variadas. O espaço da função de recompensa é parametrizado em um espaço latente de baixa dimensão $\mu \in \mathbb{R}^d$ usando Análise de Componentes Principais (PCA) ou expansões analíticas (por exemplo, Karhunen–Loève). O PCA é preferido porque seus componentes ortogonais aproximam entradas Gaussianas independentes, um pré-requisito para uma decomposição analítica de Sobol válida.
2. Ajuste do Surrogado PCE: Para cada ação $k$ no passo de trajetória $t$, um PCE é ajustado para mapear os parâmetros de recompensa de baixa dimensão $\mu$ para as probabilidades de ação transformadas por log-ratio (usando uma transformação Log-Ratio Aditiva em relação a uma ação de referência). Os coeficientes são estimados via regressão ridge regularizada.
3. Decomposição Analítica: Os coeficientes do PCE ajustado permitem o cálculo analítico dos índices de sensibilidade de Sobol. Diferentemente de métodos de Monte Carlo, que exigem integração dispendiosa, esses índices são derivados diretamente dos coeficientes polinomiais, fornecendo uma decomposição exata baseada na variância de quanto cada componente da incerteza da recompensa contribui para a variância na distribuição da política.

A estrutura inclui garantias formais de convergência para o surrogado e seus índices de Sobol. Quatro de cinco propriedades teóricas chave (convergência do PCE, convergência do índice de Sobol, preservação do simplex e limites de propagação de incerteza) foram formalmente verificadas no assistente de prova Lean 4.

Contribuições Principais

1. Estrutura Geral: Um método computacionalmente eficiente para propagar e decompor a incerteza epistêmica em modelos generativos sequenciais condicionados usando surrogados PCE ajustados a pequenos ensembles.
2. Verificação Formal: Garantias teóricas de convergência para o surrogado e seus índices de sensibilidade, com quatro de cinco teoremas formalmente verificados no Lean 4.
3. Decomposição Interpretável de Sensibilidade: O primeiro método analítico para decompor a incerteza da política em modelos generativos, revelando exatamente quais componentes da recompensa impulsionam decisões generativas específicas. Essa capacidade não está disponível em ensembles profundos padrão, redes neurais bayesianas ou dropout de Monte Carlo.
4. Validação Empírica: Demonstrações em sete tarefas de complexidade crescente, incluindo três aplicações do mundo real e quatro tarefas de validação sintéticas.

Resultados e Aplicações

A estrutura foi avaliada em três domínios primários do mundo real:

• Otimização de Reações Químicas (Buchwald–Hartwig): No conjunto de dados Doyle–Dreher, a estrutura analisou a seleção de catalisador, base, haleto de arila e aditivo. Revelou que a seleção de catalisador é robusta (baixo índice de Sobol, $D \approx 71$), indicando sinal químico genuíno, enquanto a seleção de aditivo é frágil (alto índice de Sobol, $D \approx 179$, 2,5$\times$ maior), indicando suscetibilidade ao viés do conjunto de treinamento. Isso aborda diretamente o "debate Chuang–Keiser" sobre artefatos estatísticos em preditores de rendimento.
• Descoberta Causal Bayesiana (Sinalização de Proteínas Sachs): Aplicado à rede Sachs, a análise mostrou que, embora todas as arestas exibissem alta variância devido a dados limitados (200 observações), a incerteza se decomporia em regimes distintos. As arestas da via MAPK foram impulsionadas por um componente principal da variação de covariância, enquanto as arestas do hub PKA/PKC foram impulsionadas por outro. Isso fornece um mapa direcionado para experimentos de perturbação.
• Design Molecular Baseado em Fragmentos: Em uma tarefa de design de farmacabilidade, a estrutura derrubou a suposição convencional de que andaimes são robustos e decorações são frágeis. Os resultados mostraram que a posição do ligante era o componente mais frágil ($D \approx 28,2$), enquanto as posições de andaime e decoração eram mais robustas. Isso sugere que os esforços experimentais devem priorizar a diversidade de ligantes.

Em tarefas sintéticas (mundo em grade, regressão simbólica e proxies de raciocínio de LLM), o surrogado PCE alcançou cobertura de calibração de 95% entre 0,97 e 1,00 para passos dominantes. Computacionalmente, o surrogado avaliou 10.000 amostras de política em milissegundos, oferecendo uma aceleração de $10^3$ a $10^4\times$ comparado ao retreinamento exaustivo do ensemble.

Significado e Alegações

O artigo alega que esta estrutura oferece uma capacidade única: decomposição interpretável da incerteza. Enquanto outros métodos podem afirmar que uma previsão é incerta, esta abordagem explica por que, atribuindo variância a componentes específicos da recompensa. Isso permite que cientistas:

• Distinguam entre sinais robustos e artefatos estatísticos em dados de treinamento.
• Identifiquem quais partes de um espaço de design (por exemplo, ligantes versus andaimes) requerem validação experimental adicional.
• Aloquem recursos experimentais de forma eficiente, visando as fontes específicas de incerteza epistêmica.

Os autores enfatizam que o método é agnóstico à arquitetura do GFlowNet e ao objetivo de treinamento. Eles notam limitações, como a dependência de uma estrutura de baixo rank na paisagem de recompensa para que o PCA seja eficaz e a necessidade de reajustar o surrogado para diferentes trajetórias (embora não seja necessário retreinar o ensemble). O trabalho posiciona-se como uma ponte entre a quantificação de incerteza em engenharia computacional e a IA generativa moderna, fornecendo um caminho escalável para entender a "caixa preta" da tomada de decisão sequencial sob incerteza.