LoCoFlow: Local-Context Flow Matching for Single-Cell Dynamics and Fate Inference

Resumo Técnico: LoCoFlow

Declaração do Problema

O sequenciamento de célula única fornece instantâneos de alta resolução de estados celulares em pontos temporais discretos, mas a natureza destrutiva da medição impede a observação de correspondências longitudinais reais entre células individuais. A reconstrução de trajetórias de desenvolvimento contínuas a partir desses instantâneos é um problema mal posto com múltiplas soluções possíveis. Além disso, os métodos existentes frequentemente tratam as células como vetores isolados, falhando em representar adequadamente a geometria do manifold local do espaço de estado celular ou a incerteza inerente na diversificação de linhagem. Embora células com estados transcricionais semelhantes possam exibir tendências de destino divergentes, as abordagens atuais de transporte ótimo e flow-matching tipicamente inferem acoplamentos temporais ao nível do indivíduo celular, carecendo de um framework robusto para modelar a evolução dos vizinhanças locais de células como unidades coesas.

Metodologia: LoCoFlow

Os autores introduzem o LoCoFlow (Local-Context Flow Matching), um framework que integra a geometria de nuvem de pontos local com a incerteza de destino condicional para modelar a dinâmica de célula única. O método opera em duas etapas primárias:

1. Construção de Patches Locais e Correspondência Hierárquica

Em vez de parear células isoladas, o LoCoFlow constrói "patches" sobrepostos (vizinhanças locais) ao redor de células âncoras.

• Definição de Patch: Para dados de expressão, os patches são vizinhanças geodésicas em um grafo de vizinhos mais próximos; para dados espaciais, são vizinhanças de tamanho fixo em coordenadas medidas.
• Transporte Hierárquico: O método emprega um esquema de transporte ótimo (OT) de dois níveis. Primeiro, resolve um problema de transporte ao nível de patch usando a distância de Bures para alinhar centros de patch e covariâncias (capturando tanto localização quanto dispersão). Segundo, dentro dos pares de patches retidos, computa um plano de transporte ao nível de célula regularizado por entropia.
• Atribuição Suave (Soft Assignment): A agregação desses planos gera uma distribuição condicional de células alvo para cada origem, em vez de uma atribuição rígida, preservando múltiplas correspondências potenciais de sucessores.

2. Pipeline de Treinamento de Duas Etapas

• Etapa 1: Transporte de Ponte Regularizado pela Geometria:
  • Usando os conjuntos de endpoints pareados derivados do transporte hierárquico, o modelo aprende "pontes" (caminhos intermediários) entre os patches de origem e alvo.
  • Uma rede neural prevê um campo de velocidade residual $f_\theta$ que, combinado com interpolação linear, garante que o caminho passe exatamente pelos endpoints.
  • Priors Geométricos: A função de perda inclui termos de energia cinética e uma penalidade radial que restringe a deformação relativa (expansão/contração) do patch sem restringir a translação comum. Também penaliza velocidades ortogonais à geometria do manifold local estimada.
• Etapa 2: Destilação de Velocidade Condicional:
  • Como o estado futuro é desconhecido durante a inferência, a ponte (que requer ambos os endpoints) não pode ser usada diretamente. O LoCoFlow destila as velocidades da ponte em um campo de velocidade condicional $V_\omega$.
  • Latentes Hierárquicos: O modelo aprende a gerar códigos de alvo $c_i$ a partir de um contexto de origem $Z_0$ usando uma estrutura latente hierárquica: um latente compartilhado pelo patch $\eta_P$ (capturando resultados ao nível da população) e latentes específicos da célula $\eta_i$ (capturando variação individual).
  • Treinamento: Um codificador de alvo aprende primeiro a mapear os endpoints observados para os códigos $c^\star$. Em seguida, um gerador aprende a produzir esses códigos a partir do contexto de origem e dos latentes. A rede de velocidade é treinada para corresponder às velocidades da ponte condicionadas a esses códigos gerados.
  • Inferência: Na inferência, os códigos latentes são extraídos de um prior (sem observações futuras) e mantidos fixos durante a integração. Isso permite a amostragem de trajetórias alternativas que representam a incerteza de destino.

