Resumen Técnico: Recompensas de Distancia de Conjuntos para la Generación de Informes Radiológicos (SDR)

Planteamiento del Problema

La generación de informes de radiografías de tórax mediante modelos de lenguaje-visión (VLM) presenta un desafío único para el post-entrenamiento mediante Aprendizaje por Refuerzo (RL), específicamente en lo que respecta al diseño de las funciones de recompensa. Los enfoques estándar de RL para tareas de razonamiento suelen depender de modelos de recompensa de proceso que verifican los pasos intermedios en una cadena causal. Sin embargo, los informes radiológicos no siguen una estructura de razonamiento secuencial o causal; en su lugar, consisten en hallazgos desordenados y ortogonales (por ejemplo, "el tamaño del corazón puede estar aumentado" y "el pulmón derecho está despejado" son observaciones independientes).

En consecuencia, las recompensas discretas estándar basadas en la precisión de coincidencia exacta son incompatibles con este dominio. Estas inducen ruido al penalizar informes que contienen hallazgos correctos pero en un orden diferente o con ligeras variaciones de fraseo, fallando en capturar la realidad clínica de que el conjunto de hallazgos importa más que su secuencia. Además, las anotaciones a nivel de paso requeridas para las recompensas de proceso a menudo no están disponibles, y verificar cada paso mediante fuentes externas es computacionalmente costoso.

Metodología

Los autores proponen las Recompensas de Distancia de Conjuntos (SDR), un marco que trata los informes de radiografías de tórax como conjuntos desordenados de incrustaciones (embeddings) de oraciones en lugar de texto secuencial. La metodología opera en dos fases: post-entrenamiento (GRPO) y escalado en tiempo de inferencia.

1. Representación a Nivel de Oración

Un informe de radiografía de tórax $y$ se divide en dos secciones: Hallazgos ($y_F$) e Impresión ($y_I$). Cada sección se segmenta en oraciones individuales. Un transformador de oraciones pre-entrenado y congelado (específicamente all-mpnet-base-v2) incrusta cada oración $s$ en un vector de dimensión fija $e \in \mathbb{R}^d$. El informe se representa así como dos conjuntos desordenados de incrustaciones:
$$E_F(y) = \{E_\phi(s^F_i) \mid 1 \le i \le n_F\}, \quad E_I(y) = \{E_\phi(s^I_j) \mid 1 \le j \le n_I\}$$
Esta representación es invariante a la permutación, alineándose con la intuición clínica de que el orden de los hallazgos no tiene significado diagnóstico.

2. Recompensas de Distancia de Conjunto a Conjunto (Post-Entrenamiento)

Durante la Optimización de Política Relativa de Grupo (GRPO), el modelo genera un informe $\hat{y}$. La recompensa semántica se calcula computando la distancia entre el conjunto de incrustaciones generadas y el conjunto de incrustaciones de referencia (verdad de campo o ground truth).

• Métricas de Distancia: Los autores utilizan distancias simétricas de conjunto a conjunto normalizadas al intervalo $[0, 1]$, específicamente la distancia de Chamfer y la distancia de Hausdorff, construidas sobre la distancia de coseno entre incrustaciones de norma unitaria.
  • Chamfer: Promedia los costos de los vecinos más cercanos, actuando como una recompensa de "cobertura suave".
  • Hausdorff: Toma el costo máximo de los vecinos más cercanos, actuando como una penalización de "peor caso" para los hallazgos críticos omitidos.
• Formulación de la Recompensa: La recompensa total $R(\hat{y}, y)$ es una suma ponderada de una recompensa de formato binaria (que asegura que la salida siga una plantilla estructurada) y la recompensa semántica continua derivada de la distancia de conjunto:
  $$R(\hat{y}, y) = \lambda_{fmt} R_{fmt}(\hat{y}) + \lambda_{sem} (1 - D(E_S(\hat{y}), E_S(y)))$$
  donde $D$ es la distancia de conjunto. Esto proporciona una señal continua y diferenciable que recompensa la corrección parcial y gestiona los hallazgos desordenados.

3. Escalado en Tiempo de Inferencia y Poda

Las mismas métricas de distancia de conjunto se aplican en el tiempo de inferencia sin actualizaciones de gradiente:

• Selección Best-of-N: Para una imagen dada, el modelo genera $K$ informes candidatos. El conjunto de incrustaciones de cada candidato se compara contra los conjuntos de incrustaciones pre-calculados del corpus de entrenamiento. El candidato con la distancia mínima a la distribución de entrenamiento es seleccionado. Esto sirve como un prior libre de gradientes para salidas clínicamente plausibles.
• Poda por Flujo (Streaming Pruning): Para reducir el costo computacional, el sistema calcula la distancia de conjunto continua de las generaciones parciales durante la decodificación. Los candidatos cuya distancia parcial excede un umbral son podados tempranamente. Este bucle de "podar-luego-decodificar" reduce la generación de tokens en más del 50% manteniendo una calidad comparable a la selección completa de Best-of-N.

