Resumo Técnico: Modelagem Hierárquica de Códigos CID em Modelos de Fundação de EHR

Declaração do Problema

Modelos de fundação de Registros Eletrônicos de Saúde (EHR), como o BEHRT e o Med-BERT, tipicamente tratam os códigos de diagnóstico da Classificação Internacional de Doenças (CID) como tokens planos e independentes. Essa abordagem negligencia a estrutura hierárquica clinicamente significativa inerente aos sistemas de codificação CID, que organiza doenças em famílias, subcategorias e diagnósticos cada vez mais específicos. Consequentemente, os métodos de aprendizado de representação existentes falham em explorar explicitamente essa ontologia, tratando diagnósticos relacionados como símbolos isolados. Isso é particularmente problemático em dados clínicos onde os diagnósticos são frequentemente esparsos, de cauda longa e específicos de cada instituição. O artigo hipotetiza que modelar explicitamente essa hierarquia como um viés indutivo poderia melhorar o aprendizado de representação ao permitir que códigos raros se beneficiem de dados associados a códigos clinicamente semelhantes, capturando trajetórias de cuidado mais amplas e melhorando a robustez sob mudanças de domínio.

Metodologia

Os autores propõem e avaliam dois mecanismos complementares para injetar a hierarquia do CID-10-CM no aprendizado de representação contemporâneo de EHR, coletivamente denominados modelos de CID Hierárquico (HICD). Eles definem três níveis de hierarquia aninhados ($G_0, G_1, G_2$) variando da granularidade mais grosseira para a mais fina.

1. HICD-BERT (Injeção ao Nível de Token)

Esta abordagem estende a arquitetura Transformer BEHRT. Em vez de tratar os códigos CID como identificadores atômicos, o modelo aumenta os embeddings de token padrão com embeddings específicos de hierarquia derivados via truncamento de prefixo de string.

• Mecanismo: Para um código de diagnóstico, o modelo adiciona embeddings aditivos correspondentes a diferentes níveis da string do código (por exemplo, prefixo de 1 caractere, prefixo de 2 caracteres, categoria de 3 caracteres).
• Fonte de Dados: Este método é orientado por dados, não requerendo acesso à ontologia externa durante o treinamento; a hierarquia é derivada diretamente das strings dos códigos.
• Treinamento: Pré-treinado usando Modelagem de Linguagem Mascarada (MLM) em sequências de diagnósticos, seguido de ajuste fino (fine-tuning) para tarefas de downstream.

2. HICD-Graph (Codificação ao Nível de Grafo)

Esta abordagem injeta a hierarquia ao nível do grafo de codificação antes de inicializar um codificador Transformer.

• Mecanismo: Um grafo híbrido de coocorrência de diagnósticos é construído. Ele combina:
  • Arestas empíricas: Pesadas por Informação Mútua de Ponto a Ponto (PMI) baseada na coocorrência de diagnósticos dentro de visitas de pacientes.
  • Arestas ontológicas: Arestas de peso uniforme conectando nós de diagnóstico folha aos seus nós ancestrais na hierarquia CID ($G_0, G_1, G_2$).
• Treinamento: Uma Rede Neural Convolucional de Grafos (GCN) é treinada com um objetivo de predição de ligação para aprender embeddings de códigos conscientes da hierarquia. Esses embeddings então inicializam a tabela de embedding de um codificador Transformer de nível de paciente.

Configuração Experimental

• Datasets: MIMIC-IV (usado para pré-treinamento e avaliação in-domain) e eICU (usado para avaliação de transferência entre datasets).
• Tarefas:
  1. Readmissão de 30 Dias (MIMIC-IV): Classificação binária.
  2. Predição de Admissão de Emergência (MIMIC-IV): Classificação binária.
  3. Readmissão em UTI (eICU): Classificação binária para avaliação de transferência (encoder congelado, classificador ajustado).
• Ablação: O estudo avalia sistematicamente todas as oito combinações dos três níveis de hierarquia ($G_0, G_1, G_2$) para ambas as arquiteturas contra uma linha de base sem hierarquia.

Principais Resultados

