Resumo Técnico: Geração de Expressão Gênica de Célula Única Escalável com Modelos de Difusão Latente

Declaração do Problema

A modelagem computacional da expressão gênica de célula única é essencial para compreender a heterogeneidade celular, diferenciação e respostas a perturbações. No entanto, gerar perfis de expressão realistas permanece um desafio significativo devido a dois fatores primordiais: a natureza baseada em contagens dos dados de expressão gênica e as complexas dependências latentes entre os genes.

Modelos generativos existentes enfrentam limitações críticas:

1. Violação da Permutabilidade (Exchangeability): Muitos métodos impõem ordenações fixas artificiais de genes ou dependem de subconjuntos restritos de genes altamente variáveis (HVGs). Isso contradiz a realidade biológica de que os perfis de expressão gênica são conjuntos permutáveis onde a ordem dos genes não possui significado intrínseco. Arquiteturas de ordem fixa vinculam as dimensões de entrada a genes específicos, impedindo o modelo de lidar com vocabulários de genes variáveis entre tecidos ou espécies sem retreinamento.
2. Instabilidades Arquiteturais e de Treinamento: Abordagens baseadas em Redes Adversárias Generativas (GANs) sofrem de instabilidades de treinamento e colapso de modo (mode collapse). Enquanto isso, VAEs (Autoencoders Variacionais) padrão frequentemente dependem de arquiteturas rasas ou exigem estruturas latentes hierárquicas complexas para lidar com tamanhos de conjuntos, o que pode ser desnecessário se a parametrização for tratada corretamente.

Metodologia: scLDM

Os autores introduzem o scLDM (Modelo de Difusão Latente Escalável), um modelo generativo profundo totalmente baseado em transformer projetado para respeitar a propriedade fundamental de permutabilidade dos dados de célula única. O framework consiste em duas etapas principais:

1. VAE Baseado em Transformer com Design Invariante/Equivariante à Permutação

O núcleo do scLDM é uma arquitetura VAE que utiliza um Bloco de Atenção Cruzada Multi-cabeça (MCAB) para lidar com a natureza não ordenada dos conjuntos de genes.

• Processamento de Entrada: Para gerenciar a complexidade computacional dada a alta esparsidade dos dados de RNA-seq de célula única (frequentemente >70% de zeros), o modelo filtra as entradas para incluir apenas genes expressos e realiza o preenchimento (padding) do contexto para um comprimento fixo $d$. Isso foca o modelo em sinais biologicamente significativos sem restringir a capacidade do decodificador de modelar contagens zero.
• Encoder (Invariante à Permutação): O encoder processa o conjunto de contagens e IDs de genes para produzir variáveis latentes $Z$ de tamanho fixo. Ele emprega MCAB onde as "queries" são pseudo-entradas aprendíveis (pontos indutores), e as "keys/values" são derivadas dos embeddings dos genes. Este design garante que a representação latente seja invariante à permutação, independentemente da ordem dos genes de entrada.
• Decoder (Equivariante à Permutação): O decodificador reconstrói as contagens de expressão gênica para um conjunto solicitado de IDs de genes. Ele utiliza a mesma arquitetura MCAB, mas modifica as queries para serem embeddings específicos de genes indexados pelos IDs de genes solicitados. Isso garante que o output seja equivariante à permutação, o que significa que permutar os IDs de genes de entrada resulta em uma permutação correspondente dos parâmetros de saída.
• Verossimilhança (Likelihood): O modelo parametriza a verossimilhança condicional usando uma distribuição Binomial Negativa (NB) ou Gaussiana, capturando efetivamente a sobredispersão e a natureza de contagem dos dados.

2. Modelo de Difusão Latente (LDM) com Flow Matching

Em vez de um prior Gaussiano padrão, os autores substituem o prior por um Modelo de Difusão Latente treinado no espaço latente do VAE.

