Resumo Técnico: Fluindo com Confiança (FMwC)

Declaração do Problema

Modelos generativos, particularmente modelos de Flow Matching (FM), alcançaram desempenho state-of-the-art em domínios de linguagem, visão e ciência. No entanto, eles produzem saídas (por exemplo, citações alucinadas, figuras anatomicamente incorretas ou cristais termodinamicamente instáveis) sem nenhuma indicação nativa de confiabilidade. Soluções existentes para quantificação de incerteza, como Deep Ensembles ou MC-Dropout, dependem da execução de $k$ modelos separados ou $k$ trajetórias estocásticas por amostra. Essa abordagem incorre em um custo computacional de $k\times$, tornando a confiança por amostra proibitivamente cara para implantação em larga escala. Além disso, métodos baseados em verossimilhança demonstraram ser pouco confiáveis para distinguir entradas in-distribution de out-of-distribution em modelos generativos profundos. A questão central abordada é se modelos de flow matching podem aprender a quantificar sua própria incerteza para saídas geradas sem aumentar os custos de inferência ou exigir acesso a verossimilhanças.

Metodologia: Flow Matching com Confiança (FMwC)

O FMwC é uma abordagem de passagem única que equipa modelos de flow matching com uma pontuação de confiança por amostra e nível de trajetória, ao mesmo custo computacional do flow matching determinístico padrão. O método integra o Dropout Adaptativo Variacional (VAD) com propagação de momentos em forma fechada através da trajetória da Equação Diferencial Ordinária (ODE).

1. Formulação Variacional

Em vez de aprender um único campo de velocidade determinístico $v_{t,\theta}(x_t)$, o FMwC trata os pesos da rede $\omega$ como variáveis aleatórias extraídas de uma posterior variacional $q_\psi(\omega)$. O modelo maximiza o Limite Inferior da Evidência (ELBO):
$$\text{ELBO}(\psi) = \mathbb{E}_{\omega \sim q_\psi} [\log \mathcal{L}(D | \omega)] - \text{KL}[q_\psi(\omega) \parallel \pi(\omega)]$$
onde o termo de verossimilhança corresponde à perda de regressão padrão de flow matching condicional. Crucialmente, o gradiente da verossimilhança logarítmica esperada sob a posterior coincide com o gradiente da perda de regressão padrão avaliada em amostras de pesos, permitindo treinamento sem simulação.

2. Posterior Dependente do Estado via VAD

Para garantir que a incerteza seja sensível à geometria local do fluxo, o FMwC emprega o Dropout Adaptativo Variacional (VAD). Neste framework, a posterior sobre os pesos na camada $l$ é uma Gaussiana $q_\psi(\omega_l | x, t) = \mathcal{N}(\theta_l, \alpha_l(x, t)^2 \theta_l^2)$, onde a escala de dropout $\alpha_l(x, t)$ é prevista por uma pequena rede de inferência $E_\gamma$. Isso permite que a dispersão da posterior se alargue em regiões de ambiguidade geométrica (onde o campo de velocidade é difícil de ajustar) e se estreite em outros lugares, aprendido exclusivamente via ELBO.

3. Propagação Analítica de Variância

A inovação central do FMwC é propagar a variância do campo de velocidade analiticamente através tanto das camadas da rede neural quanto da trajetória de integração da ODE em uma única passagem direta:

• Através da Rede: Usando o truque de reparametrização local, a média ($\mu_l$) e a variância ($\sigma_l^2$) das ativações são propagadas camada por camada. Para camadas lineares, a recursão de variância é em forma fechada. Para não linearidades, os momentos são computados analiticamente usando quadratura Gauss-Hermite ou linearização de Taylor, evitando a necessidade de amostragem Monte Carlo.
• Através da ODE: Em cada passo de Euler, a rede produz a velocidade da média posterior $\mu_t$ e sua variância $\sigma_t^2$. O integrador avança a posição usando $\mu_t$, enquanto $\sigma_t^2$ é registrado. Isso produz uma trajetória de variância $\{\sigma_t^2\}_{t=0}^T$ em vez de um único escalar.

4. Leituras de Confiança

A trajetória de variância pode ser agregada em pontuações de confiança escalares para tarefas a jusante:

• Confiança do Ponto Final: A variância no passo final ($t=1$).
• Confiança Integrada: A soma dos desvios padrão da velocidade ao longo da trajetória.
• Razão Temporal: A razão entre a variância da janela tardia e a variância da janela inicial, que ajuda a distinguir amostras que convergem (baixa variância tardia) daquelas que permanecem ambíguas.
• Cabeças Aprendidas: Uma pequena cabeça supervisionada pode mapear a sequência completa de variância para uma pontuação específica da tarefa.

