Integrating gene regulatory priors into Transformer attention with scTransformer for interpretable scRNA-seq analysis

Resumo Técnico: scTransformer – Integrando Priores de Regulação Gênica na Atenção de Transformers

Declaração do Problema

Os dados de sequenciamento de RNA de célula única (scRNA-seq) são caracterizados por alta dimensionalidade, esparsidade e ruído técnico, o que impõe desafios significativos para o aprendizado de representação. Embora modelos de fundação baseados em Transformer (ex: scGPT, Geneformer) tenham demonstrado escalabilidade ao serem pré-treinados em milhões de células, eles tipicamente tratam os genes como tokens independentes com autoatenção não restrita. Esse design permite que qualquer gene atenda a qualquer outro, baseando-se apenas na covariância estatística para descobrir dependências. Consequentemente, esses modelos frequentemente sofrem de:

1. Interpretabilidade Limitada: Padrões de atenção são difíceis de mapear para processos biológicos mecanísticos.
2. Instabilidade: Em regimes de baixos dados ou na presença de efeitos de lote (batch effects), modelos não restritos podem sofrer overfitting para correlações específicas do conjunto de dados, gerando previsões que são robustas em acurácia, mas instáveis na seleção de características (ou seja, diferentes execuções de treinamento selecionam conjuntos de genes disjuntos).
3. Conhecimento Biológico Ignorado: Abordagens existentes ignoram amplamente redes regulatórias curadas de fator de transcrição (TF)–gene alvo (TG), que codificam estruturas causais plausíveis no controle transcricional.

Metodologia: Arquitetura scTransformer

Os autores propõem o scTransformer, um codificador Transformer biologicamente informado que integra redes regulatórias curadas de TF–TG diretamente no mecanismo de autoatenção via mascaramento rígido (hard masking).

1. Representação de Entrada e Embedding:

• Embedding de Token Gênico: Cada gene em uma célula é mapeado para um token de $d$ dimensões combinando um embedding de identidade gênica treinável ($e^{id}$) e uma codificação de expressão contínua ($e^{expr}$).
• Codificação Sinusoidal Contínua: Diferente da discretização por bins, os valores de expressão são codificados usando um esquema sinusoidal contínuo ($\sin(x \cdot \omega_k), \cos(x \cdot \omega_k)$). Isso preserva a informação quantitativa ao longo da amplitude das medições de scRNA-seq, garantindo suavidade e limitação.
• Pooling: Um mecanismo de pooling baseado em similaridade agrega sequências de tokens de comprimento variável em um embedding de célula de dimensão fixa. O vetor de contexto é computado como a média dos embeddings de tokens e usado como uma consulta (query) em uma atenção cruzada sobre a sequência de tokens.

2. Mecanismo de Atenção com Gate de Priores:
A inovação central é a substituição da atenção não restrita por um mecanismo de atenção com gate de priores (prior-gated):

• Rede Regulatória: Uma rede TF–TG curada (do CollecTRI) é pré-processada para filtrar TFs com alvos suficientes e genes presentes no conjunto de dados.
• Mascaramento Rígido (Hard Masking): Uma matriz de permissão booleana ($M_{allow}$) é construída onde $M_{allow}(i, j) = 1$ apenas se o gene $i$ for um TF e o gene $j$ for seu alvo (ou $i=j$ para auto-loops). Todas as outras interações são definidas como 0.
• Restrição de Atenção: Esta matriz é convertida em uma máscara aditiva ($\tilde{M}_{prior}$) onde as interações proibidas recebem $-\infty$. Antes da operação softmax na atenção de produto escalar escalonada, esta máscara é adicionada aos escores de atenção:
  $$A_h = \text{softmax}\left(\frac{Q_h K_h^\top}{\sqrt{d_k}} + \tilde{M}_{prior} + M_{pad}\right)$$
  Isso garante que a informação flua apenas ao longo das arestas regulatórias conhecidas (e auto-loops), restringindo efetivamente o espaço de hipóteses para interações biologicamente plausíveis.

3. Métricas de Interpretabilidade:
Para quantificar a relevância biológica da atenção aprendida, os autores definem métricas de nível de módulo:

• Escore de Concentração ($\phi$): Mede o quão seletivamente um TF concentra a atenção em seus alvos anotados usando entropia normalizada. Uma entropia menor (maior $\phi$) indica uma atenção focada e estruturada.
• Importância do Módulo: Combina o escore de concentração com a massa total de atenção direcionada aos alvos para quantificar a força dos sinais regulatórios.

4. Detalhes de Implementação:

