Resumo Técnico: Flow Matching para Dados de Contagem (count-FM)

Declaração do Problema

Dados de contagem de alta dimensão, prevalentes em áreas como sequenciamento de RNA de célula única (scRNA-seq) e análise de trens de impulsos neurais, apresentam desafios únicos para a modelagem generativa. Esses tipos de dados são caracterizados por discreticidade, esparsidade e estruturas de correlação complexas. Abordagens de geração profunda existentes frequentemente lutam neste domínio:

• Abordagens de Estado Categórico: Métodos que tratam contagens como variáveis categóricas (por exemplo, tokens) exigem dimensões de saída que escalam com o intervalo máximo de contagem. Isso leva a ineficiência de parâmetros e trata contagens adjacentes como estados não relacionados, ignorando a geometria inerente do espaço de contagem.
• Transformações Contínuas: Abordagens que mapeiam contagens para espaços contínuos via dequantização ou representações latentes substituem a massa de probabilidade discreta por densidades contínuas, potencialmente borrando a estrutura discreta subjacente.
• Limitações dos Modelos de Contagem Existentes: Embora existam alguns modelos de salto específicos para contagem (por exemplo, Poisson-JUMP), eles frequentemente não são projetados para transporte entre distribuições de contagem arbitrárias ou carecem das propriedades eficientes de treinamento e geração dos métodos modernos baseados em fluxo.

O artigo aborda a necessidade de um framework generativo que opere diretamente no espaço de contagem, respeite a geometria discreta dos dados e permita transporte eficiente entre distribuições de origem e destino arbitrárias, sem a sobrecarga de parâmetros dos modelos de estado categórico.

Metodologia: count-FM

Os autores propõem o count-FM, um framework de flow matching adaptado para dados de contagem baseado em um processo de nascimento-morte de tempo contínuo com saltos unitários locais.

1. Processo de Nascimento-Morte de Tempo Contínuo

O modelo define a transição entre uma distribuição de origem $p_0$ e uma distribuição de destino $p_1$ como um processo de salto de Markov de tempo contínuo (CTMC) no espaço de vetores de contagem $\mathbb{N}_0^d$. O processo evolui via nascimentos unitários locais ($+1$) e mortes ($-1$) para cada coordenada $i$:

• Taxa de Nascimento: $\lambda_{t,i}(x)$
• Taxa de Morte: $\mu_{t,i}(x)$
• Dinâmica: Sobre um pequeno intervalo $h$, a probabilidade de transição do estado $x$ para $x + e_i$ é $h \lambda_{t,i}(x)$, e para $x - e_i$ é $h \mu_{t,i}(x)$. A morte é proibida quando uma contagem é zero.

Esta parametrização é altamente eficiente: em vez de prever uma distribuição categórica completa sobre todos os níveis de contagem possíveis (o que escala com a contagem máxima), o modelo prevê apenas duas taxas por dimensão ($2d$ saídas), independentemente do intervalo de contagem.

2. Ponte Binomial Condicional e Treinamento

Para treinar o modelo, os autores definem uma ponte condicional entre pontos fixos $(x_0, x_1)$. Para cada coordenada $i$, a ponte é construída usando uma distribuição Binomial:
$$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$
Esta ponte garante que a média condicional se mova linearmente de $x_0$ para $x_1$ enquanto mantém estados intermediários com valores inteiros.

Ao impor a preservação de massa (a equação de Kolmogorov para frente) nesta ponte condicional, os autores derivam expressões de forma fechada para as taxas de nascimento e morte condicionais:
$$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$$
onde $(\cdot)^+$ denota retificação positiva.

3. Objetivo de Treinamento

O modelo é treinado minimizando uma divergência Kullback-Leibler (KL) no espaço de caminhos entre o processo de ponte condicional e o processo aprendido pelo modelo. Isso se decompõe em um objetivo local de correspondência de taxas:
$$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$$
onde $\ell(u, v) = v - u \log v$ é um termo KL generalizado. Isso permite o treinamento sem simulação das taxas de transição marginais aprendendo as taxas condicionais identificáveis.

4. Amostragem e Acoplamento

• Amostragem: Amostras são geradas simulando o processo de nascimento-morte aprendido para frente de $t=\epsilon_t$ até $t=1-\epsilon_t$ usando uma discretização de salto local de primeira ordem.
• Acoplamento de Pontos Extremos: O artigo explora duas estratégias de acoplamento para emparelhar amostras de origem e destino:
  • Acoplamento Independente: Emparelhamento aleatório ($p_0(x_0)p_1(x_1)$).
  • Transporte Ótimo (OT) Acoplamento: Emparelhamento de minibatch OT usando um custo de Poisson simétrico. Isso tende a induzir transições com menor curvatura, melhorando a interpretabilidade e a eficiência de amostragem.

