Federated Learning for ICD Classification with Lightweight Models and Pretrained Embeddings

Resumen Técnico: Aprendizaje Federado para Clasificación ICD con Modelos Ligeros y Representaciones Preentrenadas

Planteamiento del Problema
La clasificación automatizada de notas clínicas en códigos de la Clasificación Internacional de Enfermedades (ICD) es crítica para la investigación epidemiológica y la administración sanitaria, pero sigue siendo un cuello de botella debido al esfuerzo manual requerido. Si bien los Modelos de Lenguaje Grandes (LLM) han demostrado ser prometedores para automatizar esta tarea, su despliegue práctico en el sector sanitario se ve obstaculizado por regulaciones de privacidad (por ejemplo, GDPR) y preocupaciones éticas que impiden la centralización de datos sensibles de los pacientes. Los enfoques tradicionales de Aprendizaje Federado (FL) a menudo requieren el ajuste fino de modelos transformadores masivos en nodos distribuidos, lo cual es computacionalmente costoso, genera una gran carga de comunicación y es difícil de implementar en entornos hospitalarios con recursos limitados. Además, los estudios existentes de FL en el ámbito sanitario suelen centrarse en tareas más simples (por ejemplo, predicción de mortalidad) en lugar de la codificación ICD compleja y multietiqueta.

Metodología
Los autores proponen un pipeline modular de dos etapas que desacopla la representación semántica de la clasificación para abordar los desafíos de privacidad y escalabilidad:

1. Extracción de Representaciones Offline: En lugar de ajustar finamente los LLM a través de instituciones, el estudio utiliza seis modelos de representación de vanguardia y de código abierto (seleccionados del ranking Massive Text Embedding Benchmark) para codificar resúmenes de alta clínica crudos en vectores densos de dimensión fija (768–1536 dimensiones). Este proceso ocurre localmente y de forma offline en cada nodo, garantizando que ningún texto crudo abandone el centro. Los modelos seleccionados incluyen Qwen3-Embed-0.6B, gte-Qwen2-1.5B-instruct, gte-multilingual-base, snowflake-arctic-embed-l/m-v2.0 y nomic-embed-text-v1.5.
2. Clasificación Ligera: Las representaciones resultantes sirven como entrada para clasificadores simples y ligeros de Perceptrón Multicapa (MLP). Se evaluaron tres arquitecturas: un MLP estándar, un DeepMLP (tres capas apiladas con normalización por lotes y dropout) y un DeepResMLP (que incorpora conexiones residuales).
3. Configuración de Entrenamiento Federado: El estudio simula un entorno de FL horizontal utilizando el framework Flower a través de 20 clientes virtuales. Los datos se partitionan del conjunto de datos MIMIC-IV (específicamente resúmenes de alta etiquetados con códigos ICD-9 e ICD-10). La estrategia de entrenamiento emplea FedAvg para la agregación durante 30 rondas de comunicación.
4. Diseño Experimental: Los autores realizaron experimentos extensos comparando el entrenamiento centralizado (línea base) contra el entrenamiento federado. Evaluaron el rendimiento a través de seis modelos de representación y tres arquitecturas de clasificador bajo dos estándares de codificación (ICD-9 e ICD-10). Además, se realizó un estudio de ablación utilizando 10 particiones estratificadas aleatorias para evaluar la estabilidad.

Contribuciones Clave

• Arquitectura que Preserva la Privacidad: El artículo demuestra un pipeline viable donde el texto clínico crudo nunca abandona el nodo local. Al congelar los modelos de representación y entrenar únicamente cabezas de clasificación ligeras de manera federada, el enfoque se adhiere a estrictos requisitos de soberanía de datos.
• Desacoplamiento de Representación y Clasificación: El estudio cuestiona la necesidad del ajuste fino de extremo a extremo de modelos masivos para la codificación ICD. Plantea que representaciones de alta calidad y congeladas pueden proporcionar suficiente riqueza semántica para que clasificadores simples (como MLPs) logren un rendimiento competitivo.
• Evaluación Exhaustiva: El trabajo proporciona una evaluación rigurosa de seis modelos de representación distintos y tres profundidades de clasificador en entornos tanto centralizados como federados, ofreciendo un análisis detallado de cómo la calidad de la representación frente a la complejidad del clasificador impacta el rendimiento multietiqueta.
• Análisis de Escalabilidad: La arquitectura propuesta utiliza modelos con significativamente menos parámetros (menos de 7.1M para el clasificador más grande) en comparación con los transformadores de extremo a extremo de vanguardia (por ejemplo, PLM-ICD con ~139M de parámetros), lo que la hace viable para su despliegue en hardware hospitalario estándar.

Resultados

• Dominio de la Calidad de la Representación: Los experimentos revelan que la elección del modelo de representación es el determinante principal del rendimiento predictivo, superando la complejidad del clasificador. Por ejemplo, Qwen3-Embed-0.6B combinado con un MLP simple logró la puntuación micro F1 más alta para ICD-10 (0.3942), mientras que gte-Qwen2-1.5B-instruct con DeepResMLP obtuvo el mejor rendimiento para ICD-9.
• Complejidad del Clasificador: Las arquitecturas más profundas (DeepMLP, DeepResMLP) no superaron consistentemente al MLP simple. En muchas configuraciones, la complejidad añadida produjo rendimientos decrecientes o incluso caídas en el rendimiento, lo que sugiere que la carga semántica es manejada eficazmente por las representaciones preentrenadas.
• Federado frente a Centralizado: Los resultados del aprendizaje federado coincidieron estrechamente con las líneas base de entrenamiento centralizado. Por ejemplo, la configuración Qwen3-Embed-0.6B + MLP logró una micro F1 de 0.3944 en FL frente a 0.3942 en entrenamiento centralizado para ICD-10, lo que indica que la descentralización de datos no degrada significativamente el rendimiento cuando se utilizan representaciones estandarizadas.
• Rendimiento en Códigos Frecuentes: En los 10 y 50 códigos ICD más frecuentes, los modelos lograron puntuaciones F1 significativamente más altas (a menudo >0.60 para los 10 primeros), demostrando una fuerte capacidad discriminativa para diagnósticos comunes.
• Estabilidad: El estudio de ablación a través de 10 particiones aleatorias mostró una baja varianza en el rendimiento (desviaciones estándar a menudo <0.003 para las puntuaciones F1), confirmando la robustez del enfoque.

Significado y Afirmaciones
Los autores afirman que este trabajo ofrece "una vía viable hacia sistemas de codificación médica escalables y conscientes de la privacidad". Argumentan que se pueden construir sistemas de codificación médica de alto rendimiento sin depender de arquitecturas costosas, rígidas y de extremo a extremo que requieran centralizar datos sensibles. El estudio sugiere que la "paradoja de datos abundantes pero aislados" en el sector sanitario puede abordarse aprovechando representaciones congeladas y estandarizadas localmente y entrenando únicamente cabezas ligeras de manera colaborativa.

El artículo mantiene un tono modesto respecto a sus limitaciones. Reconoce que la dependencia de representaciones estáticas y no ajustadas finamente puede limitar el rendimiento en códigos raros o altamente específicos del dominio. Además, la configuración de FL asume condiciones idealizadas (distribución homogénea de datos, actualizaciones síncronas) que pueden no reflejar completamente la heterogeneidad de las redes hospitalarias reales. Los autores concluyen que, si bien su enfoque logra resultados competitivos con una fracción de los parámetros de los modelos SOTA, el trabajo futuro debería explorar estrategias federadas adaptativas, ajuste fino conjunto y la integración de conocimiento médico específico del dominio para mejorar aún más la robustez.