Resumen Técnico: Grafos de Conocimiento y LLMs de Razonamiento para el Hallazgo de Predictores de Perturbación Transcriptómica Simples pero Efectivos

1. Planteamiento del Problema

Predecir los efectos transcriptómicos de perturbaciones de knockout de genes no observadas sigue siendo un desafío significativo para los modelos de células virtuales. Aunque enfoques recientes aprovechan los grafos de conocimiento biológico para mejorar la extrapolación más allá de las perturbaciones de entrenamiento, los modelos complejos de aprendizaje profundo suelen tener dificultades para superar a los baselines simples (por ejemplo, la media de las muestras de entrenamiento). La dificultad central radica en generalizar a nuevas distribuciones de intervención, $P(x|do(i))$, donde cada perturbación induce una distribución única constreñida por procesos biológicos complejos. Los métodos existentes que utilizan Redes Neuronales de Grafos (GNN) o preentrenamiento a gran escala han mostrado ser prometedores, pero los mecanismos específicos mediante los cuales el conocimiento biológico permite la generalización a nuevas intervenciones siguen sin estar claros. Además, optimizar la estructura causal subyacente directamente es computacionalmente intratable debido al espacio de búsqueda exponencial de posibles grafos de genes.

2. Metodología

2.1 Fundamento Teórico: Un Modelo Jerárquico Disperso

Los autores postulan que las respuestas celulares a los knockouts de genes siguen un grafo causal jerárquico disperso. En este modelo, los genes forman complejos proteicos, los cuales forman vías, que finalmente conducen a orgánulos y respuestas celulares.

• Implicación: Las perturbaciones que afectan al mismo complejo de bajo nivel deberían producir perfiles de expresión similares.
• Predicción: El promedio de la expresión de una perturbación objetivo puede aproximarse mediante un promedio simple de los k-vecinos más cercanos sobre el vecindario local definido por este grafo causal.

2.2 Baseline: El Predictor de Vecinos de Grafo

Los autores evalúan primero un predictor de "Vecino de Grafo", el cual aprovecha los grafos de conocimiento biológico existentes (por ejemplo, StringDB, GO) como priors.

• Mecanismo: Para un gen objetivo $g$, el predictor identifica sus vecinos $N(g)$ en el grafo de conocimiento. La delta de la perturbación predicha ($\hat{\delta}_g$) se calcula como la media de las deltas de entrenamiento de estos vecinos:
  $$\hat{y}_g = \bar{x}_g + \frac{1}{|N_{Tr}(g)|} \sum_{\tilde{g} \in N_{Tr}(g)} \bar{\delta}_{\tilde{g}}$$
  donde $\bar{x}_g$ es la expresión de control emparejada y $N_{Tr}(g)$ representa los vecinos presentes en el conjunto de entrenamiento.
• Hallazgo: Este enfoque simple de $K$-vecinos más cercanos supera a los modelos complejos de aprendizaje profundo (por ejemplo, GEARS, scLAMBDA) y se aproxima al rendimiento de los modelos de estado del arte (SOTA) como TxPert, particularmente en entornos fuera de la distribución (OOD).

2.3 Refinamiento: Aprendizaje por Refuerzo con LLMs de Razonamiento

Reconociendo que los grafos de conocimiento son incompletos y contienen aristas subóptimas, los autores proponen utilizar Modelos de Lenguaje de Gran Escala (LLMs) para refinar estos vecindarios.

