Arc Virtual Cell Challenge: Un primer sobre la generalización de contexto en transcriptómica

TL;DR

• El Arc Institute presenta el Virtual Cell Challenge, un reto para entrenar un modelo que prediga el efecto de silenciar un gen en tipos celulares no vistos, un objetivo denominado generalización de contexto.
• El conjunto de datos contiene aproximadamente 300k perfiles de transcriptómica de células individuales; el conjunto de entrenamiento incluye 220k células, con cerca de 38k controles no perturbados.
• Arc publicó STATE, una baseline robusta compuesta por dos modelos basados en transformadores: State Transition Model (ST) y State Embedding Model (SE).
• La evaluación se centra en Perturbation Discrimination, Differential Expression y Mean Average Error, con PDiscNorm definido como 1 - 2 * PDisc.
• El objetivo es acelerar las pruebas in silico de perturbaciones, con potencial para acelerar la investigación biológica y la descubrimiento de fármacos al reducir los bucles de retroalimentación.

Contexto y antecedentes

Arc Institute anunció recientemente el Virtual Cell Challenge. Los participantes deben entrenar un modelo capaz de predecir el efecto de silenciar un gen en un tipo de célula (parcialmente) no visto, una tarea que llaman generalización de contexto. Para ingenieros de ML con poco conocimiento en biología, el reto se presenta de modo accesible: entrenar un modelo que pueda predecir exactamente qué sucede a una célula cuando se modifica un parámetro. Para entrenar la red, Arc ha recopilado un conjunto de datos con aproximadamente 300k perfiles de ARN de célula única (scRNA-seq). El conjunto de entrenamiento consta de 220k células, y de esas, aproximadamente 38k son células de control no perturbadas. Cada célula se representa por su transcriptoma, que es un vector disperso en el que cada entrada es la cuenta bruta de moléculas de ARN del gene correspondiente. Un ejemplo concreto es el gen TMSB4X, el más frecuentemente silenciado. Se observa que las células con silenciación de TMSB4X presentan una reducción notable en el número de transcritos frente a las células de control. El fenómeno del observador impide medir el estado de la célula antes y después de la perturbación, por lo que se usan células basales como referencia y hay que separar la señal de la perturbación del ruido por la heterogeneidad. Antes del Virtual Cell Challenge, Arc publicó STATE, su intento de solución que utiliza un par de modelos basados en transformadores. STATE consta de dos modelos: el State Transition Model (ST) y el State Embedding Model (SE). ST es un “simulador celular” que toma ya sea un transcriptoma de control, o una embedding producida por SE, junto con un vector de perturbación binaria, y devuelve el transcriptoma perturbado. ST utiliza un transformer con backbone Llama y opera sobre un conjunto de controles covariados pareados para ayudar a discernir el efecto real de la perturbación. Ambos tensores, de control y perturbación, se procesan por codificadores independientes, que son MLPs de 4 capas con activaciones GELU. Si se trabaja en el espacio de expresión génica, se pasa por un decodificador aprendido. ST se entrena con Maximum Mean Discrepancy para minimizar la diferencia entre las distribuciones. SE, por su parte, es un autoencoder tipo BERT para producir embeddings semánticos de células. Para entenderlo, SE crea embeddings de genes a partir de las secuencias de aminoácidos de todas las isoformas proteicas codificadas por el gen. Estas secuencias se procesan en ESM, un modelo de lenguaje proteico de 15 mil millones de parámetros. ESM genera un embedding por aminoácido y se promedia para obtener un embedding de isoforma, que luego se proyecta al espacio del modelo mediante un codificador aprendido. Se obtiene un embedding de gen y, para la célula, se toma el top 2048 de genes por expresión logarítmica para formar una “oración celular” con tokens especiales [CLS] y [DS] para separar efectos del set de datos. El token [CLS] se usa como embedding de la célula. Un mecanismo de codificación de expresión (parecido a embeddings posicionales) añade “encodings” de expresión que modulan la magnitud de cada embedding de gen. Durante el entrenamiento, se enmascaran 1280 genes por célula y se entrena para predecirlos. La evaluación utiliza tres métricas: Perturbation Discrimination, Differential Expression y Mean Average Error. Perturbation Discrimination evalúa cuán bien su modelo puede distinguir diferencias relativas entre perturbaciones. Para ello, calculamos las distancias de Manhattan entre los transcriptomas perturbados medidos y los transcriptomas perturbados predichos, y clasificamos el lugar de la verdad entre todas las perturbaciones; el score se normaliza como PDiscNorm = 1 - 2 * PDisc. Differential Expression estima la fracción de genes verdaderamente afectados que se identifican como significativos, calculando p-valores con una prueba de Wilcoxon con corrección de empates y aplicando el procedimiento de Benjamini-Hochberg. Estas descripciones provienen de la publicación de Arc sobre STATE y su uso como baseline para el Virtual Cell Challenge. Hugging Face Blog