Contribuições Principais

1. Estimador de Confiança de Passagem Única: O FMwC é o primeiro método a fornecer uma pontuação de confiança por amostra calibrada para flow matching ao custo de uma única passagem direta determinística, igualando o orçamento de inferência do FM padrão. Ele alcança isso propagando analiticamente o ruído aprendido e dependente da entrada através da ODE.
2. Melhorias de Inferência sem Treinamento: A pontuação de confiança permite aplicações práticas sem retreinamento:
   • Filtragem: Melhora monotonicamente a qualidade da amostra descartando saídas de baixa confiança.
   • Edição: Identifica o instante exato ($t^*$) onde o modelo se compromete com uma estrutura (pico de variância), permitindo edições direcionadas e que respeitam restrições (por exemplo, trocas químicas em cristais) injetando ruído em $t^*$.
   • Passos Adaptativos: Concentra os passos de computação da ODE em regiões de alta variância (ambiguidade), recuperando amostras de alta qualidade com orçamentos de passos baixos onde passos uniformes falham.
3. Interpretabilidade Geométrica: O artigo estabelece que a pontuação de confiança do FMwC correlaciona-se com a divergência local do campo de velocidade aprendido. Isso fornece uma janela para o processo generativo, ligando a incerteza epistêmica à instabilidade geométrica (pontos de bifurcação) do fluxo.

Resultados Experimentais

Os autores avaliaram o FMwC em três modalidades: estimativa de densidade 2D (Checkerboard), geração de imagens (MNIST) e geração de cristais inorgânicos de novo (Flow Matching Riemanniano).

• Qualidade da Amostra: O FMwC preserva ou melhora a qualidade da amostra em comparação com o FM padrão. No Checkerboard e no MNIST, alcançou taxas de deslocamento incorreto e pontuações FID mais baixas, respectivamente, mantendo qualidade comparável em cristais.
• Custo de Inferência: O FMwC incorre em sobrecarga negligenciável em backbones maiores (por exemplo, ~1,003$\times$ para MNIST, ~1,08$\times$ para Cristais) em comparação com a linha de base determinística. Em contraste, as linhas de base de ensemble e MC-Dropout exigem $5\times$ o poder de computação.
• Desempenho de Filtragem: A filtragem baseada em pontuações FMwC alcançou pontuações de Área Sob a Curva de Precisão-Recall (AUPRC) comparáveis a ensembles de $k=5$ e MC-Dropout, mas ao custo de uma única trajetória. Por exemplo, no conjunto de dados Checkerboard, a pontuação de razão temporal alcançou uma AUPRC de 0,55, igualando a linha de base de ensemble (0,55) com 1/5 do poder de computação.
• Sucesso na Edição: Ao editar cristais injetando ruído no pico de variância ($t^*$), o FMwC alcançou taxas de sucesso significativamente mais altas (por exemplo, 39% vs 22% para trocas químicas) em comparação com a seleção aleatória de instantes, demonstrando que regiões de alta variância correspondem a pontos críticos de decisão.
• Correlação com Divergência: A trajetória de variância $\sigma_t^2$ mostrou forte correlação com a magnitude da divergência do campo de velocidade ($|\nabla \cdot v_{t,\theta}|$). No conjunto de dados Crystal, a correlação integrada atingiu um Spearman $\rho$ de 0,96, validando a interpretação geométrica da pontuação de confiança.

Significado e Alegações

O artigo alega que o FMwC transforma o flow matching de um gerador de amostras em um gerador de amostras emparelhadas com uma medida de confiabilidade. Ao aprender uma posterior dependente da entrada sobre os pesos e propagar a variância analiticamente, o FMwC torna a geração consciente da confiança um padrão prático para aplicações científicas e críticas para a segurança. Os autores enfatizam que a pontuação de confiança é uma propriedade geométrica intrínseca do fluxo único treinado, distinta do desacordo inter-modelo medido por ensembles. Isso permite um direcionamento "cirúrgico" do processo generativo, visando momentos de instabilidade sem os custos proibitivos dos métodos existentes de quantificação de incerteza. O trabalho sugere que entender e alavancar a geometria do campo de velocidade aprendido pode levar a modelos generativos mais confiáveis, interpretáveis e controláveis.