Graph-Constrained Policy Learning for Extreme Clinical Code Prediction

Resumen Técnico: Aprendizaje de Políticas con Restricción de Grafo para la Predicción de Códigos Clínicos Extremos

Definición del Problema
La predicción de códigos clínicos consiste en mapear resúmenes de altas no estructurados a códigos de diagnóstico estandarizados ICD-10-CM. Esta tarea presenta un desafío dual: comprender documentos médicos largos y complejos y realizar una clasificación multietiqueta extrema sobre un espacio de etiquetas que es grande (hasta 15.761 códigos hoja), disperso y profundamente jerárquico. La mayoría de los sistemas automatizados existentes tratan esto como un problema de clasificación multietiqueta plano, puntuando cada código hoja de forma independiente. Este enfoque tiene dificultades con las etiquetas raras debido a las señales de entrenamiento limitadas y no logra aprovechar las restricciones estructurales de la jerarquía ICD, lo que a menudo conduce a predicciones estructuralmente inválidas o a un rendimiento deficiente en la larga cola de los códigos raros.

Metodología
Los autores proponen reformular la predicción de ICD como un problema de recorrido con restricción de grafo, modelado como un proceso de decisión de Markov (MDP) de horizonte finito sobre una jerarquía ICD-10-CM podada.

• Formulación del MDP:

  • Estado: Comprende la nota de alta, la frontera actual de nodos pendientes de expansión y el conjunto de códigos hoja facturables registrados hasta el momento.
  • Acción: La política selecciona un subconjunto de nodos hijos válidos de la frontera actual que estén respaldados por la nota.
  • Restricción: Un entorno ligero impone restricciones de validez estrictas. El modelo solo puede seleccionar hijos del conjunto expuesto por el nodo actual, asegurando que los códigos estructuralmente inválidos reciban probabilidad cero por construcción.
  • Terminación: El proceso continúa nivel por nivel hasta que se alcanzan los códigos hoja facturables y la frontera queda vacía.

• Implementación:

  • Arquitectura: Un único modelo de lenguaje de solo decodificador (Qwen3-4B-Instruct, cuantizado en 4 bits con QLoRA) actúa como la política de recorrido. Opera dentro de una ventana de contexto continua, condicionando cada decisión al historial completo acumulado del recorrido.
  • Entorno: El modelo interactúa con un entorno basado en herramientas utilizando expand_level() para revelar hijos y select() para elegir nodos válidos. Esto convierte el problema de clasificación multietiqueta extrema en una secuencia corta de decisiones de subconjuntos dispersos y con restricción de grafo.
  • Entrenamiento: El enfoque principal es el Ajuste Fino Supervisado (SFT) sobre "trayectorias de oro" —secuencias deterministas generadas mediante la reproducción de las rutas de ancestros correctas de los códigos de oro a través del entorno—. El modelo es entrenado para imitar estas demostraciones de oráculo.
  • Comparaciones: El estudio evalúa esta política compartida contra:
    1. Baselines Planos: CAML, LAAT y PLM-ICD (reentrenados en los mismos cortes).
    2. Baseline de Cascada (SFT-3): Tres modelos especialistas independientes (capítulo, subsección, código) que operan de forma unidireccional por nivel.
    3. Aprendizaje por Refuerzo (GRPO): Optimización de Política Relativa de Grupo (GRPO) aplicada a la política compartida.

Contribuciones Clave y Diseño Experimental
El artículo realiza un estudio factorial controlado en el conjunto de datos MIMIC-IV (122.197 notas validadas) para aislar los efectos de la arquitectura (política compartida vs. cascada), el algoritmo de entrenamiento (SFT vs. GRPO) y el presupuesto de datos (2.000 vs. 5.000 trayectorias de oro).

• Espacios de Etiquetas: Los experimentos se ejecutan en dos escalas: un subconjunto "Frecuente" (50 códigos más comunes) y la configuración "Todos" (espacio completo de 15.761 códigos).
• Métricas: Se utilizan las puntuaciones F1 promedio micro y macro, junto con el Path-F1 (crédito parcial jerárquico) y el análisis de recuperación por profundidad.

Resultos

• Rendimiento vs. Baselines Planos: La política de recorrido con restricción de grafo supera significativamente a los baselines planos, particularmente en la configuración "Todos".
  • En el espacio completo de 15.761 códigos, el mejor modelo (SFT-1+) logra un 0.527 micro-F1 y un 0.256 macro-F1, superando al baseline plano más fuerte (LAAT) por +0.044 micro-F1 y +0.157 macro-F1.
  • Los baselines planos (CAML, PLM-ICD) colapsan en el espacio completo (macro-F1 ~0.002), lo que indica una incapacidad para discriminar entre 15.761 candidatos sin guía estructural.
• Escalado de Datos: Aumentar los datos de entrenamiento supervisado de 2.000 a 5.000 notas es la única intervención que mejora el rendimiento de manera consistente y significativa. Las ganancias se duplican aproximadamente al pasar de la configuración de 50 códigos a la de espacio completo.
• Arquitectura: Una única política de recorrido compartida iguala la precisión de la cascada de tres especialistas, pero evita un fallo estructural crítico: la cascada sufre de desbordamiento de la ventana de contexto (truncamiento) en el 28–32% de las notas de prueba en la configuración de espacio completo debido a las grandes listas de candidatos en el nivel de hoja. La política compartida tiene un 0% de truncamiento.
• Aprendizaje por Refuerzo: GRPO no proporciona ningún beneficio sobre la continuación supervisada cuando se entrena con la misma cantidad de datos. El modelo de RL iguala su inicio cálido supervisado y cae significativamente por debajo del modelo entrenado con trayectorias supervisadas adicionales.

Significancia y Reclamaciones
El artículo argumenta que, para la predicción de códigos clínicos extremos, la simplicidad y la calidad de los datos superan a la complejidad arquitectónica.

1. Descomposición con Restricción de Grafo: Reformular el problema como un recorrido restringido mitiga eficazmente el "cuello de botella de los códigos raros" al reducir el espacio de búsqueda en cada paso y asegurar la validez estructural, ofreciendo ganancias sustanciales sobre los métodos de puntuación plana.
2. Supervisión sobre RL: Contrario a la tendencia de utilizar aprendizaje por refuerzo para tareas complejas, el estudio encuentra que simplemente añadir más trayectorias supervisadas de alta calidad es más efectivo que optimizar las recompensas mediante GRPO.
3. Política Unificada vs. Cascadas: Se prefiere una única política de recorrido compartida sobre los modelos especialistas en cascada. Logra una precisión comparable mientras elimina las limitaciones de escalabilidad (desbordamiento de la ventana de contexto) inherentes al enfoque de cascada.
4. Recomendación Práctica: Los profesionales deben desplegar una única política de recorrido compartida e invertir recursos en la adquisición de más datos de anotación supervisada en lugar de perseguir arquitecturas de cascada complejas o estrategias de aprendizaje por refuerzo.

Los autores señalan las limitaciones, incluyendo el uso de un único modelo de 4B, un diseño de recompensa específico para RL y un presupuesto de anotación relativamente pequeño (5.000 notas) en relación con los datos disponibles, lo que sugiere que se necesitan más estudios de escalado para determinar los puntos de saturación.