• Reducción del Espacio de Búsqueda: En lugar de buscar en el espacio intratable de todos los grafos posibles ($O(2^{\#genes^2})$), la búsqueda se restringe a aristas candidatas plausibles bajo las creencias a priori del LLM.
• Formulación de RL:
  • Política: Un LLM ($\pi_\theta$) actúa como una política para seleccionar acciones (Añadir $A$, Eliminar $R$) a partir de un conjunto inicial de vecinos del grafo de conocimiento ($N_{KG}$).
  • Espacio de Acción: El conjunto refinado es $S_{LLM} = (N_{KG} \setminus R) \cup A$.
  • Función de Recompensa: El modelo se entrena mediante la Optimización de Política Relativa de Grupo (GRPO) para maximizar una recompensa $R(s, a)$ compuesta por:
    1. Recompensa de Formato ($R_{Fmt}$): Asegura que la salida sea procesable (parsable).
    2. Recompensa de Validez ($R_{Valid}$): Incentiva la adición de genes del conjunto de entrenamiento y la eliminación de genes que no estaban en el vecindario original.
    3. Recompensa de Desempeño ($R_\Delta$): Incentiva la mejora en la métrica de Pearson Delta en comparación con el predictor de la línea base del grafo de conocimiento.
• Entrenamiento: El LLM se entrena con datos de perturbación de las líneas celulares de Replogle et al. (2022), con una división estricta de entrenamiento/prueba para asegurar la evaluación sobre perturbaciones no observadas.

3. Contribuciones Clave

1. Demostración de Baselines Simples: El artículo establece que un predictor simple de $K$-vecinos más cercanos basado en grafos de conocimiento biológico es un baseline altamente competitivo para la predicción de perturbaciones OOD, superando a la mayoría de los métodos complejos de aprendizaje profundo.
2. El LLM como Prior de Refinamiento: Los autores demuestran que los LLMs, cuando se optimizan mediante Aprendizaje por Refuerzo (RL), pueden editar eficazmente los vecindarios de los grafos de conocimiento (añadiendo relaciones faltantes y podando las irrelevantes) para mejorar el rendimiento predictivo.
3. Desempeño de Estado del Arte: El LLM ajustado mediante RL alcanza un rendimiento equivalente a los métodos SOTA actuales (específicamente TxPert) en las líneas celulares de Replogle et al. (2022), partiendo de una arquitectura simple de vecinos más cercanos.
4. Generalización de Tareas Aguas Abajo: El estudio muestra que el entrenamiento por RL mejora el desempeño del LLM en tareas posteriores (específicamente la predicción de la expresión diferencial en el benchmark PerturbQA), lo que sugiere que el modelo aprende patrones de razonamiento generalizables en lugar de sobreajustarse a la tarea de predicción específica.

4. Resultados

• Desempeño en Benchmarks: En cuatro líneas celulares (K562, RPE1, HEPG2, Jurkat) de Replogle et al. (2022), el LLM entrenado por RL ($LLM + KG + RL$) logró puntuaciones de Pearson Delta comparables a TxPert (por ejemplo, 0.625 vs 0.626 en K562) y superó significativamente al LLM no entrenado y al baseline del grafo de conocimiento bruto.
• Generalización OOD: El método generaliza a perturbaciones no observadas, manteniendo un alto rendimiento incluso cuando los knockouts de genes específicos no estaban presentes en el entrenamiento.
• Aprendizaje por Transferencia: Cuando se entrenó solo en tres de las cuatro líneas celulares, el modelo aún mejoró el desempeño en la línea celular no utilizada (K562) en comparación con el modelo base, lo que indica una generalización de la política de razonamiento entre líneas celulares.
• Tarea Aguas Abajo: En el conjunto de datos PerturbQA (Wu et al.), el modelo post-entrenado superó al LLM base en la predicción de la expresión diferencial y la dirección del cambio, incluso sin un ajuste fino específico para esta tarea.

5. Significado y Reivindicaciones

El artículo afirma que los grafos de conocimiento biológico sirven como efectivos priors de modelos, y que el modelo más simple que aprovecha estos priors (vecino más cercano) es sorprendentemente robusto. El significado principal radica en mostrar que el Aprendizaje por Refuerzo puede refinar los LLMs para convertirlos en herramientas generalizables para la predicción biológica compleja.

Los autores argumentan que, si bien el entrenamiento por RL es computacionalmente costoso y no es necesariamente el método de despliegue más práctico para la predicción inmediata, sirve como una prueba de concepto:

• Valida que los LLM poseen señales predictivas significativas respecto a las relaciones entre genes más allá de lo que está explícitamente codificado en los grafos de conocimiento estáticos.
• Demuestra que el RL puede ajustar un LLM para que actúe como un posterior sobre los cambios en el vecindario de los genes, aprendiendo efectivamente una política para seleccionar conjuntos de genes predictivos.
• Destaca que las mejoras en el rendimiento predictivo pueden transferirse a tareas de razonamiento aguas abajo, sugiriendo un camino hacia modelos de "célula virtual" más robustos.

Los autores mantienen la modestia, señalando limitaciones tales como el costo computacional del RL y la dependencia de tener un conjunto de entrenamiento para la línea celular específica, lo que impide predicciones de "cero disparos" (zero-shot) reales en líneas celulares nuevas. No pretenden haber resuelto el problema de la generalización por completo, sino que han mostrado un camino viable para refinar modelos interpretables y simples utilizando el razonamiento avanzado de los modelos de lenguaje.