Resumo Técnico: Resíduos de Wasserstein: Aprendendo Fluxos de Gradiente a partir de Dinâmicas Populacionais

Declaração do Problema

O artigo aborda o problema inverso de reconstruir dinâmicas populacionais a partir de um conjunto esparso de instantâneos observacionais. Em muitos campos científicos, como a biologia computacional e a dinâmica de multidões, a evolução populacional é modelada como um Fluxo de Gradiente de Wasserstein (WGF): uma curva de distribuições de probabilidade $\rho_t$ impulsionada pelo descida mais íngreme de um funcional de energia $F$. O objetivo é recuperar o funcional de energia subjacente $F$ de modo que seu fluxo de gradiente se ajuste às distribuições marginais observadas $\{q_t\}$ em tempos de observação discretos $t \in \mathcal{T}_{obs}$.

As abordagens dominantes existentes dependem do esquema Jordan–Kinderlehrer–Otto (JKO), que define WGFs como uma sequência de passos proximais minimizando uma energia mais uma penalidade de transporte ótimo (OT). No entanto, os métodos baseados em JKO sofrem de duas limitações primárias:

1. Inflexibilidade à Discretização Temporal: Eles exigem a resolução de problemas de OT custosos entre instantâneos consecutivos.
2. Sensibilidade a Lacunas: Eles têm dificuldade com grandes lacunas temporais entre observações, frequentemente falhando em capturar trajetórias curvas (por exemplo, partículas movendo-se ao longo de um vale sinusoidal) porque dependem de interpolação por cordas lineares entre os instantâneos.

Metodologia: A Estrutura de Resíduo de Wasserstein

Os autores propõem uma abordagem baseada em resíduos que contorna o esquema JKO e os acoplamentos de transporte ótimo. Em vez de impor a condição JKO, eles impõem as equações diferenciais fundamentais que caracterizam os fluxos de gradiente diretamente via uma função de perda não negativa.

1. A Formulação de Resíduo

O artigo utiliza as formulações de Tangente e de Igualdade de Dissipação de Energia (EDE) de fluxos de gradiente. Especificamente, foca na restrição de velocidade:
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
onde $v_t$ é o campo de velocidade e $\frac{\delta F}{\delta \rho_t}$ é a primeira variação do funcional de energia.

Os autores definem um resíduo de velocidade $R_{vel}$ que é zero se, e somente se, $\rho$ é um fluxo de gradiente de $F$:
$$R_{vel}[F, (\rho, v)] = \int_0^T \int_{\mathbb{R}^D} \left\| v_t(x) + \nabla_x \frac{\delta F}{\delta \rho_t}(x) \right\|^2 \rho_t(x) \, dx \, dt$$

O objetivo global combina este resíduo com uma divergência de ajuste de dados $D$ nos tempos observados:
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
Esta estrutura unifica métodos existentes como Path-Finding e Action Matching como instanciacões de um único princípio de resíduo.

2. O Algoritmo "Stitching"

Para minimizar este objetivo, os autores introduzem o Stitching, um método baseado em partículas e livre de simulação (simulation-free), com as seguintes características:

• Trajetórias Aprendíveis: Ao contrário dos métodos JKO onde a curva $\rho$ é implícita, o Stitching parametriza explicitamente a curva $\rho_\theta$ como uma Estimativa de Densidade de Kernel (KDE) sobre um conjunto de $N$ trajetórias de partículas aprendíveis $x_{t,k}^\theta$.
• Livre de Simulação: O método não requer a resolução de EDOs (Equações Diferenciais Ordinárias) ou a integração de EDEs (Equações Diferenciais Estocásticas). As posições e velocidades das partículas são parâmetros aprendíveis diretos.
• Aproximação Empírica: Para tornar a computação eficiente, a KDE é aproximada por uma medida empírica nos centros das partículas. Isso simplifica o resíduo de velocidade para uma soma sobre as partículas:
  \mathcal{L}_{stitch} \approx \int_0^T \sum_{k=1}^N w_k \left\| \dot{x}_{t,k} + \nabla_x \frac{\delta F_\theta}{\delta \rho_\theta_t}(x_{t,k}) \right\|^2 dt + \sum_{t \in \mathcal{T}_{obs}} D(\rho_\theta_t, q_t)
• Parametrização Funcional: O funcional de energia $F_\theta$ é parametrizado como uma soma de um termo de potencial $V_\theta$, um termo de entropia (densidade logarítmica) e um termo de interação $W_\theta$, todos representados por redes neurais.

