Resumo Técnico: Pós-pré-treinamento Guiado por Protótipos para Aprendizado de Representação de Células Únicas

Declaração do Problema

O aprendizado de representação de células únicas (SCRL) visa decodificar dados de expressão gênica de alta dimensão para revelar a lógica regulatória celular. Embora modelos de base recentes inspirados em modelos de linguagem grandes (tratando genes como tokens e células como sentenças) tenham mostrado promessa, eles enfrentam duas limitações fundamentais:

1. Distribuições de Tipo Celular de Cauda Longa: Conjuntos de dados de células únicas exibem desequilíbrio de classes extremo, onde populações celulares raras, mas biologicamente críticas (por exemplo, células-tronco iniciadoras de doenças), constituem uma fração ínfima das amostras. Isso leva a embeddings ruidosos ou colapsados para tipos celulares raros, pois os modelos são tendenciosos em relação às populações abundantes.
2. Mudanças de Covariáveis: Artefatos técnicos, como efeitos de lote e profundidades de sequenciamento variáveis, criam lacunas de distribuição entre os dados de treinamento e inferência. Modelos existentes frequentemente superajustam-se a padrões específicos de lote em vez de capturar a lógica regulatória universal, limitando sua generalização entre diferentes plataformas ou protocolos experimentais.

Estratégias atuais de ajuste fino frequentemente falham em resolver completamente essas questões porque começam a partir de uma variedade latente de embeddings instável estabelecida durante o pré-treinamento inicial.

Metodologia: CellRefine

Para abordar esses desafios, os autores introduzem o CellRefine, um framework de pós-pré-treinamento projetado para operar entre a fase inicial de pré-treinamento e o ajuste fino downstream. O CellRefine injeta vieses indutivos fundamentados biologicamente para remodelar o espaço de embeddings latentes em uma estrutura mais estável e discriminativa.

O método emprega uma função objetivo multifacetada ($L_{Total}$) combinando quatro componentes:

1. Aprendizado Guiado por Protótipos (Programas de Genes Marcadores):

   • Conceito: Utiliza conjuntos conhecidos de genes marcadores como priores estruturais. Como conjuntos completos de marcadores são frequentemente desconhecidos, o método usa subconjuntos curados de marcadores para criar "protótipos" para cada tipo celular.
   • Implementação: Genes marcadores para um tipo celular $c_i$ são ordenados com base na hierarquia da ontologia celular (de subtipos específicos a linhagens mais amplas) usando uma função $H(c_i)$. Essas sequências ordenadas são tokenizadas e codificadas em embeddings de protótipo ($z_m$).
   • Perda: Uma perda de regularização guiada por protótipo ($L_{prototype}$) maximiza a similaridade de cosseno entre o embedding de uma célula ($z_c$) e seu embedding de protótipo correspondente, enquanto minimiza a similaridade com outros protótipos. Isso injeta estrutura biológica na representação.

2. Regularização Consciente de Linhagem:

   • Conceito: Aborda a questão em que tipos celulares intimamente relacionados (compartilhando uma linhagem parental na ontologia) permanecem insuficientemente separados no espaço de embeddings.
   • Implementação: Uma perda de separação consciente de linhagem ($L_{lineage}$) penaliza a similaridade excessiva entre pares de células que compartilham a mesma linhagem parental e nível de ontologia, incentivando o modelo a distinguir identidades biológicas de alta granularidade.

3. Codificação Variacional de Mistura Gaussiana (GMVE):

   • Conceito: Modela explicitamente o espaço latente como uma estrutura agrupada para melhorar a qualidade da representação.
   • Implementação: Um Codificador Variacional de Mistura Gaussiana é acoplado ao modelo. Uma perda de divergência KL ($L_{GMVE}$) regulariza o posterior variacional contra um prior de mistura gaussiana aprendido, impondo uma estrutura latente agrupada.

4. Modelagem de Linguagem Mascarada (MLM):

   • O objetivo padrão de pré-treinamento é mantido para garantir que o modelo continue a aprender padrões gerais de expressão gênica.

O objetivo total é:
$$L_{Total} = L_{MLM} + \lambda_1 L_{prototype} + \lambda_2 L_{lineage} + \lambda_3 L_{GMVE}$$

Contribuições Principais

