Tokenization Tradeoffs in Structured EHR Foundation Models

Resumo Técnico: Trade-offs de Tokenização em Modelos de Base para Registros Eletrônicos de Saúde (EHR)

1. O Problema

Os modelos de base (foundation models) para Registros Eletrônicos de Saúde (EHR) estruturados são pré-treinados em sequências longitudinais de eventos clínicos temporizados para aprender representações adaptáveis de pacientes. Um componente crítico, porém pouco estudado, é a tokenização: o processo de converter a linha do tempo clínica de um paciente em uma sequência discreta de tokens.

As decisões de design na tokenização determinam:

• Quais informações são preservadas (ex: associações código-valor, intervalos temporais).
• A eficiência do código e o tamanho do vocabulário.
• Quais relações o modelo deve aprender dos dados versus quais são pré-calculadas.

Existe um trade-off fundamental: codificar a estrutura clínica no momento da tokenização (pré-computação) pode expandir o vocabulário e os parâmetros, mas reduz a carga de aprendizado do modelo. Por outro lado, uma tokenização fatorada (separada) mantém um vocabulário compacto, mas exige que o modelo aprenda associações complexas (como o significado de um valor numérico dependendo do código clínico anterior) a partir de dados limitados, um problema conhecido como problema de ligação local (local binding problem). Até este estudo, o impacto dessas escolhas no desempenho downstream e na eficiência computacional não havia sido quantificado de forma controlada.

2. Metodologia

Os autores conduziram um experimento fatorial 2 × 2 × 2 para isolar os efeitos independentes de três eixos de design de tokenização:

1. Codificação de Eventos (Event Encoding):
   • Joint (Conjunta): Cada evento clínico é mapeado para um único token que funde o código e seus atributos (ex: "Glicose_Serica_Quantil3").
   • Factorized (Fatorada): O evento é decomposto em tokens separados (código base + tokens de atributos/valores).
2. Codificação de Tempo (Time Encoding):
   • Time-Positions: Usa a idade do paciente em dias como índice posicional, codificado via Rotary Positional Embeddings (RoPE).
   • Time-Tokens: Insere tokens discretos de intervalo de tempo entre eventos clínicos consecutivos.
3. Anotações de Fluxo de Trabalho (Workflow Annotations):
   • Com Workflow: Eventos clínicos com múltiplos passos (ex: pedido, coleta, resultado) são representados como eventos distintos em timestamps diferentes.
   • Sem Workflow: Apenas o evento final (ex: resultado) é representado.

Configuração Experimental:

• Dados: Pré-treinamento realizado no conjunto de dados do Hospital SickKids (criança) e avaliação externa no conjunto MIMIC-IV (adulto, UTI).
• Modelo: Arquitetura Transformer decodificador-only (28 camadas, 768 dimensões, ~116M parâmetros fixos no backbone). O número total de parâmetros variou (132M–170M) devido às diferenças no tamanho do vocabulário nas camadas de entrada e saída.
• Tarefas: Avaliação em 74 tarefas de predição clínica (transfusões, procedimentos, imagens, laboratório, medicação e desfechos) no SickKids e 13 tarefas no MIMIC.
• Métricas: Área sob a curva ROC (AUROC) e eficiência computacional (FLOPs de pré-treinamento).
• Ablações: Experimentos específicos para dissecar se as diferenças de desempenho eram devidas à eficiência de tokens (tamanho da sequência) ou à eficiência de ligação local.

3. Principais Contribuições e Resultados

A. Desempenho e Eficiência Computacional

• Codificação Conjunta (Joint) vs. Fatorada: A codificação conjunta superou a fatorada em 73/74 tarefas (99%).
  • Eficiência: A codificação conjunta exigiu 39,5% menos operações de ponto flutuante (FLOPs) no pré-treinamento, pois gerou sequências mais curtas.
  • Mecanismo: As ablações mostraram que a vantagem não se deveu apenas ao tamanho da sequência, mas à eficiência de ligação local. Na codificação fatorada, o modelo precisa aprender que o mesmo token de valor (ex: "Quantil 2") tem significados clínicos diferentes dependendo do código que o precede. A codificação conjunta pré-computa essa associação, resolvendo o problema de ligação local que é difícil de aprender com dados limitados de EHR.
• Codificação de Tempo: A estratégia Time-Positions (RoPE baseado na idade) superou Time-Tokens em 71/74 tarefas (97%).
  • Eficiência: Time-Positions exigiu 9,6% menos FLOPs.
  • Observação: Tokens de tempo explícitos degradaram o desempenho, provavelmente ao fragmentar sequências e consumir posições de contexto sem ganho informativo proporcional.

B. Generalização e Transferência

• Avaliação Externa (SickKids para MIMIC):
  • Houve uma grande incompatibilidade de vocabulário (69,8% de tokens fora do vocabulário no MIMIC).
  • A vantagem da codificação conjunta generalizou-se para o conjunto de dados externo (efeito similar ao local), indicando que a eficiência de ligação local é uma propriedade fundamental da modelagem de eventos estruturados.
  • Os efeitos de Time-Positions e Workflow não generalizaram significativamente. As anotações de fluxo de trabalho são específicas da instituição, e as codificações temporais calibradas para a população pediátrica não se adaptaram bem à população adulta de UTI.

C. Análise de Eficiência de Amostra (Sample Efficiency)

• As vantagens da codificação conjunta e posicional foram consistentes em todos os tamanhos de conjunto de dados (de 2 a 32.768 exemplos), embora a vantagem temporal tenha diminuído ligeiramente em conjuntos de treinamento maiores.

4. Significado e Conclusões

Este estudo estabelece a tokenização como uma alavanca tratável para melhorar simultaneamente o desempenho e a eficiência de modelos de base para EHR.

• Recomendação de Design: Para dados de EHR estruturados, onde os dados de treinamento são ordens de magnitude menores que os corpora de linguagem, codificar a estrutura clínica no momento da tokenização (usando tokens conjuntos e posições temporais) é superior. Isso evita que o modelo gaste capacidade de aprendizado limitante tentando descobrir associações composicionais básicas.
• Eficiência: Estratégias que capturam estrutura de domínio significativa resultam em sequências mais compactas e modelos mais rápidos de treinar, sem sacrificar a precisão.
• Generalização: Embora a estrutura local (código-valor) seja transferível entre instituições, características institucionais específicas (fluxos de trabalho) e demográficas (distribuição de idade) limitam a transferência direta de certas anotações de tokenização.
• Implicação Prática: Ao desenvolver modelos de base para saúde, a escolha da tokenização deve priorizar a resolução de problemas de ligação local e a compactação de informações estruturais, em vez de depender puramente da capacidade do modelo para aprender essas relações a partir de dados brutos.

Em suma, o artigo demonstra que "menos é mais" em termos de tokens brutos, mas "mais estrutura" no design do token (junção de código e valor) leva a modelos mais robustos, eficientes e generalizáveis.