Resumo Técnico: Modelos de Transporte de Beckmann

Declaração do Problema

A modelagem generativa padrão via flow matching e modelos de difusão depende de campos de velocidade dependentes do tempo (drifts) para transportar amostras de uma distribuição base $\mu_0$ para uma distribuição alvo $\mu_1$. Embora eficaz, essa abordagem requer a integração de uma equação diferencial ordinária (ODE) dependente do tempo ao longo de um intervalo fixo $[0, 1]$, o que é computacionalmente caro durante a inferência.

Uma abordagem alternativa, proposta por Wang e Du [2025] como "Equilibrium Matching" (EM), restringe o drift a ser independente do tempo (autônomo). Isso simplifica a arquitetura e o sistema dinâmico. No entanto, a formulação original de EM carecia de uma prova rigorosa de que um fluxo autônomo com um drift independente do tempo poderia transportar exatamente $\mu_0$ para $\mu_1$. Além disso, a função de perda proposta na formulação original de EM mostrou introduzir vieses no transporte, falhando em impor a alocação correta de massa, particularmente quando a distribuição alvo está suportada em uma variedade (manifold) de dimensão inferior.

Metodologia

O artigo introduz os Modelos de Transporte de Beckmann (BTM), um framework que instancia o flow matching usando um campo de velocidade independente do tempo $b: \mathbb{R}^d \to \mathbb{R}^d$. A principal percepção teórica é que, se a distribuição alvo $\mu_1$ for singular (suportada em uma variedade de dimensão inferior $M_1$), um drift independente do tempo pode definir um mapa de transporte válido.

1. Fluxo Autônomo e Condição de Divergência

Os autores estabelecem que um drift independente do tempo $b$ define um transporte válido se satisfizer uma equação de divergência específica envolvendo uma função de peso positiva $\nu$:
$$\nabla \cdot (\nu b) = \mu_0 - \mu_1$$
Aqui, $\mu_0$ atua como a fonte e $\mu_1$ como o sumidouro. O peso $\nu$ representa a distribuição do interpolante média no tempo na configuração padrão de flow matching.

• Convergência: Sob esta condição, as trajetórias da ODE autônoma $\dot{X}_t = b(X_t)$ começando de $\mu_0$ convergem para o suporte $M_1$ de $\mu_1$.
• Mapa de Transporte: O tempo de chegada $\tau(x_0)$ (que pode ser finito ou infinito dependendo do comportamento do drift em $M_1$) define um mapa $T(x_0) = \lim_{t \to \tau} X_t(x_0)$ que leva $\mu_0$ adiante para $\mu_1$ ($T_\sharp \mu_0 = \mu_1$).

2. Correção do Equilibrium Matching

O artigo identifica que a perda original de Equilibrium Matching minimiza um objetivo de regressão que não satisfaz estritamente a condição de divergência necessária para o transporte exato. Ao tratar o cronograma (schedule) do interpolante como parte do próprio interpolante (em vez de um fator de escala separado no alvo), os autores derivam uma perda consistente (Equação 5) que garante que o drift aprendido satisfaça a condição de divergência. Esta correção elimina o viés de alocação de massa observado no método EM original.

3. Mapa de Passo Único via Equação de Conservação

Uma contribuição fundamental é a caracterização do mapa de transporte $T$ não como o resultado da integração de uma ODE, mas como a solução única de uma equação de conservação ao longo das trajetórias de fluxo:
$$b(x) \cdot \nabla T(x) = 0 \quad \text{para } x \notin M_1$$
com a condição de contorno $T(x) = x$ para $x \in M_1$.
Esta equação implica que o valor do mapa $T$ é constante ao longo das linhas de fluxo de $b$. Esta percepção estrutural permite a aprendizagem direta de $T$ sem resolver a ODE durante a inferência.

4. Objetivo de Aprendizagem Direta

Para aprender $T$ diretamente de amostras, os autores discretizam a equação Euleriana que governa o mapa de fluxo. Usando um passo de Euler progressivo, eles derivam um objetivo de regressão (Equação 15) que treina uma rede $T_\theta$ para prever o próximo estado na sequência de fluxo:
$$\mathcal{L}(T) = \mathbb{E}_{t, x_0, x_1} \left[ \left| T(I_t) - \text{sg}\left( T(I_t) + \dot{I}_t \cdot \nabla T(I_t) \right) \right|^2 \right] + \lambda \mathbb{E}_{x_1} \left[ |T(x_1) - x_1|^2 \right]$$
onde $\text{sg}$ denota o operador stop-gradient. Isso permite a geração em um passo ($x_1 = T(x_0)$) ou o refinamento iterativo ($x_1 = T^k(x_0)$) sem integração numérica de ODE.