Principais Contribuições

1. Representação de Contexto Local: Muda a unidade fundamental de modelagem dinâmica de vetores de expressão isolados para estados de população local (nuvens de pontos), modelando explicitamente a evolução da geometria da vizinhança.
2. Transporte Hierárquico: Introduz um mecanismo de transporte ótimo acoplado nos níveis de patch e de célula que respeita tanto o movimento global das populações celulares quanto suas mudanças estruturais internas (ex: heterogeneidade, expansão).
3. Incerteza de Destino Condicional: Desenvolve um processo de destilação de duas etapas que separa as restrições geométricas do caminho da incerteza do resultado, permitendo a amostragem de múltiplas trajetórias de destino plausíveis sem observações futuras.
4. Framework de Análise de Perturbação: Fornece um mecanismo para distinguir entre efeitos de células focais e efeitos de vizinhança (background) em simulações de perturbação gênica, revelando como o contexto local molda as respostas previstas.

Resultados Experimentais

O artigo avalia o LoCoFlow em três conjuntos de dados: diferenciação de corpos embrioides (EB), reprogramação de iPSC, desenvolvimento de C. elegans e regeneração cerebral de Axolote.

• Reconstrução de Dados Retidos (Held-Out): Nos conjuntos de dados EB e iPSC, o LoCoFlow alcança distâncias de Wasserstein ($W_1, W_2$) menores em comparação com baselines como OT-CFM, MFM e PACE na maioria das configurações de alta dimensão (ex: PCA50). No entanto, o artigo observa que o LoCoFlow às vezes apresenta uma MMD (Maximum Mean Discrepancy) mais alta, sugerindo que as previsões podem estar mais concentradas em regiões específicas do espaço de estado celular.
• Predição de Destino (Dados LARRY): Usando dados de linhagem rastreada LARRY, o LoCoFlow reduz o erro bruto médio de Brier de 0.7990 (scDiffEq) para 0.6620, indicando melhor precisão na predição de vetores de probabilidade de destino completos.
• Respostas de Perturbação:
  • Patches Geométricos: Ao perturbar genes em patches definidos geometricamente, as intervenções de apenas background e de todos os membros produzem respostas alinhadas (ex: mudando o destino de monócitos para neutrófilos).
  • Minibatches Aleatórios: Em contraste, modelos que utilizam minibatches montados aleatoriamente mostram respostas opostas de célula focal e de background, destacando a importância do contexto geométrico.
• Interpolação Espacial (Axolote): O LoCoFlow melhora a predição espacial ($W_2$) para a regeneração do cérebro de Axolote comparado ao OT-CFM, embora os autores notem que as distribuições espaciais previstas permanecem visualmente similares ao ponto temporal de origem, sugerindo a necessidade de um alinhamento espacial adicional.

Significância e Alegações

O artigo afirma que o LoCoFlow fornece um framework robusto para estudar como a estrutura da população celular local molda as trajetórias inferidas, enquanto representa explicitamente a incerteza no destino celular. Os autores enfatizam que:

• O Contexto Local é Crítico: Os padrões de resposta distintos observados entre patches geométricos e minibatches aleatórios apoiam a hipótese de que o contexto local é uma representação útil para a inferência probabilística de destino.
• Suposições de Modelagem: Os autores mantêm uma postura modesta, reconhecendo que dados de snapshots não identificam unicamente as histórias celulares. Os componentes geométricos e latentes do LoCoFlow impõem suposições de modelagem em vez de revelar relações de linhagem verdadeiras.
• Hipóteses Testáveis: O framework gera hipóteses testáveis sobre como o contexto local molda as respostas previstas a perturbações, distinguindo entre mudanças intrínsecas da célula e influências ambientais.

O trabalho não pretende alegar que resolve a ambiguidade fundamental da inferência de trajetória, mas oferece um método que se alinha melhor à geometria do manifold dos estados celulares e fornece uma maneira estruturada de quantificar e amostrar a incerteza do destino.