Resumo Técnico: Geradores de Boltzmann de Poucos Passos via Mapas de Fluxo de Verossimilhança Escaláveis (SCALLOP)

Declaração do Problema

Avanços recentes em modelagem generativa baseada em fluxo permitiram a geração de amostras de alta qualidade com poucas avaliações de função (NFEs). No entanto, estimar a verossimilhança do modelo (densidade da amostra) continua sendo um gargalo. As abordagens existentes enfrentam um compromisso:

1. Arquiteturas Restritivas: Fluxos de normalização garantem o cálculo exato da verossimilhança via fórmula de mudança de variáveis, mas exigem restrições arquiteturais (ex: camadas invertíveis) que limitam a expressividade, particularmente em domínios com simetrias, como a modelagem molecular.
2. Estimadores de Alta Variância: Fluxos de Normalização Contínua (CNFs) e seus contrapartes destilados (ex: Mapas de Fluxo) oferecem maior expressividade, mas requerem a integração da divergência do campo de velocidade. Isso é tipicamente estimado usando o estimador de traço de Hutchinson, o que introduz uma variância substancial e exige orçamentos computacionais que escalam mal com a dimensionalidade.
3. Limitações de Destilação: Métodos recentes como o F2D2 (Ai et al., 2026) tentam destilar tanto a amostragem quanto a verossimilhança em poucos passos, mas ainda dependem de estimadores de divergência baseados em Hutchinson durante o treinamento, herdando sua alta variância e sobrecarga computacional.

Consequentemente, há uma falta de métodos escaláveis e de baixa variância para estimar verossimilhanças para modelos generativos expressivos de poucos passos, dificultando sua aplicação como geradores de Boltzmann para estimativa de energia livre e amostragem aprimorada em domínios científicos.

Metodologia: SCALLOP

Os autores propõem o SCAlable LikeLihood distillation of flOw maPs (SCALLOP), um framework que se baseia no modelo de mapa de fluxo de verossimilhança F2D2, mas substitui o objetivo de treinamento baseado em Hutchinson por uma alternativa vetorizada e escalável.

1. Mapas de Fluxo de Verossimilhança

O SCALLOP utiliza um mapa de fluxo de verossimilhança $\Psi$, um operador que mapeia conjuntamente um estado $x_t$ no tempo $t$ para um estado $x_s$ no tempo $s$ e a mudança de log-densidade associada. Ele é parametrizado como:
$$\Psi_\theta(x_t, t, s) = \begin{bmatrix} x_t \\ 0 \end{bmatrix} + (s-t) \begin{bmatrix} u_\theta(x_t, t, s) \\ D_\theta(x_t, t, s) \end{bmatrix}$$
onde $u_\theta$ prevê a velocidade e $D_\theta$ prevê a divergência (mudança de log-densidade).

2. Correspondência de Divergência Condicional (CDM)

A inovação central é substituir o estimador de traço de Hutchinson pela Correspondência de Divergência Condicional (CDM).

• Derivação: Usando a regra da cadeia no log-densidade ao longo da trajetória do fluxo ($dt \log p_t(x_t) = \nabla \log p_t(x_t) \cdot v_t(x_t) + \partial_t \log p_t(x_t)$), os autores reformulam o termo de divergência.
• Mecanismo: Ao condicionar no ponto de dado alvo $x_1$ (em um cenário de interpolação linear $x_t = (1-t)x_0 + tx_1$), a distribuição condicional $p_t(x_t|x_1)$ torna-se tratável (Gaussiana). Isso permite que a divergência seja expressa sem operadores de divergência explícitos ou estimadores de traço estocásticos.
• Objetivo: O objetivo da CDM minimiza a diferença entre a divergência prevista $D_\theta$ e um termo derivado do campo de velocidade conhecido e dos scores condicionais, contornando a necessidade do estimador de Hutchinson.

3. Correspondência de Divergência Vetorizada

Para melhorar ainda mais a eficiência, os autores introduzem a Correspondência de Divergência Vetorizada (CDM-v).

• Em vez de prever a divergência escalar $\nabla \cdot v_t$, a rede prevê o vetor das entradas diagonais do Jacobiano, $\text{diag}(\nabla v_t)$.
• Isso transforma o alvo de aprendizado de um escalar para um sinal vetorial, que é mais eficiente em termos de dados e permite uma formulação totalmente vetorizada.
• A análise teórica mostra que minimizar a perda vetorizada é equivalente a minimizar a perda de correspondência de divergência padrão até uma constante.

