DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning

O artigo apresenta o DT-BEHRT, uma arquitetura de Transformer aprimorada por grafos que utiliza modelagem de trajetórias de doenças e pré-treinamento orientado por ontologia para gerar representações de pacientes interpretáveis e robustas a partir de registros eletrônicos de saúde.

O essencial

Imagine que o prontuário eletrônico de um paciente (o histórico médico digital) é como um romance gigante e bagunçado, escrito por dezenas de autores diferentes (médicos, enfermeiros, laboratórios) ao longo de muitos anos.

O problema é que, para tentar prever o futuro de saúde desse paciente (se ele vai ter uma recaída, se precisa de mais um dia no hospital, etc.), os computadores antigos tentavam ler esse livro como se fosse uma lista de compras simples. Eles olhavam para os códigos médicos (como "diabetes" ou "fratura") e tentavam adivinhar o que viria a seguir, mas muitas vezes perdiam a história e a conexão entre as coisas.

É aqui que entra o DT-BEHRT, a nova solução proposta pelos pesquisadores. Vamos explicar como ele funciona usando algumas analogias do dia a dia:

1. O Problema: A Lista de Compras vs. A História de Vida

Os métodos antigos tratavam todos os códigos médicos da mesma forma. Era como se, em uma receita de bolo, eles tratassem "ovos", "farinha" e "forno" como se tivessem a mesma importância e a mesma ordem.

• A realidade: Alguns códigos são como o motor do carro (os diagnósticos, como "diabetes" ou "hipertensão"). Eles dirigem a história da saúde do paciente. Outros são como o combustível ou a manutenção (medicamentos e exames), que são importantes, mas seguem o motor.
• O erro: Os modelos antigos misturavam tudo, sem entender que o "diagnóstico" é o que realmente define para onde a saúde do paciente está indo.

2. A Solução: O Detetive de Trajetórias (DT-BEHRT)

O DT-BEHRT é como um detetive médico superinteligente que não apenas lê a lista, mas entende a trajetória da doença. Ele usa três "superpoderes" (módulos) para entender o paciente:

A. O Organizador de Sistemas (Módulo de Agregação de Doenças)

Imagine que o corpo humano é uma cidade com vários bairros (sistemas): o bairro do Coração, o bairro dos Pulmões, o bairro do Fígado.

• O DT-BEHRT não olha apenas para um código isolado. Ele agrupa os problemas por "bairro".
• Se o paciente tem vários problemas nos pulmões, o modelo cria um "resumo do bairro dos pulmões". Isso ajuda a entender que, embora os códigos sejam diferentes, eles pertencem à mesma história de doença. É como dizer: "Ah, todos esses sintomas são sobre o sistema respiratório, vamos tratá-los como um grupo".

B. O Cronista do Tempo (Módulo de Progressão da Doença)

A saúde não é estática; ela muda com o tempo.

• Este módulo funciona como um filme em time-lapse. Ele conecta as visitas do paciente ao longo do tempo.
• Ele vê que o paciente começou com um "histórico de cirurgia" na visita 1, evoluiu para "artérias entupidas" na visita 2 e agora precisa de "medicação forte" na visita 3.
• Ele entende que a doença é um caminho (uma trajetória), não apenas pontos soltos no mapa.

C. O Resumo Final (Representação do Paciente)

No final, o modelo junta todas essas peças: o resumo dos "bairros" do corpo e o "filme" da evolução da doença. Ele cria uma identidade digital única para aquele paciente, que captura não apenas o que ele tem hoje, mas para onde a doença está indo.

3. O Treinamento: Aprendendo a Ler o "Subtexto"

Para que esse detetive fosse tão bom, os pesquisadores criaram um método de treinamento especial (Pré-treinamento):

• Jogo do "O que falta?": Eles escondiam partes do prontuário e pediam para o modelo adivinhar o que estava faltando. Isso força o modelo a entender o contexto e as conexões entre as doenças.
• O "Avô" da Doença: Eles também pediam para o modelo adivinhar a "categoria pai" de uma doença (ex: se o código é "diabetes tipo 2", o modelo deve saber que isso pertence à categoria "Doenças Endócrinas"). Isso ajuda o modelo a entender a hierarquia médica, como um aluno que aprende não só a palavra, mas a família dela.

4. Por que isso é importante? (A Interpretabilidade)

A maior vantagem do DT-BEHRT não é apenas ser preciso, mas ser transparente.

• Muitos modelos de IA são "caixas pretas": você dá os dados e eles dão um número, mas não dizem o porquê.
• O DT-BEHRT funciona como um mapa de atenção. Quando ele faz uma previsão, ele mostra: "Olha, eu estou focando no sistema respiratório e na evolução da doença cardíaca nas últimas visitas".
• Isso é como se o médico pudesse olhar para o computador e ver: "Ah, a IA está preocupada com o coração dele porque ela viu essa sequência específica de eventos". Isso gera confiança, pois a IA pensa de forma semelhante a um médico humano.

