Tokenization Tradeoffs in Structured EHR Foundation Models

Resumen Técnico: Compromisos de Tokenización en Modelos Fundacionales de EHR

1. El Problema

Los modelos fundacionales para Registros Electrónicos de Salud (EHR) estructurados se entrenan en secuencias longitudinales de eventos clínicos con marca de tiempo. Un paso crítico, pero insuficientemente estudiado, es la tokenización: la conversión de la línea temporal clínica del paciente en una secuencia de tokens discretos para el modelo.

Las decisiones de diseño en la tokenización determinan:

• Qué información se preserva (asociaciones código-valor, intervalos temporales).
• La eficiencia computacional (longitud de la secuencia y tamaño del vocabulario).
• Qué relaciones debe aprender el modelo desde cero versus cuáles están precalculadas.

Existe un compromiso fundamental (trade-off) entre codificar la estructura clínica en el momento de la tokenización (lo que expande el vocabulario pero simplifica el aprendizaje) versus usar una codificación factorizada (vocabulario compacto, pero obliga al modelo a aprender asociaciones complejas a partir de datos limitados). La literatura carecía de una comparación controlada sobre cómo estas elecciones afectan el rendimiento y la eficiencia.

2. Metodología

Los autores diseñaron un experimento factorial 2 × 2 × 2 para evaluar de forma independiente tres ejes de diseño de tokenización:

1. Codificación de Eventos (Event Encoding):
   • Codificación Conjunta (Joint): Un solo token combina el código clínico y sus atributos (ej. "Glucosa-Alta").
   • Codificación Factorizada (Factorized): Tokens separados para el código y los atributos (ej. "Glucosa" + "Alta"), requiriendo que el modelo aprenda la asociación.
2. Codificación del Tiempo (Time Encoding):
   • Posiciones (Time-Positions): Uso de Rotary Positional Embeddings (RoPE) basado en la edad del paciente en días, sin tokens adicionales de tiempo.
   • Tokens (Time-Tokens): Inserción explícita de tokens discretos que representan intervalos de tiempo entre eventos.
3. Anotaciones de Flujo de Trabajo (Workflow):
   • Con Flujo: Representación de eventos clínicos como secuencias de etapas (ej. Orden -> Muestra -> Resultado).
   • Sin Flujo: Representación de un evento como un solo punto temporal (ej. solo el resultado final).

Datos y Entrenamiento:

• Fuente Principal: Datos de EHR de The Hospital for Sick Children (SickKids, Toronto), con 2 millones de pacientes y 169 millones de eventos.
• Evaluación Externa: Cohorte de UCI de adultos de Beth Israel Deaconess Medical Center (MIMIC-IV) para probar la generalización.
• Modelo: Transformer decodificador (28 capas, 768 dimensiones, ~116M parámetros de backbone). Se entrenaron 8 modelos (uno por configuración factorial) con predicción de siguiente token.
• Evaluación: 74 tareas de predicción clínica local (SickKids) y 13 tareas externas (MIMIC), evaluadas mediante el área bajo la curva ROC (AUROC) con probes lineales.

3. Contribuciones Clave

1. Cuantificación de Compromisos: Se estableció que la tokenización no es solo un paso de preprocesamiento, sino un "palanca" tratable que impacta directamente tanto en el rendimiento predictivo como en el costo computacional.
2. Descubrimiento de Eficiencia de Enlace Local: Se demostró que la ventaja de la codificación conjunta se debe a la eficiencia de enlace local (local binding efficiency). Al precalcular la asociación código-valor en un solo token, se evita que el modelo tenga que aprender estas asociaciones contextuales complejas a partir de datos limitados, un problema conocido como el "problema de enlace" en redes neuronales.
3. Generalización Transversal: Se evidenció que las ventajas de la codificación estructural (conjunta) se generalizan entre instituciones (pediatría a adultos) a pesar de un alto desajuste de vocabulario, mientras que los efectos temporales y de flujo de trabajo son específicos de la institución.

4. Resultados Principales

• Rendimiento y Eficiencia Computacional:

  • La Codificación Conjunta superó a la factorizada en 73 de 74 tareas y requirió un 39.5% menos de operaciones de punto flotante (FLOPs) durante el preentrenamiento.
  • La codificación Time-Positions superó a Time-Tokens en 71 de 74 tareas y requirió un 9.6% menos de FLOPs.
  • La inclusión de anotaciones de flujo de trabajo mejoró el rendimiento local pero aumentó el costo computacional en un 35.6% y no mejoró la transferencia externa.

• Mecanismos (Ablaciones):

  • La ventaja de la codificación conjunta desaparece cuando se elimina la necesidad de "enlace" (asociar código con valor), confirmando que el beneficio proviene de resolver el problema de enlace local en el momento de la tokenización.
  • Los tokens de tiempo explícitos (Time-Tokens) degradaron el rendimiento al fragmentar las secuencias y consumir posiciones de contexto sin ganancia informativa proporcional.

• Evaluación Externa (MIMIC):

  • A pesar de una tasa de tokens fuera de vocabulario (OOV) del 69.8% al aplicar el modelo de SickKids a MIMIC, la ventaja de la codificación conjunta se mantuvo con un tamaño de efecto similar.
  • Las codificaciones temporales calibradas a la edad pediátrica no generalizaron bien a la población adulta, indicando que la codificación temporal es sensible a la distribución demográfica.

5. Significado e Implicaciones

Este estudio redefine las mejores prácticas para el desarrollo de modelos fundacionales de EHR:

• Preferencia por Codificación Conjunta: En entornos con datos limitados (típicos en medicina comparado con el lenguaje natural), es más eficiente precalcular las asociaciones semánticas (código + valor) en el token que esperar a que el modelo las aprenda. Esto mejora tanto la precisión como la eficiencia.
• Codificación Temporal: Utilizar la posición relativa (RoPE) es superior a insertar tokens de tiempo explícitos. Sin embargo, la calibración temporal debe ser robusta a cambios demográficos.
• Despliegue Multicéntrico: La tokenización basada en la estructura clínica (codificación conjunta) es más robusta para la transferencia entre instituciones que las características dependientes del flujo de trabajo local o la demografía específica.
• Recomendación: Los autores sugieren que la codificación conjunta de eventos y la codificación posicional del tiempo deberían ser los valores predeterminados para futuros modelos fundacionales de EHR, priorizando la alineación de vocabulario y la eficiencia computacional.

En conclusión, la tokenización es una herramienta crítica para optimizar modelos de IA en salud, permitiendo obtener mejores representaciones de pacientes con menos recursos computacionales y mayor robustez en la generalización.