4. Treinamento e Arquitetura

• Arquitetura: O SCALLOP emprega um backbone de Transformer de Difusão (DiT) com camadas de codificação compartilhadas e cabeças de leitura separadas para a velocidade da amostra e a mudança de log-densidade vetorizada.
• Simetria: Para sistemas moleculares, a arquitetura impõe equivariância SE(3) para a velocidade e invariância SE(3) para a log-densidade (via soma do output vetorial).
• Estabilização: Para lidar com singularidades em $t=1$, o treinamento amostra $t$ de $[0, 1-\epsilon]$ e aplica ponderação linear ao termo de correspondência de divergência.

Principais Contribuições

1. Objetivo de Treinamento Escalável: O SCALLOP introduz um objetivo de destilação de verossimilhança vetorizado e livre de Hutchinson (CDM-v) que reduz significativamente a variância do treinamento e o custo computacional em comparação ao F2D2.
2. Superioridade Empírica: O SCALLOP demonstra melhorias substanciais sobre o F2D2, reduzindo a variância do treinamento em até 100$\times$, convergindo mais rápido e alcançando um melhor desempenho no conjunto de teste.
3. Eficiência de Inferência: Em domínios moleculares, o SCALLOP alcança até 10$\times$ mais rapidez na inferência que fluxos de normalização e 100$\times$ mais rápido que o FALCON, mantendo uma qualidade de amostra competitiva ou superior.
4. Generalização: O método é validado em conjuntos de dados de imagens de alta dimensão (CelebA-64), mostrando estimativa de verossimilhança precisa e dinâmica de treinamento estável em comparação ao F2D2.

Resultados

Sistemas Moleculares (Geradores de Boltzmann)

Avaliado em sistemas de Alanina Dipeptídeo (ALA-2) até Hexapeptídeo (ALA-6):

• Desempenho: O SCALLOP supera consistentemente o F2D2 em todas as métricas (Tamanho de Amostra Efetiva, Wasserstein-2 de Energia, Wasserstein-2 de Ângulo Torsional). Ele atinge um desempenho competitivo com o estado da arte FALCON e FALCON-A.
• Estabilidade: Os gradientes de treinamento para o SCALLOP exibem variância significativamente menor e uma Razão Sinal-Ruído (SNR) mais alta em comparação ao F2D2.
• Velocidade: O SCALLOP suporta grandes lotes (batch sizes) de $10^4$ durante a inferência, enquanto o FALCON exige lotes fortemente restritos devido aos custos de computação do Jacobiano. O SCALLOP é 10$\times$ mais rápido que fluxos de normalização e 100$\times$ mais rápido que o FALCON.

Geração de Imagens (CelebA-64)

• O SCALLOP alcança Bits Por Dimensão (BPD) e Fréchet Inception Distance (FID) comparáveis ao F2D2 de 1 a 8 passos de Euler.
• Ele reduz consistentemente o erro de BPD por amostra em comparação ao F2D2.
• O cálculo da perda é 16,9% mais rápido (83,1% do tempo de execução do F2D2) com dinâmica de treinamento mais estável.

Significância e Alegações

O artigo afirma que o SCALLOP resolve o antigo compromisso entre expressividade e avaliação de verossimilhança tratável em modelos generativos de poucos passos. Ao eliminar o estimador de traço de Hutchinson, o SCALLOP permite o treinamento de mapas de fluxo de verossimilhança altamente expressivos com baixa variância e alta escalabilidade.

Os autores posicionam o SCALLOP como uma solução prática para aplicações científicas que exigem geradores de Boltzmann, onde a estimativa precisa de verossimilhança é crítica para cálculos de energia livre e amostragem aprimorada. Eles observam que, embora o SCALLOP possa teoricamente gerar amostras em 1 passo, ele utiliza na prática cerca de 10 passos para tarefas de alta dimensão. O método depende da premissa de que o campo de velocidade é bem aprendido, e a lacuna teórica referente ao viés introduzido por essa premissa permanece não quantificada. No entanto, os resultados empíricos demonstram que o SCALLOP é uma abordagem competitiva e de estado da arte para acelerar tanto o treinamento quanto a inferência para modelagem generativa consciente de verossimilhança.