Resumen Técnico: MapPFN - Aprendizaje de Mapas de Perturbación Causal en Contexto

1. Planteamiento del Problema

La comprensión mecanicista de las poblaciones celulares requiere identificar las dependencias causales en las redes de regulación génica (GRN) mediante datos de perturbación de células individuales. Sin embargo, existen limitaciones críticas en el estado actual:

• Escasez de datos: Los conjuntos de datos de perturbación de células individuales abarcan solo unos pocos contextos biológicos (líneas celulares).
• Falta de adaptabilidad: Los métodos existentes no pueden aprovechar nuevas evidencias intervencionistas en el momento de la inferencia para adaptarse a contextos no vistos durante el entrenamiento.
• Complejidad de la predicción: Predecir la respuesta a una perturbación es un problema de mapeo entre distribuciones no emparejadas (ya que la secuenciación destruye las células individuales), lo que dificulta la generalización a nuevos genes o contextos biológicos sin reentrenamiento costoso.

El objetivo es desarrollar un modelo que pueda meta-aprender mapas de perturbación, permitiendo la adaptación a contextos biológicos no vistos y a conjuntos de genes arbitrarios mediante aprendizaje en contexto (In-Context Learning, ICL).

2. Metodología: MapPFN

MapPFN es la primera Red Ajustada a Datos Previos (Prior-Data Fitted Network, PFN) diseñada específicamente para la predicción de perturbaciones.

2.1. Enfoque General

En lugar de entrenar un modelo específico para cada tarea o contexto biológico, MapPFN se entrena previamente en un prior sintético masivo generado a partir de modelos causales. Durante la inferencia, utiliza el aprendizaje en contexto para predecir distribuciones post-perturbación basándose en un conjunto de experimentos observacionales e intervencionistas proporcionados en el momento.

2.2. Arquitectura del Modelo

• Base: Utiliza una arquitectura Transformador de Difusión Multimodal (MMDiT), modificada para manejar datos de células individuales.
• Tratamiento de Datos: Las células se tratan como "tokens". El modelo procesa tres flujos de modalidades separadas:
  1. Ruido (input para la difusión).
  2. Estados celulares (datos observacionales e intervencionistas).
  3. Tratamientos (codificados one-hot).
• Mecanismo: Utiliza atención cruzada para permitir interacciones entre modalidades. Se eliminan las codificaciones de posición sinusoidales (ya que las células son conjuntos no ordenados) y se utilizan incrustaciones aprendibles para diferenciar modalidades y condiciones.
• Objetivo de Entrenamiento: Se entrena mediante Flow Matching (Correspondencia de Flujos), aprendiendo un campo de velocidad que transporta el ruido hacia la distribución post-perturbación predicha.

2.3. Proceso de Pre-entrenamiento (Prior Sintético)

Para superar la escasez de datos reales, MapPFN se pre-entrena exclusivamente en datos sintéticos generados a partir de un prior sobre Modelos Causales Estructurales (SCM):

1. Generación de Prior: Se muestrean diversos SCMs (redes de regulación génica con topologías realistas, esparsidad y modularidad) y se simulan dinámicas de expresión génica utilizando el simulador SERGIO (que incluye ruido técnico y no linealidades).
2. Generación de Contexto: Para cada paso de entrenamiento, se genera un contexto $C$ que incluye:
   • Una distribución observacional $Y_{obs}$.
   • Un conjunto de experimentos intervencionistas $\{(t_k, Y^k_{int})\}$ (perturbaciones conocidas).
3. Objetivo: Dado $Y_{obs}$ y el contexto $C$, el modelo debe predecir la distribución post-perturbación $Y^q_{int}$ para una perturbación de consulta no vista $t_q$.
4. Ventaja Causal: Al condicionar en datos intervencionistas, el modelo reduce la clase de equivalencia de Markov de los mecanismos causales subyacentes, mejorando la identificabilidad en comparación con métodos que solo usan datos observacionales.

2.4. Inferencia y Adaptación

En el momento de la inferencia (test-time):

