Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

Resumo Técnico: Fluxo de Correspondência Multi-Marginal com Interpolantes Aprendidos Adversariamente

Definição do Problema
O artigo aborda o desafio de aprender a dinâmica dependente do tempo de um sistema dados apenas observações amostradas (instantâneos ou snapshots) em pontos de tempo discretos, sem acesso a trajetórias de verdade fundamental (ground-truth). Esta é uma tarefa central em campos como a biologia (ex: sequenciamento de RNA de célula única, rastreamento celular) e a medicina. Formalmente, o objetivo é inferir um campo vetorial dependente do tempo $v_t$ que governa uma Equação Diferencial Ordinária (ODE) $dx_t = v_t(x_t)dt$, de modo que o mapa de integração realize o transporte da distribuição marginal inicial $q_0$ para uma marginal alvo $q_1$, enquanto satisfaz restrições marginais intermediárias $p_{t_i} = q_{t_i}$ para um conjunto de passos de tempo $0 = t_1 < \dots < t_K = 1$.

Abordagens existentes, como a aplicação do Flow Matching (FM) padrão a pares consecutivos de marginais ou o uso de interpolantes lineares por partes ou splines cúbicas, frequentemente resultam em curvas de interpolação não suaves. Esses problemas de não suavidade levam a uma alta variância de gradiente no objetivo de aprendizagem e dificuldades na marginalização dos campos vetoriais resultantes, particularmente em configurações de alta dimensão ou quando os dados são ruidosos.

Metodologia: ALI-CFM
Os autores propõem o ALI-CFM (Adversarially Learnt Interpolants with Conditional Flow Matching), um método de dois estágios que desacopla a aprendizagem dos interpolantes da aprendizagem do campo vetorial.

1. Aprendizagem Adversária de Interpolantes (ALI):
   Em vez de forçar os interpolantes a passar por pontos de dados específicos (correspondência ponto a ponto), o ALI aprende um interpolante neural $G_\phi(x_0, x_1, t)$ que corresponde à distribuição das marginais intermediárias. O interpolante é parametrizado como uma interpolação linear mais um termo de correção neural:
   $$G_\phi(x_0, x_1, t) = (1-t)x_0 + tx_1 + t(1-t)f_\phi(x_0, x_1, t)$$
   Para garantir que a distribuição de transporte (pushforward) do interpolante corresponda às marginais intermediárias observadas $q_{t_i}$, os autores empregam um objetivo adversarial inspirado em GANs. Um discriminador $D_\gamma$ distingue entre amostras da marginal verdadeira $q_{t_i}$ e amostras geradas pelo interpolante $G_\phi$. O gerador (interpolante) é treinado para minimizar a divergência de Jensen-Shannon entre essas distribuições.

   Para garantir a unicidade dos interpolantes e evitar caminhos arbitrariamente curvos, os autores introduzem regularizadores:

   • Referência Linear: Penaliza o desvio de uma linha reta entre os pontos finais acoplados.
   • Referência Linear por Partes: Regride o interpolante para um caminho linear por partes definido por acoplamentos de Transporte Ótimo (OT) de cadeia de Markov.
   • Norma da Segunda Derivada: Impõe suavidade ao penalizar a curvatura do interpolante.

2. Marginalização via Conditional Flow Matching (CFM):
   Uma vez treinados os interpolantes $G_\phi$, eles são fixados. Uma segunda rede neural $u_\theta$ é treinada para aproximar o campo vetorial que gera esses interpolantes. Isso é feito usando o objetivo padrão de Conditional Flow Matching (CFM), onde o campo vetorial alvo é a derivada temporal do interpolante aprendido:
   $$L_{CFM} = \mathbb{E}_{t, (x_0, x_1) \sim \pi} \| u_\theta(G_\phi(x_0, x_1, t)) - \frac{d}{dt}G_\phi(x_0, x_1, t) \|^2$$
   O acoplamento $\pi$ entre $q_0$ e $q_1$ pode ser independente ou um plano de Transporte Ótimo.

Principais Contribuições

1. Nova Técnica de Interpolação: O artigo introduz um método para aprender interpolantes usando um objetivo adversarial que corresponde a distribuições em vez de pontos. Isso permite interpolantes suaves que são únicos sob suposições leves (provadas nos apêndices).
2. Robustez a Ruído e Complexidade: O método foi projetado para lidar com conjuntos de dados ruidosos e geometrias complexas e variantes no tempo, onde métodos de interpolação ponto a ponto (como splines ou lineares por partes) falham ou tornam-se instáveis.
3. Escalabilidade e Versatilidade: Os autores demonstram que o ALI-CFM escala para dados de alta dimensão e um grande número de distribuições marginais ($K$), superando baselines existentes em cenários complexos específicos, mantendo-se competitivo em outros.

Resultados Experimentais
Os autores avaliam o ALI-CFM em quatro tarefas distintas:

• Distribuição Sintética "Knot" (Nó): Em um conjunto de dados sintético com 1.200 marginais formando um nó, o ALI-CFM foi o único método capaz de capturar com precisão a geometria dependente do tempo. Métodos concorrentes (OT-MFM, splines cúbicas, linear por partes) produziram interpolantes não suaves com alta variância de gradiente.
• Rastreamento Celular (Glioblastoma): Usando dados de microscopia de uma célula em movimento, o ALI-CFM aprendeu com sucesso a dinâmica espaço-temporal a partir de dados ruidosos e subamostrados. Em contraste, as trajetórias do OT-CFM divergiram, e o OT-MFM independente do tempo falhou em capturar dependências temporais.
• Inferência de Trajetória de Célula Única (scRNA-seq): Nos conjuntos de dados Embryoid Body (EB), Cite-seq e Multiome, o ALI-CFM teve desempenho comparável aos baselines de última geração (OT-MFM, OT-CFM) em termos de Distância de Earth Mover (EMD). Os autores observam que, embora o treinamento adversarial torne difícil superar estritamente métodos que sofrem overfitting a dados ponto a ponto nessas métricas específicas, o método permanece altamente competitivo.
• Inferência de Coordenadas de Tumor (Transcriptômica Espacial): Em um conjunto de dados de câncer de mama com seções de tecido espacialmente ordenadas, o ALI-CFM superou significativamente todos os baselines (incluindo OT-CFM, OT-MMFM e OT-MFM) na previsão da distribuição de regiões tumorais em seções não utilizadas no treinamento (held-out). Os autores atribuem esse sucesso à capacidade do método de lidar com distribuições multimodais e aos desafios de alinhar seções de tecido.

Significância e Alegações
O artigo afirma que o ALI-CFM oferece uma abordagem conceitualmente nova para o flow matching multi-marginal ao mudar de uma atribuição pontual estrita para a correspondência de distribuições. Essa mudança permite a aprendizagem de interpolantes suaves e únicos que são robustos ao ruído e capazes de modelar dinâmicas complexas onde métodos anteriores falham. Os autores destacam que seu método é particularmente eficaz em cenários com centenas ou milhares de distribuições marginais e dados ruidosos, um regime onde nenhum outro algoritmo de flow matching demonstrou desempenho comparável. Embora o método tenha desempenho equivalente aos atuais estados da arte em benchmarks padrão de célula única, sua performance superior em transcriptômica espacial e tarefas de geometria sintética complexa ressalta sua versatilidade e escalabilidade.