AdaPCLA: Adaptive Prior-Calibrated Logit Adjustment for Long-Tailed Longitudinal EHR Generation

Resumen Técnico: AdaPCLA para la Generación de EHR Longitudinales de Cola Larga

Declaración del Problema

Los Registros Electrónicos de Salud (EHR) longitudinales son fundamentales para aplicaciones clínicas como la predicción de riesgos y la recomendación de medicamentos. Sin embargo, el intercambio de datos de EHR del mundo real se ve obstaculizado por las regulaciones de privacidad, lo que hace necesaria la generación de datos sintéticos de alta calidad. Los modelos generativos existentes (por ejemplo, GAN, VAE, modelos autorregresivos) tienen dificultades con la distribución de cola larga inherente a los códigos médicos. Mientras que los patrones clínicos comunes son frecuentes, las condiciones raras pero clínicamente significativas (eventos de cola) proporcionan señales de aprendizaje dispersas. Los modelos autorregresivos estándar tienden a optimizar los patrones dominantes de la "cabeza", lo que provoca que subrepresenten la estructura de coocurrencia de los eventos raros. Esto resulta en datos sintéticos que parecen realistas en promedio, pero que no logran preservar la fidelidad y la veracidad requeridas para las subpoblaciones raras, limitando su utilidad downstream y su adaptabilidad entre poblaciones.

Metodología: AdaPCLA

Los autores proponen AdaPCLA (Ajuste de Logit Calibrado por Prior Adaptativo), un marco basado en currículo diseñado para internalizar el conocimiento de la distribución de los datos en los parámetros del modelo a través de un proceso de recocido simulado (simulated annealing). El marco consta de tres componentes principales:

1. Backbone del Generador de EHR: Un Transformer autorregresivo que modela secuencias de visitas de pacientes como vectores multi-hot de códigos médicos. Captura dependencias temporales y entre eventos mediante estados ocultos compartidos.
2. Modelado de Prior Estadístico (El Andamio): Para abordar la dispersión de los códigos raros, el método construye un término de sesgo por etiqueta ($b_{stat}$) derivado de las frecuencias marginales empíricas. Este sesgo desplaza los logits efectivos de los códigos raros hacia cero (la región de alta sensibilidad de la función sigmoide) durante las etapas tempranas del entrenamiento, amplificando las actualizaciones de gradiente de las ocurrencias positivas dispersas en relación con los códigos dominantes de la cabeza.
3. Inyección de Prior por Currículo (Recocido): A diferencia de los métodos de ajuste de logit estáticos que dependen de correcciones externas durante la inferencia, AdaPCLA trata el sesgo del prior como un andamio de optimización temporal. Un coeficiente de programación $\alpha(\tau)$ reduce gradualmente el sesgo de 1 a 0 a lo largo de los pasos de entrenamiento.
   • Entrenamiento: El modelo optimiza utilizando logits efectivos $\tilde{z} = z + \alpha(\tau)b_{stat}$.
   • Inferencia: El sesgo se elimina por completo ($\alpha=0$), obligando al modelo a depender únicamente de sus parámetros internos aprendidos para generar eventos raros, asegurando una inferencia libre de sesgos.

Adaptación de Zero-Shot entre Poblaciones:
AdaPCLA permite la adaptación sin entrenamiento a nuevas poblaciones con diferentes frecuencias de códigos. Al aplicar una corrección de prior objetivo ($\Delta b$) a los logits en el momento de la inferencia, el modelo puede ajustar las probabilidades de generación para coincidir con una distribución objetivo sin necesidad de reentrenamiento, siempre que las dependencias condicionales permanezcan consistentes.

Análisis Teórico

El artículo proporciona una justificación teórica de la eficacia del método:

• Transferencia de Gradiente: Los autores analizan cómo el andamio del prior afecta las actualizaciones intrínsecas de los logits a través del Kernel de Tangente Neuronal (NTK) empírico por etiqueta. El andamio reformatea los residuales de la Entropía Cruzada Binaria (BCE), y estos efectos se transfieren a los logits intrínsecos a través de la estructura del NTK, fortaleciendo las señales de aprendizaje de los códigos raros en las etapas iniciales.
• Límite de Internalización del Prior: Los autores derivan un límite que muestra que el error entre los logits intrínsecos aprendidos y el objetivo móvil (a medida que se elimina el andamio) depende de la velocidad de recocido y del autovalor mínimo del NTK empírico. Un esquema de recocido más suave y un NTK bien condicionado aseguran que las señales del prior se internalicen con éxito en los parámetros del modelo en lugar de perderse cuando se retira el andamio.

Contribuciones Clave

1. Propuesta de Marco: Introducción de AdaPCLA, un marco de currículo guiado por prior que fortalece el aprendizaje de códigos raros al tiempo que permite la inferencia libre de sesgos y el control de distribución zero-shot.
2. Perspectiva Teórica: Caracterización de las actualizaciones de logit de códigos raros mediante la dinámica del NTK empírico por etiqueta y la derivación de un límite de internalización de prior que vincula la velocidad de recocido y el condicionamiento del NTK con la retención de la señal.
3. Validación Empírica: Demostración de mejoras significativas en la plausibilidad de la cola, la utilidad downstream y la adaptación entre poblaciones utilizando conjuntos de datos del mundo real.

Resultados Experimentales

Los experimentos se realizaron en los conjuntos de datos MIMIC-III y MIMIC-IV, comparando AdaPCLA contra baselines que incluyen modelos tipo GPT, LSTM, EVA, SynTEG y HALO.

• Plausibilidad Clínica de la Cola: AdaPCLA logró el mejor desempeño en métricas enfocadas en la cola.
  • En MIMIC-III, mejoró TailPairSeen en un 114.2% respecto a HALO (0.1114 vs. 0.0520).
  • En MIMIC-IV, mejoró TailPairSeen en un 65.1% respecto a HALO (0.0822 vs. 0.0498).
  • También demostró una consistencia de contexto de cola superior (menor TailCtxJSD) en comparación con los baselines.
• Utilidad Downstream: Utilizando el protocolo Train-on-Synthetic, Test-on-Real (TSTR), AdaPCLA alcanzó los puntajes más altos de Accuracy, F1 y AUPRC para la clasificación de diagnóstico de 25 etiquetas en ambos conjuntos de datos, superando al baseline de tipo GPT por un 3.5% de F1 en la adaptación zero-shot entre poblaciones.
• Estudios de Ablación:
  • Sesgo Estático vs. Recocido: Los modelos que utilizan un sesgo fijo durante el entrenamiento pero lo eliminan en la inferencia (Static-PCLA Train-only) fallaron completamente (TailPairSeen cayó casi a cero), confirmando que el sesgo por sí solo no enseña la estructura.
  • Internalización: El enfoque de recocido de AdaPCLA internalizó con éxito el prior, logrando las mejores métricas sin sesgo externo en la inferencia.
  • Control Zero-Shot: El modelo se adaptó con éxito de MIMIC-IV a MIMIC-III utilizando solo correcciones de prior objetivo, logrando un AUPRC de 0.922 sin reentrenamiento.

Significancia y Reivindicaciones

El artículo afirma que AdaPCLA resuelve el compromiso entre el ajuste de la distribución general de los datos y la preservación de las dependencias de eventos raros. Al internalizar el conocimiento del prior a través de un proceso de recocido de currículo, el método genera EHR sintéticos de alta fidelidad que son fiables para las subpoblaciones raras. Esto es significativo porque permite la investigación que preserva la privacidad y facilita la transferencia de modelos generativos a nuevos entornos clínicos (diferentes hospitales o poblaciones) sin la necesidad de un reentrenamiento costoso, abordando un cuello de botella crítico en el desarrollo de la IA médica.

Limitaciones: Los autores señalan que la evaluación actual está limitada a secuencias de códigos de diagnóstico de dos benchmarks públicos. Se espera que el trabajo futuro valide el método en diversos hospitales y modalidades de EHR más ricas más allá de los códigos de diagnóstico.