5. Geração Condicional

Para tarefas condicionais, o modelo incorpora covariáveis $y$ e utiliza Guia Livre de Classificador (CFG). Durante o treinamento, as variáveis de condicionamento são aleatoriamente descartadas. Na amostragem, as taxas condicionais e incondicionais são combinadas via uma escala de guia $w$ para controlar a força da condição.

Contribuições Principais

1. Framework count-FM: Um framework de flow matching para dados de contagem baseado em dinâmicas de nascimento-morte locais, permitindo transporte diretamente no espaço de contagem com estados intermediários permanecendo com valores de contagem.
2. Esquema de Treinamento Tratável: Um método para aprender taxas de transição marginais via uma ponte binomial condicional, mantendo a eficiência do flow matching enquanto evita a complexidade da modelagem de estado categórico.
3. Eficiência de Parâmetros: Uma formulação que prevê apenas taxas locais de nascimento e morte ($2d$ saídas) em vez de distribuições categóricas completas, tornando-a adequada para dados de alta dimensão com grandes intervalos de contagem.

Resultados Experimentais

1. Simulação (Mistura Gamma-Poisson 2D)

• Desempenho: count-FM e count-FM-OT alcançaram a melhor qualidade de amostra (menor distância 2-Wasserstein e MMD$^2$) comparado a baselines incluindo D3PM, discrete-FM, tauLDR, SEDD e Poisson-JUMP.
• Eficiência: count-FM alcançou esses resultados com substancialmente menos parâmetros treináveis (2.372) comparado a baselines de estado categórico (14.000–67.000+).
• Geometria: Estados intermediários evoluíram suavemente no espaço de contagem. A versão acoplada por OT seguiu um caminho de transição visivelmente mais reto e alcançou qualidade de destino com menos avaliações de função (NFEs).

2. RNA-seq de Célula Única (Dentate Gyrus)

• Geração Incondicional: Em um conjunto de dados de 2.930 células e 13.913 genes, count-FM superou baselines de variáveis latentes (scVI) e modelos de difusão latente (scDiffusion, scLDM) em qualidade de amostra (W2 e MMD$^2$) enquanto usava menos parâmetros (aprox. 9,8M vs. 29M–39M para concorrentes).
• Transporte (P12 para P35): O modelo modelou com sucesso o transporte de desenvolvimento do dia pós-natal 12 para o dia 35. Usando acoplamento OT restrito a linhagens, os caminhos de transporte respeitaram estruturas de desenvolvimento conhecidas (por exemplo, progressão de neuroblasto para célula granular), produzindo distribuições intermediárias interpretáveis que permaneceram no espaço de contagem.

3. Trens de Impulsos Neurais (Hipocampo/Entorrinal)

• Geração Condicional: O modelo foi aplicado para gerar contagens de impulsos condicionadas à posição linearizada e direção de corrida.
• Comparação: count-FM com escala de guia $w=1$ superou um regressor determinístico MLP e uma baseline Poisson MLP.
• Estrutura de Correlação: Diferente do Poisson MLP, que subestimou correlações populacionais, count-FM preservou melhor a estrutura de dependência entre neurônios.
• Compromisso de Guia: $w=1$ forneceu o melhor equilíbrio entre ajuste espacial e calibração distribucional. $w=2$ afiou campos de lugar, mas distorceu a calibração de amplitude.

Significado e Alegações

O artigo alega que count-FM fornece um framework eficaz para modelar dados de contagem de alta dimensão, permitindo geração e transporte diretamente no espaço de contagem. Seu significado principal reside em:

• Fidelidade Geométrica: Ao modelar transporte via saltos unitários locais, o método respeita a geometria discreta dos dados de contagem, evitando os efeitos de borrão de aproximações contínuas.
• Escalabilidade: A eficiência de parâmetros torna-o viável para aplicações de alta dimensão (por exemplo, milhares de genes) onde métodos de estado categórico se tornam computacionalmente proibitivos.
• Interpretabilidade: Os estados intermediários no caminho de transporte são amostras de contagem válidas, permitindo a inspeção de distribuições transitórias significativas em aplicações como biologia do desenvolvimento.
• Versatilidade: O framework suporta geração incondicional, transporte entre distribuições arbitrárias e geração condicional com guia controlável.

Os autores reconhecem limitações, observando que a ponte binomial condicional é relativamente simples e pode não explorar todos os caminhos intermediários plausíveis, e que a atual discretização de primeira ordem poderia ser melhorada com esquemas de simulação mais precisos. No entanto, os resultados empíricos em simulações e conjuntos de dados biomédicos demonstram que count-FM alcança alta qualidade de amostra com maior eficiência de modelagem do que abordagens concorrentes.