Contribuciones Clave

1. Funciones de Recompensa de Distancia de Conjuntos para GRPO: El artículo aborda la inviabilidad del modelado de recompensa de proceso para la radiología al tratar los informes como conjuntos desordenados. Utilizando las distancias de Chamfer y Hausdorff como recompensas continuas, el método supera consistentemente al Ajuste Fino Supervisado (SFT) y al GRPO de coincidencia exacta en tres backbones de VLM (Qwen3-VL-2B/4B, Gemma3-4B) y dos conjuntos de datos (MIMIC-CXR, ReXGradient).
2. Selección Guiada por Distancia de Conjunto en Tiempo de Prueba: Los autores demuestran que las mismas métricas de distancia de conjunto pueden servir como criterio de selección tanto para modelos ajustados como para LLMs generalistas de código cerrado (Mistral-Small, Gemini-2.5 Flash-Lite, GPT-4o-mini, GPT-5 mini). Este enfoque mejora el rendimiento sobre la selección aleatoria sin requerir acceso a los gradientes del modelo o un ajuste fino específico del dominio.
3. Poda Basada en Distancia sobre la Marcha: El artículo introduce un método para podar candidatos con baja puntuación a mitad de la generación. Esto reduce el número total de tokens generados en más del 50% mientras preserva la calidad de la selección final, ofreciendo una forma más eficiente de escalado en tiempo de prueba.

Resultados Experimentales

La evaluación cubre métricas principales incluyendo BERTScore, RadGraph F1 y CheXbert F1.

• Rendimiento en Post-Entrenamiento:
  • En el conjunto de datos ReXGradient, la recompensa basada en Chamfer logró el mayor rendimiento en todos los modelos y métricas.
  • En MIMIC-CXR, tanto las recompensas de Chamfer como las de Hausdorff alcanzaron el mayor rendimiento promedio.
  • Las mejoras relativas sobre SFT y GRPO de coincidencia exacta fueron significativas: mejoras promedio de 6.80% (BERTScore), 7.82% (RadGraph F1) y 4.45% (CheXbert F1).
• Selección en Tiempo de Inferencia:
  • La aplicación de la selección por distancia de conjunto a LLMs de código cerrado produjo ganancias sustanciales sobre la selección aleatoria. Por ejemplo, en BERTScore, las mejoras oscilaron entre +11.2% y +18.7% dependiendo del modelo y la plantilla de prompt.
  • El método mejoró con éxito el rendimiento de modelos generalistas (ej. GPT-4o mini, Gemini) que no fueron ajustados finamente para el dominio médico específico.
• Eficiencia de la Poda:
  • La estrategia de poda ahorró entre el 42.1% y el 60.1% de los tokens generados.
  • Las generaciones podadas mantuvieron una calidad comparable a la selección completa de Best-of-N, con mejoras relativas promedio sobre la selección aleatoria de +12.7% (BS-F1), +17.1% (RG-F1) y +6.2% (CXB-F1).

Significado y Reivindicaciones

El artículo afirma establecer las recompensas de distancia de conjuntos como una señal unificada tanto para el post-entrenamiento como para el escalado en tiempo de prueba en la generación de informes de radiografía de tórax.

• Señal Unificada: La misma formulación matemática (distancia de conjunto a conjunto) impulsa eficazmente el proceso de post-entrenamiento por RL y sirve como un criterio de selección robusto y libre de gradientes en el tiempo de inferencia.
• Adaptación al Dominio: Al reconocer la naturaleza desordenada de los hallazgos radiológicos, el método evita el ruido introducido por las restricciones secuenciales, haciéndolo específicamente adecuado para la generación de informes médicos donde los hallazgos suelen ser independientes.
• Eficiencia: El trabajo demuestra que la generación de informes de alta calidad no requiere necesariamente generar todos los tokens para todos los candidatos; la poda temprana basada en la distancia de conjunto ofrece una vía viable para reducir los costos de inferencia.
• Generalizabilidad: La capacidad de aplicar estas reglas de selección a LLMs de código cerrado sugiere un camino para mejorar los resultados de la IA clínica sin necesidad de un costoso ajuste fino del dominio o acceso a los pesos del modelo.

Los autores concluyen que su enfoque proporciona una señal de recompensa continua y práctica que se alinea con la estructura clínica de los informes de radiografía, superando los mecanismos tradicionales de recompensa discreta o secuencial.