• Fase de Pós-pré-treinamento Inovadora: Propõe uma fase intermediária de adaptação (CellRefine) que refina a variedade latente de modelos de base antes do ajuste fino específico da tarefa, em vez de depender exclusivamente de pré-treinamento seguido de ajuste fino direto.
• Integração de Priores Biológicos: Integra sistematicamente conhecimento biológico mecanicista (programas de genes marcadores e ontologias celulares) diretamente no processo de otimização para guiar o aprendizado de representação.
• Tratamento de Cauda Longa e Mudanças de Covariáveis: Demonstra que fundamentar o modelo em priores biológicos mitiga o viés em relação a tipos celulares abundantes e melhora a robustez a mudanças técnicas.

Resultados Experimentais

Os autores avaliaram o CellRefine em três tarefas centrais de biologia computacional em 11 conjuntos de dados abrangendo vários tecidos (sangue, pâncreas, fígado, cérebro, pulmão, coração, etc.):

1. Previsão de Identidade Celular:

   • No Domínio: O CellRefine superou consistentemente as linhas de base mais avançadas (incluindo ajuste fino completo, LoRA e alinhamento MLM padrão) em 10 conjuntos de dados. Alcançou melhorias de macro-F1 de até 15% sobre a linha de base mais forte (MLM $\to$ FF) em conjuntos de dados desbalanceados como Sangue e Pâncreas.
   • Fora do Domínio (Zero-Shot): Na transferência cross-modal (treinamento em scRNA-seq, teste em transcriptômica espacial), o CellRefine mostrou desempenho superior de recall@k, indicando representações mais robustas e generalizáveis.

2. Imputação de Transcriptômica Espacial:

   • O CellRefine melhorou a previsão de expressões gênicas não medidas, alcançando maiores pontuações de correlação de Pearson e similaridade de cosseno em comparação com as linhas de base em conjuntos de dados espaciais de fígado.

3. Previsão de Resposta a Perturbação Gênica:

   • Na previsão de mudanças na expressão gênica sob perturbação, o CellRefine produziu melhorias consistentes, embora menores, sobre as linhas de base, alcançando o melhor desempenho geral entre todos os métodos testados.

Estudos de Ablação:

• A remoção de qualquer um dos três termos de regularização ($L_{prototype}$, $L_{lineage}$, $L_{GMVE}$) resultou em degradação de desempenho, confirmando que a abordagem multiobjetivo é sinérgica.
• Eficiência de Amostra: Em cenários de aprendizado com poucos exemplos (5–50 exemplos), o CellRefine demonstrou ganhos de desempenho significativamente mais acentuados em baixas contagens de amostras em comparação com o ajuste fino padrão, sugerindo que ele aproveita dados rotulados limitados de forma mais eficaz.

Análise Qualitativa:

• Visualizações UMAP do espaço latente revelaram que o CellRefine reorganiza a variedade de embeddings, tornando populações celulares raras de "cauda" (por exemplo, subconjuntos específicos de células T, tipos celulares raros de fígado) mais distinguíveis e menos emaranhadas com grandes clusters celulares em comparação com modelos treinados apenas com MLM.

Significado e Alegações

O artigo alega que o CellRefine representa uma direção promissora para o aprendizado de representação de células únicas, demonstrando que vieses indutivos informados biologicamente podem melhorar significativamente o desempenho downstream de modelos de base.

• Utilidade Clínica: Ao capturar melhor tipos celulares raros e melhorar a robustez a mudanças de covariáveis, o método aumenta o potencial para identificar mecanismos de doenças e alvos terapêuticos em estados celulares raros ou transitórios que são frequentemente obscurecidos em modelos padrão.
• Mudança Metodológica: O trabalho argumenta contra depender exclusivamente de pré-treinamento orientado por dados seguido de ajuste fino, defendendo, em vez disso, fases intermediárias que incorporam explicitamente conhecimento de domínio (genes marcadores e ontologias) para estabilizar o processo de aprendizado.
• Limitações: Os autores observam modestamente que a abordagem depende de programas de genes marcadores curados, que podem estar incompletos ou ser ruidosos para tipos celulares novos. Além disso, os ganhos de desempenho não implicam necessariamente em eficiência de parâmetros aprimorada.

Em resumo, o CellRefine oferece um framework para preencher a lacuna entre o pré-treinamento em grande escala e as realidades específicas, frequentemente desbalanceadas, dos dados biológicos de células únicas, produzindo representações celulares mais discriminativas e generalizáveis.