Principais Contribuições

1. Justificativa Teórica: Prova que os drifts de flow matching dependentes do tempo definem transportes válidos para alvos singulares, recuperando o Equilibrium Matching mediante uma correção de perda.
2. Framework Geral: Mostra que qualquer par $(\nu, b)$ que satisfaça a equação de divergência $\nabla \cdot (\nu b) = \mu_0 - \mu_1$ produz um transporte autônomo válido. Isso unifica o Flow Matching (onde $\nu$ é o interpolante médio no tempo) e o modelo generativo de fluxo de Poisson livre de treinamento (onde $\nu \equiv 1$).
3. Equação de Conservação: Deriva uma equação de conservação única ($b \cdot \nabla T = 0$) que caracteriza o mapa de transporte, fornecendo uma interpretação dinâmica da restrição de fluxo no problema de transporte de Beckmann.
4. Aprendizagem Direta do Mapa: Propõe um método para aprender o mapa de transporte $T$ diretamente de amostras usando uma perda residual derivada da equação Euleriana do fluxo, contornando a integração de ODE na inferência.
5. Correção de Viés: Demonstra que a perda original de Equilibrium Matching introduz viés na alocação de massa e fornece um interpolante consistente e corrigido que resolve este problema sem custo adicional.

Resultados Experimentais

Os autores validam sua teoria em alvos atômicos 2D e benchmarks de imagem:

• Alvos Atômicos 2D: A perda corrigida do BTM alcança uma alocação de massa quase perfeita (MAE $\approx$ 0,005) entre os modos, enquanto a perda original de Equilibrium Matching mostrou viés significativo (MAE $\approx$ 0,102) devido a estruturas de bacia incorretas.
• Convergência Iterativa: Em um alvo em espiral incorporado em altas dimensões, um mapa $T$ parcialmente treinado produz amostras difusas, mas iterar o mapa (ex: $T^3$) torna a geração mais nítida, convergindo para a trajetória do fluxo autônomo. Um $T$ totalmente treinado coincide com o fluxo em um único passo.
• ImageNet 256x256:
  • Correção: Aplicar a perda consistente à arquitetura de Equilibrium Matching (EqM-XL/2) melhorou o FID de 1,90 para 1,87 sem alterar a arquitetura ou o orçamento de treinamento.
  • Geração em um Passo: Treinar um mapa direto no espaço latente SD-VAE alcançou um FID de 17,58 com uma passagem direta (sem classifier-free guidance). Isso supera outros métodos de um passo como o iCT (FID 30,10) e é competitivo com métodos que exigem orientação, embora permaneça abaixo de modelos de difusão de múltiplos passos.

Significância e Alegações

O artigo afirma fornecer um framework unificador que recupera modelos de forma fechada (Fluxo de Poisson) e métodos recentes baseados em equilíbrio sob um único guarda-chuva teórico. Ao ligar fluxos autônomos ao problema de transporte de Beckmann, os autores fornecem uma interpretação dinâmica para uma restrição de fluxo que anteriormente carecia de uma, distinta da formulação de Benamou–Brenier do transporte ótimo de Monge–Kantorovich.

A significância reside em:

• Rigor Teórico: Fornecer a prova que faltava de que fluxos autônomos podem transportar exatamente para variedades singulares.
• Eficiência Prática: Permitir a geração em um passo ao aprender o mapa diretamente, removendo a necessidade de integração de ODE na inferência.
• Correção de Métodos Existentes: Identificar e corrigir um viés fundamental na função de perda de Equilibrium Matching.
• Flexibilidade de Design: Introduzir o peso $\nu$ como um controle de design para alternar entre diferentes modelos generativos (ex: FM vs. Fluxo de Poisson) e controlar o comportamento do drift próximo à variedade alvo.

Os autores permanecem modestos quanto ao desempenho atual do BTM de um passo no ImageNet comparado aos modelos de difusão de estado da arte com múltiplos passos, observando que o FID de 17,58 é "modesto", mas demonstra a viabilidade da abordagem. Eles identificam a extensão do classifier-free guidance para o BTM e a otimização do compromisso entre tempo de treinamento e iterações de inferência como direções futuras de trabalho.