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

Resumen Técnico: scTransformer – Integración de Priores de Regulación Génica en la Atención de Transformers

Planteamiento del Problema

Los datos de secuenciación de ARN de célula única (scRNA-seq) se caracterizan por su alta dimensionalidad, escasez (sparsity) y ruido técnico, lo que plantea desafíos significativos para el aprendizaje de representaciones. Aunque los modelos fundacionales basados en Transformers (por ejemplo, scGPT, Geneformer) han demostrado escalabilidad al preentrenar sobre millones de células, estos suelen tratar los genes como tokens independientes con autoatención sin restricciones. Este diseño permite que cualquier gen atienda a cualquier otro, dependiendo únicamente de la covariación estadística para descubrir dependencias. En consecuencia, estos modelos suelen sufrir de:

1. Interpretabilidad Limitada: Los patrones de atención son difíciles de mapear hacia procesos biológicos mecanísticos.
2. Inestabilidad: En regímenes de bajos datos o ante la presencia de efectos de lote (batch effects), los modelos sin restricciones pueden sobreajustarse a correlaciones específicas del conjunto de datos, produciendo predicciones que son robustas en precisión pero inestables en la selección de características (es decir, diferentes ejecuciones de entrenamiento seleccionan conjuntos de genes disjuntos).
3. Ignorancia del Conocimiento Biológico: Los enfoques existentes ignoran en gran medida las redes reguladoras curadas de factor de transcripción (TF)–gen diana (TG), las cuales codifican estructuras causales plausibles en el control transcripcional.

Metodología: Arquitectura scTransformer

Los autores proponen scTransformer, un codificador Transformer biológicamente informado que integra redes reguladoras curadas de TF–TG directamente en el mecanismo de autoatención mediante el uso de enmascaramiento estricto (hard masking).

1. Representación de Entrada y Embedding:

• Embedding de Token de Gen: Cada gen en una célula se mapea a un token de $d$ dimensiones que combina un embedding de identidad de gen entrenable ($e^{id}$) y una codificación de expresión continua ($e^{expr}$).
• Codificación Sinusoidal Continua: A diferencia del agrupamiento (binning) discreto, los valores de expresión se codifican mediante un esquema sinusoidal continuo ($\sin(x \cdot \omega_k), \cos(x \cdot \omega_k)$). Esto preserva la información cuantitativa a través del rango dinámico de las mediciones de scRNA-seq, asegurando suavidad y acotación.
• Pooling: Un mecanismo de pooling basado en similitud agrega secuencias de tokens de longitud variable en un embedding de célula de dimensión fija. El vector de contexto se computa como la media de los embeddings de los tokens y se utiliza como consulta (query) en la atención cruzada sobre la secuencia de tokens.

2. Mecanismo de Atención con Puerta de Priores (Prior-Gated Attention):
La innovación central es el reemplazo de la atención sin restricciones por un mecanismo de puerta de priores:

• Red Regulatoria: Una red TF–TG curada (de CollecTRI) se preprocesa para filtrar TFs con suficientes objetivos y genes presentes en el conjunto de datos.
• Enmascaramiento Estricto (Hard Masking): Se construye una matriz de permiso booleana ($M_{allow}$) donde $M_{allow}(i, j) = 1$ solo si el gen $i$ es un TF y el gen $j$ es su objetivo (o $i=j$ para bucles de retroalimentación). Todas las demás interacciones se establecen en 0.
• Restricción de Atención: Esta matriz se convierte en una máscara aditiva ($\tilde{M}_{prior}$) donde las interacciones prohibidas se asignan como $-\infty$. Antes de la operación softmax en la atención de producto escalar escalado, esta máscara se suma a las puntuaciones de atención:
  $$A_h = \text{softmax}\left(\frac{Q_h K_h^\top}{\sqrt{d_k}} + \tilde{M}_{prior} + M_{pad}\right)$$
  Esto asegura que la información fluya solo a lo largo de los bordes regulatorios conocidos (y bucles de retroalimentación), restringiendo efectivamente el espacio de hipótesis hacia interacciones biológicamente plausibles.

3. Métricas de Interpretabilidad:
Para cuantificar la relevancia biológica de la atención aprendida, los autores definen métricas a nivel de módulo:

• Puntuación de Concentración ($\phi$): Mide qué tan selectivamente un TF concentra la atención en sus objetivos anotados utilizando entropía normalizada. Una menor entropía (mayor $\phi$) indica una atención enfocada y estructurada.
• Importancia del Módulo: Combina la puntuación de concentración con la masa total de atención dirigida hacia los objetivos para cuantificar la fuerza de las señales regulatorias.

4. Detalles de Implementación:

