BioM-JEPA: joint-embedding prediction of graph-connected gene blocks in single cells

Resumo Técnico: BioM-JEPA

Definição do Problema
O sequenciamento de RNA de célula única (scRNA-seq) fornece observações esparsas e dependentes da profundidade de programas biológicos. Os modelos de aprendizado autossupervisionado existentes frequentemente aprendem através da reconstrução de genes individuais (predição de tokens) ou valores de expressão. Os autores argumentam que essa abordagem possui limitações: uma medição de um único gene é esparsa, ruidosa e biologicamente ambígua quando isolada de seu programa coordenado. Além disso, a otimização bem-sucedida da predição ao nível de gene não garante a emergência de representações celulares robustas e de alta dimensão; os modelos podem convergir enquanto produzem embeddings que são de baixa dimensão e fortemente acoplados a artefatos técnicos, como a profundidade de sequenciamento (contagem de genes detectados). A questão central de design permanece: o que um modelo deve prever para aprender uma representação celular reutilizável e biologicamente significativa?

Metodologia: BioM-JEPA
Os autores propõem o BioM-JEPA, uma Arquitetura de Predição de Embedding Conjunto (JEPA) projetada para células únicas. Diferente das JEPAs tradicionais que predizem regiões de imagens mascaradas ou blocos espaço-temporais de vídeo, o BioM-JEPA prediz representações agregadas de blocos de genes conectados por grafos.

• Unidades de Predição (Blocos Conectados por Grafos): Em vez de mascarar genes individuais, o modelo define blocos de alvo como conjuntos de genes conectados amostrados de um grafo de genes binário. Este grafo é construído combinando associações proteicas de alta confiança (STRING v12) com evidências de coexpressão em todo o transcriptoma derivadas do corpus não rotulado de pré-treinamento.
• Arquitetura Estudante-Professor (Student-Teacher):
  • Estudante (Student): Observa o conjunto complementar de genes em uma célula (contexto) e prediz a representação agregada do bloco de alvo oculto.
  • Professor (Teacher): Recebe o conjunto completo de genes observados e fornece a representação alvo. Crucialmente, o alvo do professor é um estado agregado computado a partir dos genes observados dentro do bloco de alvo, não uma reconstrução de valores de genes individuais.
  • Objetivo de Treinamento: O modelo minimiza a distância entre a representação do bloco prevista pelo estudante e a representação alvo agregada pelo professor. Inclui termos de regularização para garantir variação entre as dimensões do embedding e reduzir a covariância redundante. Nenhuma perda de reconstrução de gene individual é utilizada.
• Atenção Linear: Para lidar com o grande número de genes observados sem o custo computacional quadrático da atenção softmax densa, o modelo emprega atenção linear em toda a rede do estudante, do professor e do preditor. Isso permite que o modelo processe a sequência completa de genes observados com computação crescendo linearmente em relação ao número de genes.
• Dados e Escala: O modelo foi pré-treinado em uma coleção local de 22,1 milhões de células (um subconjunto do recurso de 502 milhões de células do scBaseCount). A representação final foi analisada após 10,24 milhões de apresentações de células, uma exposição de treinamento significativamente menor do que modelos comparáveis como Geneformer, scGPT, scFoundation e CellFM.

Principais Contribuições

1. Predição ao Nível de Bloco: Introdução de blocos de genes conectados por grafos como a unidade fundamental de aprendizado de representação de transcriptômica JEPA, afastando-se de alvos independentes ao nível de gene.
2. Eficiência de Arquitetura: Demonstração de que a atenção linear permite o processamento eficiente de amplos contextos de genes, evitando a necessidade de limites artificiais de tokens ou matrizes de atenção quadráticas.
3. Comparações Controladas: Comparação rigorosa do BioM-JEPA contra três baselines específicos usando um backbone e pipeline de extração compartilhados:
   • Token-IJEPA: Prediz genes individuais (JEPA padrão).
   • Random Block: Utiliza conjuntos de genes amostrados aleatoriamente sem conectividade de grafo.
   • Decoder-Only: Reconstrói valores de expressão a partir de blocos de grafos.

Resultos

• Qualidade da Representação: O BioM-JEPA produziu embeddings com um rank efetivo significativamente maior (subindo de 19 para ~40 durante o treinamento) em comparação ao Token-IJEPA (19) e ao Decoder-only (<10). Isso indica que o modelo reteve variação através de muitos direcionamentos independentes.
• Robustez Técnica: Os embeddings do BioM-JEPA mostraram a associação mais fraca com a profundidade de genes detectados (profundidade de sequenciamento) em comparação com todos os controles, sugerindo que o modelo distinguiu com sucesso o sinal biológico do ruído de amostragem estocástica.
• Desempenho em Tarefas Downstream:
  • Anotação: Alcançou o maior score médio de macro-F1 Top-5 (0,835) em tarefas de anotação de tipo celular few-shot nos datasets de hPancreas e córtex.
  • Reconstrução: Embeddings congelados preservaram a estrutura de expressão contínua, recuperando mais genes altamente expressos e preservando correlações de vias do Reactome melhor do que outros modelos.
  • Resposta à Perturbação: Alcançou o menor erro quadrático médio agregado na predição de respostas de log-fold-change a perturbações genéticas (datasets Adamson, Norman, Dixit) e melhor recuperação dos genes de resposta principais.
• Eficiência: Em um experimento de ajuste fino (fine-tuning) de uma época correspondente no hPancreas (batch size 8, ~26M parâmetros treináveis), o BioM-JEPA foi 5,75× mais rápido em throughput de treinamento e 3,76× mais rápido na extração de embeddings em dados não vistos (held-out) em comparação ao scFoundation.

Significância e Alegações
O artigo afirma que blocos de genes conectados por grafos servem como unidades de predição eficazes para o aprendizado de representação do tipo JEPA em biologia de célula única. Os autores afirmam que:

• A convergência de um objetivo ao nível de gene (predição de token) não garante uma representação celular robusta; a escala da unidade de predição importa.
• O objetivo ao nível de bloco é bem adequado para dados transcriptômicos esparsos porque genes conectados fornecem evidências redundantes para um estado latente compartilhado, tornando a representação menos dependente de qualquer medição esparsa individual.
• As representações aprendidas capturam a estrutura biológica (identidades celulares canônicas, relações de vias e composições de perturbação genética) sem supervisão explícita de tipos celulares ou vias.
• A atenção linear torna esta abordagem computacionalmente prática para grandes contextos de genes, permitindo que representações úteis emerjam de uma exposição de treinamento comparativamente compacta em relação a outros modelos de grande escala.

Os autores concluem que esta abordagem oferece uma rota direta para definir unidades de predição baseadas em relações moleculares (vias, complexos) em vez de tratar genes como "patches" independentes, oferecendo uma base para futuros modelos de representação biológica.