Desempenho In-Domain (MIMIC-IV)

• Melhoria Geral: Configurações aumentadas pela hierarquia superaram a linha de base sem hierarquia em 26 de 28 comparações pareadas em ambos os modelos e tarefas.
• HICD-BERT: Os ganhos de desempenho foram impulsionados principalmente pelo nível de granularidade mais fina ($G_2$). O nível intermediário ($G_1$) sozinho não melhorou significativamente o desempenho, mas contribuiu quando combinado com outros níveis. A configuração ideal variou por tarefa (ex: todos os níveis para readmissão, nível fino único para predição de emergência).
• HICD-Graph: Mostrou uma tendência de melhoria monotônica onde a adição de mais níveis de hierarquia aumentou progressivamente o desempenho. A arquitetura baseada em grafos integrou efetivamente informações de múltiplos níveis de granularidade simultaneamente.

Transferência entre Datasets (MIMIC-IV para eICU)

• HICD-Graph: Demonstrou transferibilidade robusta. Seis de sete variantes de hierarquia melhoraram significamente sobre a linha de base, com a configuração de todos os níveis alcançando o AUROC mais forte (0,587). A estrutura relacional baseada em grafos pareceu mais invariante a mudanças de dataset.
• HICD-BERT: Os benefícios da hierarquia desapareceram amplamente sob mudança de distribuição. Apenas a configuração de nível único mais grosseira ($G_0$) mostrou uma melhoria significativa sobre a linha de base. Os autores sugerem que os padrões de tokens sequenciais aprendidos pelo HICD-BERT estão mais fortemente acoplados às estruturas de visita específicas do MIMIC-IV, enquanto a estrutura de grafo ancorada na ontologia é mais generalizável.

Análise de Embedding

• Agrupamento Semântico: Modelos conscientes da hierarquia produziram clusters de códigos mais coerentes. Especificamente, a configuração $G_1+G_2$ no HICD-Graph resultou nos clusters de códigos locais mais apertados no espaço de embedding.
• Correlação com Desempenho: As configurações que produziram os clusters clinicamente significativos mais apertados (ex: códigos compartilhando o mesmo prefixo de 2 caracteres) também alcançaram o desempenho preditivo de downstream mais forte, sugerindo que a hierarquia melhora o desempenho ao moldar a geometria do embedding para respeitar relações clínicas.

Significância e Alegações

O artigo alega que a estrutura hierárquica presente nos sistemas de codificação CID é um sinal subutilizado em modelos de fundação de EHR. As principais contribuições e alegações incluem:

1. Hierarquia como um Viés Indutivo Robusto: Codificar explicitamente a hierarquia CID melhora significativamente o desempenho de predição de downstream em diferentes arquiteturas (baseadas em Transformer e baseadas em Grafos) e tarefas.
2. Dependência de Granularidade: A utilidade dos níveis de hierarquia não é uniforme; o nível ideal depende tanto da abordagem de modelagem (nível de token vs. nível de grafo) quanto da tarefa de downstream.
3. Transferibilidade de Estrutura: A codificação de hierarquia baseada em grafos transfere-se de forma mais robusta entre datasets do que a injeção de nível de token, sugerindo que as estruturas relacionais derivadas da ontologia capturam relações clínicas mais fundamentais do que os padrões de tokens sequenciais.
4. Integração de Baixo Custo: Os benefícios são alcançados com modificações arquitetônicas mínimas (embeddings aditivos ou arestas de grafo adicionais) sem alterar os objetivos de pré-treinamento ou pipelines principais.
5. Generalização: Os insights sobre o valor das ontologias clínicas estruturadas estendem-se além do CID-10-CM e podem ser aplicados a outros vocabulários estruturados (ex: códigos ATC para medicamentos) e métodos de aprendizado de representação.

Os autores permanecem modestos, observando que, embora as melhorias absolutas de AUROC sejam modestas, elas são estatisticamente robustas e consistentes. Eles reconhecem limitações quanto à generalização para outros desfechos clínicos, instituições e sistemas de EHR, e observam que seus níveis de hierarquia foram definidos manualmente em vez de aprendidos automaticamente.