Resumo em uma frase

O DT-BEHRT é um novo sistema de inteligência artificial que entende o prontuário médico não como uma lista de compras bagunçada, mas como uma história complexa de vida, separando os sistemas do corpo e seguindo a evolução da doença ao longo do tempo, tudo isso de uma forma que os médicos conseguem entender e confiar.

Os testes mostraram que ele é muito melhor do que os sistemas antigos para prever coisas como readmissão ao hospital ou mortalidade, especialmente em pacientes com histórias médicas longas e complexas.

Resumo técnico

1. Problema e Motivação

O uso generalizado de Registros Eletrônicos de Saúde (EHR) oferece oportunidades sem precedentes para modelagem preditiva na medicina. No entanto, os dados estruturados de EHR apresentam desafios específicos que os modelos existentes não abordam adequadamente:

• Heterogeneidade dos Códigos Médicos: Os códigos médicos (diagnósticos, medicamentos, procedimentos, exames) desempenham papéis clínicos distintos. Por exemplo, diagnósticos são a força motriz da trajetória de saúde do paciente e possuem interações densas dentro de sistemas orgânicos, enquanto medicamentos e procedimentos refletem caminhos de tratamento temporalmente relacionados, mas com interações limitadas dentro de uma única visita.
• Limitações dos Modelos Atuais:
  • Abordagens Baseadas em Sequência (ex: BEHRT, Med-BERT): Tratam todos os códigos uniformemente, ignorando a ordem não confiável dos códigos dentro de uma mesma visita e falhando em capturar dependências estruturais complexas.
  • Abordagens Baseadas em Grafos: Capturam bem as relações estruturais, mas muitas vezes falham em modelar dependências temporais sequenciais entre visitas hospitalares.
  • Abordagens Híbridas: Embora tentem integrar grafos e sequências, a maioria ainda trata todos os tipos de códigos da mesma maneira, negligenciando a necessidade de módulos específicos para capturar a dinâmica de progressão de doenças versus tratamentos.

O objetivo do trabalho é desenvolver uma arquitetura que disante as trajetórias de doenças, modelando explicitamente as interações centradas em diagnósticos e padrões de progressão assíncronos, resultando em representações de pacientes interpretáveis e robustas.

---

2. Metodologia: DT-BEHRT

O DT-BEHRT (Disease Trajectory-aware Transformer) é uma arquitetura de sequência aprimorada por grafos, composta por quatro módulos principais que operam em conjunto para capturar diferentes aspectos da trajetória de saúde do paciente:

A. Representação de Sequência (Sequence Representation - SR)

• Utiliza uma arquitetura Transformer baseada no BERT.
• Processa a sequência de visitas hospitalares, onde cada token é um código médico.
• Incorpora embeddings de tipo de código (diagnóstico, medicação, laboratório, procedimento) e índice de visita para capturar o contexto temporal e semântico.
• Não assume ordem estrita dentro de uma única visita, tratando os códigos como um conjunto (bag-of-words) dentro do contexto da visita.

B. Agregação de Doenças (Disease Aggregation - DA)

• Conceito: Baseia-se na ontologia ICD-9, que organiza diagnósticos em 19 capítulos (sistemas orgânicos, ex: cardiovascular, respiratório).
• Mecanismo: Introduz tokens especiais de agregação ($A$), um para cada capítulo da ICD-9.
• Função: Se um capítulo tiver um número suficiente de códigos distintos na trajetória do paciente, o token correspondente é adicionado à sequência.
• Máscara de Atenção: Restringe a atenção para que cada token de agregação (DA) interaja apenas com os códigos de diagnóstico dentro do seu próprio capítulo e consigo mesmo.
• Regularização: Adiciona uma perda de covariância para garantir que os tokens DA capturem abstrações decorrelacionadas de diferentes sistemas orgânicos, evitando redundância.

C. Progressão de Doenças (Disease Progression - DP)

• Conceito: Modela a evolução temporal da doença através de um grafo heterogêneo.
• Estrutura do Grafo:
  • Nós de Visita (DP Nodes): Nós virtuais que representam cada visita hospitalar.
  • Nós de Diagnóstico: Nós conectados às visitas correspondentes.
• Conexões:
  • Conexões direcionadas de nós de visita para seus diagnósticos.
  • Conexões sequenciais entre nós de visita ($t \to t+1$) para capturar a evolução temporal.
  • Laços de autoconexão para preservar informações.
• Aprendizado: Utiliza uma Rede de Atenção em Grafos (GAT) para atualizar as representações dos nós de visita, agregando informações de diagnósticos e de visitas anteriores, permitindo capturar tendências de desenvolvimento de doenças ao longo do tempo.

D. Representação do Paciente (Patient Representation - PR)

