FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

Resumo Técnico: FreeBridge – Pontes de Schrödinger Variacionais para Dinâmicas de Transição Celular

Definição do Problema
Ensaios de imagem de alto conteúdo são essenciais para quantificar respostas celulares a perturbações químicas e genéticas. No entanto, existe uma limitação fundamental: as células estão quimicamente fixadas no momento da aquisição da imagem, o que torna as trajetórias temporais individuais inobserváveis. Consequentemente, a modelagem de perturbação é reduzida à inferência de transporte estocástico entre populações de controle e tratadas, onde apenas as distribuições marginais dos pontos finais são conhecidas.

Embora modelos generativos recentes (ex: modelos de difusão, flow matching) alcancem um forte alinhamento entre os pontos finais de controle e perturbados, a consistência do ponto final por si só é insuficiente para determinar a evolução intermediária. Infinitos processos estocásticos podem conectar marginais idênticas; sem restrições adicionais, esses modelos podem atravessar regiões no espaço latente não suportadas por morfologias de células individuais observadas, levando a estados intermediários biologicamente implausíveis. Métodos existentes frequentemente otimizam objetivos de nível de distribuição que deixam as estruturas celulares intermediárias subdeterminadas, podendo satisfazer métricas como FID ou KID, mas gerando artefatos.

Metodologia: O Framework FreeBridge
Os autores propõem o FreeBridge, um framework de transporte estocástico restrito formulado como um "Motor Celular" (Cell Engine) que separa explicitamente a definição de estados celulares válidos do mecanismo de transição. O método opera sob supervisão de apenas pontos finais e consiste em três componentes principais:

1. Unidade de Estado (Geometria de Célula Única Explícita):

   • Em vez de tratar imagens como misturas de nível populacional, o FreeBridge define estados atômicos como representações de célula única segmentadas por instância (usando Cellpose).
   • Um codificador pré-treinado e fixo $\phi$ incorpora esses crops em um espaço latente $z = \phi(x)$. Este codificador é congelado durante o aprendizado do transporte, estabelecendo uma geometria latente fixa definida antes do aprendizado da regra de transição.
   • Um banco de suporte empírico $Z_{emp}$ é construído a partir das incorporações (embeddings) latentes das células observadas. Este banco define o "manifold" de morfologias biologicamente válidas.

2. Regra de Transição (Transporte de Ponte de Schrödinger):

   • A transição entre a distribuição de controle $\mu_z$ e a distribuição perturbada $\nu_{c,z}$ é modelada como uma equação diferencial estocástica (SDE) controlada: $dZ_t = u_t(Z_t) dt + \sigma dW_t$.
   • O objetivo é encontrar um campo de deriva $u_t$ que transporte a distribuição inicial para a distribuição alvo em $t=1$.
   • Diferente do transporte ótimo padrão, esta formulação incorpora um custo de estado específico da tarefa para regularizar o caminho.

3. Regularização Geométrica (Restrição de Suporte Empírico):

   • Para evitar que o processo estocástico atravesse "vazios" no espaço latente, o FreeBridge introduz um custo de suporte $V(z) = \lambda \cdot d_{emp}(z)^2$, onde $d_{emp}(z)$ é a distância ao vizinho mais próximo no banco de suporte empírico $Z_{emp}$.
   • O objetivo total minimiza a soma da energia cinética (magnitude da deriva) e o custo de suporte:
     $$\min_u \mathbb{E} \left[ \int_0^1 \frac{1}{2\sigma^2}\|u_t(Z_t)\|^2 dt + \lambda \int_0^1 d_{emp}(Z_t)^2 dt \right]$$
   • Esta regularização ancora os estados intermediários ao manifold de morfologia empiricamente observado sem exigir trajetórias pareadas ao nível da célula.

Principais Contribuições

• Formulação: Introduz uma formulação de Ponte de Schrödinger para modelagem de transição de célula única que desacopla explicitamente a especificação de estado do transporte estocástico.
• Fundamentação Geométrica: Propõe um termo de regularização inovador baseado em estimativa de densidade não paramétrica (distância do vizinho mais próximo ao banco de suporte empírico) para restringir as trajetórias intermediárias dentro de regiões morfologicamente plausíveis.
• Protocolo Unificado: Estabelece um protocolo de avaliação unificado através de múltiplos conjuntos de dados (BBBC021, RxRx1, JUMP) usando crops de célula única segmentados por instância, garantindo comparação justa entre métodos.
• Desempenho: Demonstra que o constrangimento da geometria melhora não apenas a fidelidade do ponto final, mas também a retenção semântica dos mecanismos de perturbação (Mecanismo de Ação - MoA).

Resultados
Avaliado nos conjuntos de dados BBBC021, RxRx1 e JUMP, o FreeBridge demonstra o seguinte:

• Fidelidade do Ponto Final: Alcança pontuações de Distância de Inception de Fréchet (FID) e Distância de Inception de Kernel (KID) competitivas ou superiores em relação às linhas de base de estado da arte (PhenDiff, IMPA, CellFlux). No BBBC021, alcança o menor FID (15,35) e KID (1,10).
• Consistência Intermediária: Reduz significativamente a Taxa de Violação de Suporte ($R_{viol}$), uma métrica que mede a proporção de estados intermediários que caem fora do suporte morfológico empírico. No BBBC021, o FreeBridge reduz $R_{viol}$ para 0,11, comparado a 0,38 para o PhenDiff e 0,31 para o CellFlux.
• Retenção Semântica: Estudos de ablação mostram que remover a regularização geométrica diminui a precisão de classificação do Mecanismo de Ação (MoA) (de 72,08% para 65,08% no BBBC021), indicando que restringir a evolução intermediária ajuda a preservar sinais mecanísticos.
• Análise Qualitativa: Visualizações confirmam que o FreeBridge preserva melhor as características morfológicas específicas de compostos (ex: compactação nuclear, retração citoplasmática) do que as linhas de base, que frequentemente exibem suavização de textura ou padrões incompletos.

Significância e Alegações
O artigo afirma que modelar transições biológicas requer mais do que apenas o alinhamento de fronteiras; necessita da fundamentação geométrica explícita da evolução intermediária. Ao ancorar o transporte estocástico a um manifold celular empiricamente suportado, o FreeBridge fornece uma abordagem principista para gerar dinâmicas de perturbação biologicamente interpretáveis. Os autores afirmam que esta formulação oferece uma base robusta para a simulação de "células virtuais", abordando a lacuna crítica onde modelos alinhados ao ponto final falham em capturar estados celulares intermediários realistas. O trabalho destaca que, sem restrições geométricas, modelos generativos podem produzir pontos finais estatisticamente válidos que atravessam regiões do espaço latente sem suporte biológico.

Limitações
Os autores reconhecem que o FreeBridge depende de um codificador congelado; fatores morfológicos ausentes no pré-treinamento do codificador não podem ser recuperados apenas através do aprendizado de transporte. Além disso, o cálculo da distância de suporte escala quadraticamente ($O(B^2)$) com o tamanho do minibatch, o que pode exigir busca de vizinho mais próximo aproximada para bancos de suporte maiores.