• Arquitetura: O processo de difusão é modelado usando Diffusion Transformers (DiT).
• Objetivo de Treinamento: O modelo é treinado usando Flow Matching com interpolantes lineares, otimizando uma função de perda que aprende um campo vetorial para transportar o ruído para a distribuição de dados.
• Geração Controlada: O framework suporta geração multicondicional (ex: tipo de célula + perturbação) via Orientação Livre de Classificador (Classifier-Free Guidance - CFG). Os autores propõem uma estratégia de condicionamento conjunto onde um único token representa uma combinação única de atributos, contrastando com abordagens de condicionamento aditivo usadas em trabalhos anteriores.

Principais Contribuições

1. Arquitetura Permutável: O artigo propõe um novo VAE totalmente baseado em transformer que utiliza uma arquitetura unificada de MCAB para alcançar o agrupamento (pooling) invariante à permutação no encoder e o desagrupamento (unpooling) equivariante à permutação no decoder. Isso elimina a necessidade de componentes separados de pooling/unpooling ou ordenações fixas de genes.
2. Difusão Latente com Flow Matching: A integração de um LDM baseado em Diffusion Transformer treinado com flow matching e interpolantes lineares permite geração de alta qualidade e amostragem controlada via orientação livre de classificador.
3. Condicionamento Conjunto: Os autores demonstram que uma formulação de condicionamento conjunto para orientação livre de classificador (tratando combinações de atributos como tokens únicos) supera as estratégias de condicionamento aditivo em benchmarks de perturbação de múltiplos atributos.

Resultados Experimentais

Os autores avaliaram o scLDM em conjuntos de dados observacionais (Dentate Gyrus, Tabula Muris, HLCA) e conjuntos de dados perturbacionais (Parse 1M, Replogle), comparando-o com baselines de estado da arte incluindo scVI, CFGen, scDiffusion, CPA e scGPT.

• Desempenho de Reconstrução: O scLDM alcançou métricas de reconstrução superiores (menor Erro de Reconstrução, maior Coeficiente de Correlação de Pearson e menor MSE) em todos os conjuntos de dados observacionais em comparação com scVI e CFGen. Por exemplo, no dataset HLCA, o scLD

• Qualidade Generativa: Em tarefas de geração incondicional e condicional, o scLDM demonstrou melhor correspondência de distribuição (menores distâncias Wasserstein-2 e MMD2) e um trade-off precisão-recall mais equilibrado, evitando o colapso de modo observado em alguns baselines.

• Modelagem de Perturbação: Nos datasets Parse 1M e Replogle, o scLDM mostrou desempenho competitivo ou superior na geração de perfis de expressão gênica para combinações de célula-perturbação não vistas. Superou os baselines em métricas de previsão de perturbação (ex: sobreposição de DEG, correlação de Spearman de log-fold changes).

• Classificação Downstream: Os embeddings aprendidos pelo scLDM-VAE foram avaliados em tarefas de classificação (detecção de infecção por COVID-19 e classificação de tipo celular em Tabula Sapiens 2.0). O modelo scLDM de 270 milhões de parâmetros alcançou pontuações F1 de estado da arte (0.820) na tarefa de COVID-19, superando outros modelos de fundação como TranscriptFormer e scGPT.

• Estudos de Ablação: Experimentos confirmaram que a estratégia de entrada de zero-padding é computacionalmente eficiente e mantém a qualidade de reconstrução. Além disso, foi demonstrado que o condicionamento conjunto é superior ao condicionamento aditivo para geração de múltiplos atributos.

Significância e Alegações

O artigo afirma que o scLDM fornece um framework principiado para aprender modelos permutáveis de dados de célula única, indo além das restrições artificiais de ordenações fixas de genes. Ao combinar uma arquitetura transformer invariante/equivariante à permutação com difusão latente, o modelo alcança desempenho de estado da arte tanto na reconstrução quanto na geração de perfis de expressão de célula única.

Os autores enfatizam que sua abordagem suporta naturalmente o condicionamento flexível e a generalização através de contextos não vistos, permitindo a modelagem fiel de respostas celulares complexas. Eles posicionam este trabalho como um passo em direção a modelos de fundação para dados biológicos permutáveis, sugerindo que o design é aplicável além da transcriptômica para dados multi-ômicos e multi-modais. O trabalho visa acelerar o desenvolvimento terapêutico ao permitir experimentos de perturbação in silico, embora os autores observem que a validação experimental permanece necessária para aplicações clínicas.