What’s new

La publicación de STATE ofrece una baseline concreta para el Virtual Cell Challenge, con dos módulos: State Transition Model (ST) y State Embedding Model (SE).

• ST es un transformador con backbone Llama. Recibe entradas de control covariadas y perturbación objetivo, cada una procesada por codificadores independientes (MLP de 4 capas con GELU). Si se opera en el espacio de expresión génica, ST utiliza un decodificador aprendido para generar el transcriptoma perturbado. El entrenamiento utiliza Maximum Mean Discrepancy para alinear las distribuciones previstas y reales.
• SE es un autoencoder tipo BERT que genera embeddings significativos de células y genes. Los embeddings de genes se obtienen de las secuencias de aminoácidos de las isoformas proteicas y se procesan mediante ESM2 para producir embeddings de isoformas que luego se promedian para obtener embeddings de genes. La célula se representa mediante una sentencia de 2048 genes con tokens [CLS] y [DS] para separar efectos del conjunto de datos. Se reafirman las magnitudes de expresión a través de una codificación suave de expresión y dos MLPs.
• La representación de la célula se obtiene a partir de los embeddings de los 2048 genes más expresados, con un [CLS] que sirve como embedding de la célula y un [DS] para disentangle efectos específicos del conjunto de datos.
• Los datos y la evaluación muestran que la línea STATE incorpora información transcriptómica y proteica para mejorar la generalización entre tipos celulares.

Why it matters (impacto para desarrolladores/empresas)

El Virtual Cell Challenge aborda un cuello de botella en biología y descubrimiento de fármacos: probar perturbaciones genéticas sin experimentos de laboratorio. Un modelo capaz de simular las consecuencias del silenciado de genes a través de tipos celulares podría reducir el tiempo y costo de cribado de perturbaciones y orientar la experimentación al priorizar perturbaciones con efectos consistentes entre tipos celulares. Esto podría acelerar la generación de hipótesis y guiar el diseño experimental al señalar perturbaciones con mayor probabilidad de efecto en múltiples contextos celulares.

Detalles técnicos o Implementación

A continuación se resumen aspectos técnicos centrales de STATE descritos por Arc:

• El conjunto de datos contiene expresión observada en células perturbadas y de control, modelando la expresión como X^p ∼ T^p(D_basal) + H(D_basal) + ε. Las predicciones de estado perturbado ^X_p provienen de un transcriptoma de control o de una embedding generada por SE, junto con un vector de perturbación.
• State Transition Model (ST): transformer con backbone Llama. Los inputs se procesan a través de codificadores independientes (control y perturbación), cada uno con MLPs de 4 capas y activaciones GELU. Si se trabaja en espacio de expresión gênica, se utiliza un decodificador aprendido para generar el transcriptoma perturbado. El entrenamiento usa Maximum Mean Discrepancy.
• State Embedding Model (SE): autoencoder tipo BERT que produce embeddings significativos. Los embeddings de genes se obtienen a partir de las secuencias de aminoácidos de isoformas proteicas usando ESM2. Esos embeddings se agregan para formar embeddings de genes y se proyectan mediante un codificador aprendido. La célula se representa por una sentencia de 2048 genes con [CLS] como embedding de la célula y [DS] para disentangle effects. Se aplica una codificación suave de expresión para modular la magnitud de cada embedding de gen.
• Construcción de la representación celular: la célula se codifica con los embeddings de los 2048 genes más expresados, formando una secuencia que alimenta el modelo para obtener una embedding celular significativa.
• Entrenamiento y enmascaramiento: 1280 genes por célula son enmascarados y deben ser predichos.
• Evaluación: se utilizan métricas como Perturbation Discrimination, Differential Expression y MAE. Hugging Face Blog

Puntos clave

• El Virtual Cell Challenge formaliza la generalización de contexto para perturbaciones celulares, con foco en tipos no vistos.
• El conjunto de datos combina controles y perturbaciones para evaluar la separación de señal real de perturbación frente a la heterogeneidad basal.
• STATE propone una baseline que combina embeddings de genes y proteínas con un transformer para favorecer la generalización entre tipos celulares.
• Las métricas de evaluación cubren tanto la precisión predictiva como la relevancia biológica.
• El objetivo es habilitar pruebas in silico de perturbaciones para acelerar decisiones experimentales y descubrimiento de fármacos.

Referencias

• Hugging Face Blog – Arc Virtual Cell Challenge: A Primer. https://huggingface.co/blog/virtual-cell-challenge