• Datos: Evaluado en el atlas de RNA-seq de núcleo único MSSM (cohorte PsychAD), restringido a 17 subtipos e intersecciones de genes TF–TG.
• Entrenamiento: Utiliza el optimizador AdamW con una arquitectura Pre-LN. El entrenamiento emplea lotes dispersos (sparse batching) directamente sobre matrices de fila de contenido comprimido (CSR) para manejar grandes conjuntos de datos de manera eficiente.
• Evaluación: Los modelos fueron probados a través de diversos regímenes de datos de entrenamiento (del 1% al 100%) para evaluar la eficiencia de la muestra y la robustez.

Contribuciones Clave

1. Atención con Puerta de Priores (Prior-Gated Attention): El primer marco de trabajo Transformer que incorpora redes TF–TG curadas como restricciones estrictas dentro del mecanismo de atención, convirtiendo el conocimiento regulatorio en un sesgo inductivo estructural.
2. Codificación de Expresión Continua: Un esquema de codificación sinusoidal que preserva la naturaleza cuantitativa de los datos de scRNA-seq sin recurrir a la discretización.
3. Pipeline de Entrenamiento Disperso: Una implementación eficiente que opera directamente sobre matrices CSR, permitiendo el entrenamiento escalable en grandes conjuntos de datos de célula única.
4. Métricas de Interpretabilidad Formales: Un conjunto de métricas para cuantificar la concentración de la atención y la importancia regulatoria, permitiendo la evaluación directa de la coherencia biológica sin necesidad de análisis post-hoc.
5. Protocolo Reproducible: Una configuración de entrenamiento distribuido con partición de datos controlada y evaluación multi-semilla para evaluar rigurosamente la estabilidad.

Resultados

Los autores evaluaron scTransformer frente a líneas de base sin restricciones en tareas de clasificación de tipos celulares:

• Desempeño en Regímenes de Bajos Datos: En configuraciones con datos de entrenamiento limitados (1% y 25%), el modelo con puerta de priores logró una precisión significativamente mayor que la línea de base sin restricciones (por ejemplo, 0.486 vs. 0.464 al 1% de datos). Esto demuestra una mejor eficiencia de muestra.
• Desempeño en Regímenes de Datos Completos: Con disponibilidad de datos del 100%, el modelo con puerta de priores (usando solo genes de intersección TF–TG) alcanzó una precisión comparable a la de la línea de base sin restricciones que utiliza todos los genes (0.911 vs. 0.917), a pesar de operar en un espacio de interacción mucho más pequeño.
• Estabilidad y Reproducibilidad: Los modelos sin restricciones mostraron una alta variabilidad entre ejecuciones de entrenamiento, seleccionando conjuntos de genes disjuntos (índices de Jaccard cercanos a cero, correlaciones de Spearman negativas). En contraste, los modelos con puerta de priores exhibieron una alta estabilidad, seleccionando consistentemente subconjuntos de genes similares y clasificándolos de manera consistente.
• Coherencia Biológica: Los mapas de atención de los modelos con puerta de priores mostraron puntuaciones de concentración ($\phi$) más altas, indicando que los TFs enfocan la atención en subconjuntos de objetivos específicos y biológicamente relevantes. Estos patrones se alinearon con programas regulatorios específicos de subtipo (por ejemplo, el enriquecimiento de GRHL3 en poblaciones de tipo endotelial, NANOG en OPCs), mientras que los modelos sin restricciones produjeron distribuciones de atención difusas y menos interpretables.

Significancia y Reivindicaciones

El artículo sostiene que el valor principal de integrar los priores regulatorios no es meramente mejorar la precisión predictiva, sino mejorar la fiabilidad epistémica e la interpretabilidad de las representaciones aprendidas.

• Resolución de la Ambigüedad: En datos de alta dimensión y dispersos, múltiples combinaciones de genes pueden explicar el mismo fenotipo. Los modelos sin restricciones pueden converger a soluciones arbitrarias e inestables. Los priores regulatorios restringen el espacio de hipótesis, dirigiendo al modelo hacia soluciones coherentes y con base biológica.
• Sesgo Inductivo: Al embeber la estructura biológica directamente en la arquitectura, el modelo aprende representaciones que son robustas al ruido y a los efectos de lote sin sacrificar la flexibilidad cuando hay suficientes datos disponibles.
• Fundamento para la Interpretabilidad: El enfoque ofrece un paso hacia "modelos fundacionales con base biológica", donde los pesos de atención son directamente interpretables como programas regulatorios, facilitando el conocimiento biológico sin requerir análisis post-hoc.

Los autores mantienen la modestia, señalando que las limitaciones actuales incluyen el uso de máscaras binarias estrictas (que impiden el aprendizaje de nuevas interacciones), la restricción a tareas supervisadas y la evaluación en un solo conjunto de datos. Sugieren que el trabajo futuro podría explorar priores suaves o aprendibles, preentrenamiento autosupervisado y arquitecturas híbridas.