• Dados: Avaliado no atlas de RNA de núcleo único MSSM (coorte PsychAD), restrito a 17 subtipos e interseções TF–TG.
• Treinamento: Utiliza otimizador AdamW com arquitetura Pre-LN. O treinamento emprega lote esparso (sparse batching) diretamente em matrizes de Linha de Disposição Comprimida (CSR) para lidar eficientemente com grandes conjuntos de dados.
• Avaliação: Os modelos foram testados em vários regimes de dados de treinamento (1% a 100%) para avaliar a eficiência de amostragem e robustez.

Principais Contribuições

1. Atenção com Gate de Priores: O primeiro framework de Transformer que incorpora redes TF–TG curadas como restrições rígidas dentro do mecanismo de atenção, convertendo conhecimento regulatório em viés indutivo estrutural.
2. Codificação de Expressão Contínua: Um esquema de codificação sinusoidal que preserva a natureza quantitativa dos dados de scRNA-seq sem discretização.
3. Pipeline de Treinamento Esparso: Uma implementação eficiente que opera diretamente em matrizes CSR, permitindo o treinamento escalável em grandes conjuntos de dados de célula única.
4. Métricas de Interpretabilidade Formais: Um conjunto de métricas para quantificar a concentração de atenção e a importância regulatória, permitindo a avaliação direta da coerência biológica sem análise post-hoc.
5. Protocolo Reproduzível: Uma configuração de treinamento distribuído com particionamento de dados controlado e avaliação de múltiplas sementes para avaliar rigorosamente a estabilidade.

Resultos

Os autores avaliaram o scTransformer contra baselines não restritos em tarefas de classificação de tipos celulares:

• Desempenho em Regimes de Baixos Dados: Em configurações com dados de treinamento limitados (1% e 25%), o modelo com gate de priores alcançou uma acurácia significativamente maior do que o baseline não restrito (ex: 0,486 vs 0,464 em 1% de dados). Isso demonstra melhor eficiência de amostragem.
• Desempenho em Regimes de Dados Totais: Na disponibilidade de 100% dos dados, o modelo com gate de priores (usando apenas genes de interseção TF–TG) alcançou acurácia comparável ao baseline não restrito usando todos os genes (0,911 vs 0,917), apesar de operar em um espaço de interação muito menor.
• Estabilidade e Reproduzibilidade: Modelos não restritos mostraram alta variabilidade entre execuções de treinamento, selecionando conjuntos de genes disjuntos (índices de Jaccard próximos de zero, correlações de Spearman negativas). Em contraste, os modelos com gate de priores exibiram alta estabilidade, selecionando consistentemente subconjuntos de genes semelhantes e classificando-os de forma consistente.
• Coerência Biológica: Mapas de atenção de modelos com gate de priores mostraram maiores escores de concentração ($\phi$), indicando que os TFs focaram a atenção em subconjuntos de alvos específicos e biologicamente relevantes. Esses padrões alinharam-se com programas regulatórios específicos de subtipos conhecidos (ex: enriquecimento de GRHL3 em populações endoteliais, NANOG em OPCs), enquanto os modelos não restritos produziram distribuições de atenção difusas e menos interpretáveis.

Significância e Alegações

O artigo alega que o valor primário de integrar priores regulatórios não é meramente melhorar a acurácia preditiva, mas aumentar a confiabilidade epistêmica e a interpretabilidade das representações aprendidas.

• Resolução de Ambiguidade: Em dados de alta dimensão e esparsidade, múltiplas combinações de genes podem explicar o mesmo fenótipo. Modelos não restritos podem convergir para soluções arbitrárias e instáveis. Os priores regulatórios restringem o espaço de hipóteses, direcionando o modelo para soluções coerentes e biologicamente fundamentadas.
• Viés Indutivo: Ao incorporar a estrutura biológica diretamente na arquitetura, o modelo aprende representações que são robustas ao ruído e efeitos de lote sem sacrificar a flexibilidade quando há dados suficientes disponíveis.
• Fundação para Interpretabilidade: A abordagem oferece um passo principiado em direção a "modelos de fundação biologicamente fundamentados", onde os pesos de atenção são diretamente interpretáveis como programas regulatórios, facilitando o insight biológico sem a necessidade de análise post-hoc.

Os autores mantêm a modéstia, observando que as limitações atuais incluem o uso de máscaras binárias rígidas (que impedem o aprendizado de novas interações), a restrição a tarefas supervisionadas e a avaliação em um único conjunto de dados. Eles sugerem que trabalhos futuros poderiam explorar priores suaves/aprendíveis, pré-treinamento auto-supervisionado e arquiteturas híbridas.