• El modelo recibe un nuevo contexto biológico con un conjunto limitado de experimentos de perturbación.
• Mediante ICL, infiere el mapa de perturbación para nuevos tratamientos sin necesidad de conocer la estructura causal subyacente ni reentrenar el modelo.
• Soporta Test-Time Augmentation (TTA) para escalar a conjuntos de genes más grandes de los vistos en entrenamiento, promediando predicciones sobre subconjuntos de genes.

3. Contribuciones Clave

1. Nuevo Paradigma de Predicción: Enmarca la predicción de perturbaciones como un problema de mapeo de distribuciones con contexto intervencionista en tiempo de prueba, permitiendo la adaptación a contextos biológicos no vistos.
2. Primera PFN para Perturbaciones: Introduce MapPFN, pre-entrenado en un prior biológico sintético de knockouts in silico. Esto permite el meta-aprendizaje de mapas de perturbación a través de diversas estructuras causales, superando la dependencia de grandes volúmenes de datos experimentales reales.
3. Rendimiento Superior: Demuestra que el contexto intervencionista y el prior contrafactual mejoran independientemente las predicciones.
4. Generalización Zero-Shot: Logra recuperar genes diferencialmente expresados (DEG) en datos reales sin ajuste fino, igualando a modelos entrenados desde cero en datos reales.

4. Resultados Experimentales

Los resultados se evaluaron en dos escenarios: SCMs lineales sintéticos y dos conjuntos de datos reales de células individuales (Melanoma y Leucemia).

4.1. En Entornos Sintéticos (SCMs Lineales)

• MapPFN superó a todas las líneas base (CondOT, MFM, etc.) en todas las métricas (distancia de Wasserstein, MMD, RMSE).
• Hallazgo Crítico: Las líneas base mostraron un "colapso de identidad" (predicción de la distribución observacional sin cambios), con una relación de magnitud (MR) cercana a 0.1. MapPFN logró una MR cercana a 1.0, indicando una recuperación precisa del efecto causal.

4.2. En Datos Reales (Células Individuales)

• Zero-Shot (Sin ajuste): Un MapPFN pre-entrenado solo en datos sintéticos recuperó genes diferencialmente expresados (medido por AUPRC) con un rendimiento comparable a los mejores modelos entrenados desde cero en datos reales.
• Ajuste Fino (Fine-tuning): Al ajustar el modelo pre-entrenado en datos reales, MapPFN superó consistentemente a todas las líneas base en ambas líneas celulares (Melanoma y Leucemia) en métricas de distribución, discriminación de perturbaciones y recuperación de genes.
• Escalabilidad: El modelo se adaptó exitosamente a conjuntos de genes más grandes (100 genes vs. 50 en entrenamiento) mediante TTA, mejorando la estabilidad de las predicciones.

4.3. Análisis de Ablación

• Contexto Intervencionista: Eliminar el contexto de perturbaciones (usando solo datos observacionales) degradó significativamente el rendimiento, confirmando que los datos intervencionistas son cruciales para la identificabilidad causal.
• Prior Contrafactual Emparejado: Utilizar datos contrafactuales (donde las células "virtuales" mantienen la misma semilla de ruido antes y después de la perturbación) mejoró la convergencia y el rendimiento final en comparación con datos no emparejados.

5. Significado e Impacto

• Fundamentos para "Células Virtuales": MapPFN representa un paso hacia modelos fundacionales de "células virtuales" que pueden adaptarse a nuevos contextos biológicos y tipos de perturbaciones sin necesidad de reentrenamiento masivo.
• Reducción de Costos: Al permitir la predicción precisa de efectos de perturbaciones en contextos no vistos usando solo un pequeño conjunto de datos de referencia, reduce la necesidad de experimentos costosos y lentos en el laboratorio húmedo (wet lab).
• Generalización Causal: Demuestra que el meta-aprendizaje sobre priores sintéticos diversos puede transferir conocimiento causal a dominios reales complejos, superando las limitaciones de los métodos que solo aprenden de datos observacionales.
• Disponibilidad: El código, el modelo y los datos están disponibles públicamente, fomentando la reproducibilidad y el avance en el campo de la biología computacional.

En conclusión, MapPFN establece un nuevo estado del arte en la predicción de respuestas celulares a perturbaciones, combinando la potencia de los modelos generativos modernos (difusión) con el marco teórico de la inferencia causal y el aprendizaje en contexto.