• Integra as informações dos três módulos anteriores:
  1. O token [SEQ] (resumo global da sequência).
  2. Os tokens DA (progressão por sistema orgânico).
  3. Os tokens DP (tendências temporais de evolução).
• Utiliza um mecanismo de pooling baseado em atenção para ponderar a importância relativa dos tokens DA e DP em relação ao token [SEQ], gerando o vetor final de representação do paciente ($h_{[CLS]}$).

E. Framework de Pré-treinamento

Para melhorar a robustez e o alinhamento entre os módulos, o modelo utiliza duas tarefas de pré-treinamento:

1. Previsão de Máscara Global de Código (GCMP): Semelhante ao MLM (Masked Language Modeling), mas com foco na semântica de co-ocorrência ao nível da trajetória (não apenas dentro de uma visita). Todos os códigos de um tipo selecionado são mascarados em toda a sequência.
2. Previsão de Código Ancestral (ACP): Tarefa auxiliar onde o modelo deve prever o capítulo da ICD-9 (ancestral) de um código de diagnóstico mascarado. Isso força os módulos SR e DP (que não têm acesso direto à ontologia) a alinhar suas representações com a estrutura hierárquica de conhecimento do módulo DA.

---

3. Contribuições Chave

1. Arquitetura Modelo Específica: O DT-BEHRT é o primeiro a introduzir módulos especializados que tratam diferentemente códigos de diagnóstico (focados em trajetória e sistemas orgânicos) e códigos de tratamento, desenhando uma arquitetura que reflete a heterogeneidade clínica real.
2. Framework de Pré-treinamento Inovador: A combinação de máscara de código em nível de trajetória com previsão de ancestrais ontológicos promove um alinhamento semântico robusto entre os módulos de sequência, agregação e progressão.
3. Interpretabilidade Clínica: O design do modelo gera representações que são naturalmente interpretáveis por médicos, pois os tokens de atenção correspondem a sistemas orgânicos (capítulos ICD-9) e trajetórias temporais de doenças, alinhando-se ao raciocínio clínico centrado na doença.

---

4. Resultados Experimentais

O modelo foi avaliado em três conjuntos de dados públicos de EHR: MIMIC-III, MIMIC-IV e eICU.

• Tarefas de Predição Geral:
  • Desempenho: O DT-BEHRT superou consistentemente os baselines de ponta (incluindo BEHRT, Med-BERT, G-BERT, HEART e HypEHR) em tarefas de mortalidade hospitalar, permanência prolongada (PLOS) e readmissão.
  • Destaque: A maior melhoria foi observada na tarefa de readmissão, que é notoriamente difícil devido à sua natureza multifatorial. O modelo demonstrou maior robustez em subgrupos de pacientes com comorbidades específicas (diabetes, insuficiência cardíaca, etc.).
• Predição de Fenotipagem:
  • O modelo obteve o melhor desempenho global (Macro-AUPRC) na predição de múltiplas condições crônicas e agudas, especialmente em pacientes com múltiplas visitas, indicando sua capacidade de capturar dependências temporais de longo prazo.
• Estudo de Ablação:
  • A remoção do módulo de Progressão de Doenças (DP) ou de Agregação (DA) resultou em queda de desempenho, confirmando que ambos são essenciais.
  • A tarefa de pré-treinamento ACP (Previsão de Ancestral) mostrou-se crucial, especialmente para a tarefa de mortalidade, melhorando o alinhamento entre os módulos.
• Estudo de Caso:
  • Análises qualitativas mostraram que o modelo atribui atenção correta a códigos relacionados a sistemas orgânicos específicos (ex: sistema respiratório) e consegue traçar trajetórias clínicas coerentes (ex: histórico de cirurgia de bypass $\to$ aterosclerose coronariana), validando sua interpretabilidade.

---

5. Significado e Conclusão

O DT-BEHRT representa um avanço significativo na representação de pacientes baseada em EHR ao reconhecer que nem todos os códigos médicos são iguais. Ao desacoplar a modelagem de diagnósticos (focada em sistemas orgânicos e progressão) da de tratamentos, e ao integrar grafos heterogêneos para capturar a evolução temporal, o modelo supera as limitações das abordagens puramente sequenciais ou de grafos estáticos.

Impacto Principal:

• Precisão: Melhora o estado da arte em tarefas críticas de predição clínica.
• Interpretabilidade: Fornece insights que se alinham com o raciocínio dos médicos (sistemas orgânicos e evolução da doença), o que é fundamental para a adoção de IA em ambientes clínicos reais, onde a "caixa preta" é uma barreira.
• Robustez: Mantém bom desempenho em diferentes subgrupos de pacientes e bancos de dados, sugerindo boa generalização.

O código fonte do projeto está disponível publicamente, facilitando a reprodutibilidade e o avanço futuro na área de informática em saúde.