Principais Contribuições

1. Estrutura de Resíduo: Os autores reestruturam o aprendizado de WGF como a minimização de um termo de resíduo não negativo, fornecendo uma perspectiva unificada que engloba o Path-Finding e o Action Matching.
2. Algoritmo Stitching: Eles introduzem um novo método de partículas livre de simulação que trata a curva de trajetória $\rho$ como uma variável aprendível de primeira classe. Isso permite que o método tolere grandes lacunas entre instantâneos e recupere dinâmicas curvas que os preditores de corda baseados em JKO perdem.
3. Desempenho de Estado da Arte: O método alcança resultados superiores em benchmarks de inferência de trajetória, particularmente em cenários com dados esparsos ou não pareados.

Resultados Experimentais

O artigo avalia o Stitching em três principais benchmarks:

1. Exemplo Ilustrativo (Fluxo Contínuo vs. JKO)

Em um potencial sintético de "vale ondulado", as partículas seguem uma trajetória sinusoidal.

• Resultado: Métodos baseados em JKO (especificamente JKOnet*) falham em capturar a curvatura entre instantâneos esparsos ($t=10$ a $t=20$), produzindo interpolações de linha reta. O Stitching recupera corretamente a trajetória curva, levando a uma recuperação significativamente melhor do potencial subjacente ($R^2 = 0,62$ vs. $0,51$).

2. Recuperação de Potencial Sintético

O método é testado em 15 potenciais bidimensionais sob regimes pareados (trajetórias correlacionadas) e não pareados (instantâneos decorrelacionados).

• Resultado: O Stitching mantém um desempenho consistente em ambos os regimes. Em contraste, os métodos baseados em JKO degradam-se significativamente ou colapsam no cenário não pareado, pois dependem de acoplamentos de OT entre instantâneos consecutivos. O erro do Stitching permanece dentro de um fator de $\sim 2\times$ entre os regimes, enquanto os métodos JKO frequentemente falham em recuperar campos de gradiente informativos no caso não pareado.

3. Inferência de Trajetória de Célula Única (Dataset de Corpo Embrióide)

Aplicado a dados de diferenciação de células-tronco embrionárias humanas (5 instantâneos ao longo de 27 dias).

• Resultado: O Stitching alcança o melhor desempenho (menores distâncias $W_1$ e $W_2$) em comparação com baselines incluindo Neural SDEs, TrajectoryNet e vários métodos baseados em JKO. Ele supera os competidores tanto no treinamento com dados completos quanto em um protocolo de interpolação temporal "leave-two-out".

4. Recuperação de Dinâmicas de Interação

O método é usado para desmembrar potenciais ambientais ($V$) e kernels de interação ($W$) em um sistema de McKean–Vlasov.

• Resultado: O Stitching recupera com sucesso tanto o potencial de confinamento quanto o kernel de interação atrativa, alcançando altos scores de padrão $R^2$ ($0,84$ para $V$, $0,81$ para $W$). É também significativamente mais rápido por época do que os métodos baseados em JKO, que exigem a resolução de problemas de OT em cada passo.

5. Fluxos Não-Gradiente

Os autores demonstram que a estrutura de resíduo pode ser estendida além dos fluxos de gradiente ao substituir o termo de gradiente por um campo vetorial geral, rastreando com sucesso dinâmicas quirais (não conservativas).

Significância e Alegações

O artigo afirma que a perspectiva de Resíduo de Wasserstein oferece uma alternativa mais flexível e robusta ao esquema JKO para aprender dinâmicas populacionais. Ao desacoplar o aprendizado do funcional de energia das restrições de acoplamento de transporte ótimo, o Stitching permite:

• Robustez à Esparsidade: A capacidade de aprender dinâmicas a partir de observações amplamente espaçadas sem a "perda de curvatura" inerente à interpolação linear.
• Compatibilidade com Dados Não Pareados: A capacidade de aprender de conjuntos de dados onde as identidades das partículas não são rastreadas através do tempo (instantâneos não pareados), um cenário comum em imagens biológicas.
• Eficiência Livre de Simulação: Evitar o custo computacional de resolver EDOs ou problemas de OT durante o treinamento.

Os autores posicionam este trabalho como uma unificação de paradigmas existentes e um avanço prático para a descoberta científica em campos que vão desde a biologia de célula única até o comportamento coletivo, reconhecendo limitações como a complexidade computacional $O(N^2)$ para termos de interação e a falta de garantias teóricas de convergência conforme $N \to \infty$.