HypEHR: Hyperbolic Modeling of Electronic Health Records for Efficient Question Answering

Resumo Técnico: HypEHR

1. Problema e Motivação

A resposta a perguntas sobre Registros Eletrônicos de Saúde (EHR-QA) é frequentemente abordada por pipelines baseados em Grandes Modelos de Linguagem (LLMs). Embora precisos, esses métodos são:

• Custosos: Requerem bilhões de parâmetros e grande poder computacional.
• Ineficientes para Privacidade: Difíceis de implantar em ambientes clínicos locais (on-premise) devido a restrições de privacidade de dados.
• Ignorantes da Estrutura: Falham em aproveitar explicitamente a estrutura hierárquica intrínseca dos dados clínicos (ex: ontologias médicas como ICD, trajetórias de pacientes).

A motivação central do trabalho é a observação de que ontologias médicas e trajetórias de pacientes exibem propriedades de geometria hiperbólica (espaços com curvatura negativa), onde estruturas em árvore podem ser embutidas com distorção arbitrariamente baixa, ao contrário dos espaços euclidianos que distorcem essas hierarquias.

2. Metodologia

O HypEHR é um modelo compacto baseado em Lorentz (um modelo de espaço hiperbólico) que embute códigos médicos, visitas e perguntas em um espaço hiperbólico. O framework opera em duas etapas principais:

A. Codificador de Paciente Hiperbólico (Pré-treinamento)

• Embutimento: Cada conceito médico $c$ (códigos de diagnóstico, procedimentos, etc.) é embutido em uma variedade hiperbólica de Lorentz ($H^d_L$).
• Atenção Hiperbólica: Dentro de uma visita, os embeddings dos códigos são agregados via um mecanismo de atenção hiperbólica para gerar uma representação da visita ($h_t$).
• Transformador de Lorentz: A sequência de visitas é processada por um Transformador adaptado à variedade de Lorentz (com atenção, conexões residuais e normalização no espaço hiperbólico), gerando estados contextuais e uma representação global do paciente ($z_{[CLS]}$).
• Objetivos de Treinamento:
  1. Previsão de Diagnóstico: Prever o próximo diagnóstico para minimizar a perda de entropia cruzada binária ($L_{diag}$).
  2. Regularização Hierárquica: Um termo de perda ($L_{hier}$) que força embeddings de códigos que compartilham ancestrais na ontologia ICD a estarem mais próximos no espaço hiperbólico do que códigos não relacionados. Isso é decomposto em termos radiais (ordenação de profundidade) e relativos (separação local).

B. Modelo de Resposta a Perguntas (EHR-QA)

• Codificação da Pergunta: Uma pergunta em linguagem natural é codificada por um encoder biomédico (Euclidiano) e projetada para o manifold hiperbólico via mapas exponenciais.
• Atenção Cruzada Hiperbólica: O modelo calcula a atenção entre a representação da pergunta e as visitas do paciente usando distâncias geodésicas hiperbólicas negativas.
• Agregação: Gera um resumo do paciente condicionado à pergunta ($z_{p|q}^{visit}$) usando uma média de Fréchet hiperbólica.
• Heads Específicos: Dependendo do tipo de resposta (Booleano, Conceito, Numérico, Contagem), cabeças especializadas consomem os vetores de resumo para gerar logits. Para respostas de "Conceito", o modelo utiliza uma atenção de segundo nível sobre os códigos dentro das visitas mais relevantes.

3. Contribuições Principais

1. Novo Framework Compacto: Propõe o HypEHR, um modelo de apenas 22 milhões de parâmetros, que rivaliza com pipelines baseados em LLMs (que possuem trilhões de parâmetros) em tarefas de EHR-QA.
2. Alinhamento Geométrico Explícito: É a primeira abordagem que alinha explicitamente a geometria do modelo (Lorentz) com a estrutura hierárquica intrínseca dos dados de EHR (ICD), demonstrando que a geometria hiperbólica captura melhor a expansão em árvore dos códigos médicos do que a euclidiana.
3. Eficiência e Privacidade: Oferece desempenho comparável aos LLMs com uma fração do custo computacional, tornando-o viável para implantação em hospitais com restrições de privacidade e recursos limitados.
4. Validação Teórica e Empírica: Fornece uma prova teórica de que a hierarquia ICD-10-CM admite um embutimento de baixa distorção no espaço hiperbólico e valida isso empiricamente através de análises de raio vs. profundidade.

4. Resultados Experimentais

Os experimentos foram conduzidos em dois benchmarks baseados em MIMIC-IV: EHRXQA e MIMIC-IV-Ext-Instr.

• Desempenho em QA:
  • O HypEHR alcançou 89.53% de precisão no EHRXQA e 76.02% no MIMIC-Instr.
  • Superou todos os métodos que não utilizam LLMs massivos (como RETAIN e NeuralSQL-l).
  • Competiu diretamente com modelos baseados em LLMs de grande porte (como Llemr e EHRAgent), que são considerados limites superiores (upper bounds) para esses datasets.
• Tarefas Clínicas Adicionais:
  • O encoder de paciente pré-treinado também foi testado em tarefas de previsão clínica (mortalidade, readmissão, tempo de internação, fenótipos).
  • O HypEHR obteve o melhor desempenho em previsão de readmissão e fenótipos, e desempenho comparável ao Llemr nas demais tarefas.
• Estudo de Ablação:
  • A remoção do pré-treinamento ou da perda hierárquica ($L_{hier}$) resultou em quedas significativas de desempenho.
  • Um modelo equivalente em espaço Euclidiano (EucEHR) performou substancialmente pior que o HypEHR, confirmando a vantagem da geometria hiperbólica para este domínio.
• Análise Geométrica:
  • A análise mostrou que, enquanto os raios euclidianos não correlacionam bem com a profundidade da árvore ICD, o modelo de Lorentz exibe um aumento monotônico claro do raio hiperbólico conforme a profundidade da hierarquia aumenta, validando a hipótese de que o espaço hiperbólico aloca mais capacidade para conceitos de granularidade fina.

5. Significância e Limitações

• Significância: O trabalho demonstra que não é necessário recorrer a modelos gigantescos (LLMs) para tarefas complexas de EHR-QA se a estrutura geométrica dos dados for corretamente explorada. Isso abre caminho para sistemas de IA clínica mais eficientes, interpretáveis e privados.
• Limitações:
  • O modelo requer pré-processamento para categorizar respostas em tipos predefinidos (booleano, contagem, etc.), o que adiciona complexidade de engenharia.
  • Atualmente focado em respostas discriminativas, não lidando nativamente com respostas generativas abertas.
  • Redes neurais hiperbólicas são computacionalmente mais complexas que as euclidianas e as bibliotecas públicas ainda são menos maduras, podendo causar instabilidade no treinamento em larga escala.

Em conclusão, o HypEHR representa um avanço significativo ao integrar geometria hiperbólica na modelagem de registros de saúde, oferecendo uma alternativa eficiente e de alto desempenho aos métodos baseados em LLMs.