Structure-Aware Set Transformers: Temporal and Variable-Type Attention Biases for Asynchronous Clinical Time Series

O artigo apresenta o STAR Set Transformer, um modelo que melhora a previsão em séries temporais clínicas assimétricas ao incorporar vieses de atenção suaves para capturar a localidade temporal e a compatibilidade entre variáveis, superando as abordagens baseadas em grades e conjuntos de pontos em tarefas de cuidados intensivos.

O essencial

Imagine que os registros médicos de um paciente (como pressão arterial, batimentos cardíacos e níveis de glicose) são como uma caixa de legos espalhada no chão.

A maioria dos computadores e modelos de inteligência artificial atuais tentam organizar esses legos em uma prateleira rígida e quadrada (uma grade). Eles forçam cada peça a entrar em um quadrado de tempo específico (ex: "toda hora em ponto"). O problema? Os médicos não anotam coisas a cada hora exata. Às vezes, eles medem a pressão a cada 15 minutos, às vezes a cada 4 horas, e às vezes esquecem de medir algo. Para preencher os buracos na prateleira, o computador precisa "adivinhar" (imputar) os dados faltantes, o que pode criar erros ou fazer o modelo aprender truques sujos baseados apenas no que não foi medido.

Outra abordagem é tratar cada medição como um ponto solto em uma caixa de areia. Isso é mais fiel à realidade, mas o computador perde a noção de duas coisas importantes:

1. A história de cada variável: Como a pressão do paciente mudou ao longo do tempo (a linha vermelha na figura).
2. A relação entre variáveis no mesmo momento: Como a pressão e a frequência cardíaca se comportaram juntas naquele instante específico (a linha azul na figura).

A Solução: O "STAR" (Transformador Consciente da Estrutura)

Os autores deste paper criaram uma nova maneira de ensinar o computador a ler esses legos espalhados. Eles chamam seu modelo de STAR (Set Transformer Consciente da Estrutura).

Em vez de forçar os legos na prateleira ou deixá-los totalmente bagunçados, o STAR usa dois "óculos mágicos" (chamados de vieses de atenção) para ajudar o computador a ver o que é importante:

1. O Óculo do Tempo (Viés Temporal):
   Imagine que você está em uma festa e vê alguém que você conheceu há 5 minutos. É mais provável que você converse com essa pessoa do que com alguém que chegou há 2 horas. O modelo aprende a dar mais importância a medições que aconteceram perto no tempo. Ele "penaliza" o computador se ele tentar conectar uma medição de hoje com uma de 3 dias atrás, a menos que seja realmente necessário.

2. O Óculo do Tipo (Viés de Tipo de Variável):
   Imagine que você tem uma caixa de ferramentas. Você sabe que um martelo conversa melhor com um prego do que com um copo d'água. Da mesma forma, o modelo aprende que uma medição de "glicose" deve conversar mais com outra medição de "glicose" (para ver a evolução) do que com uma medição de "temperatura". Ele cria uma "afinidade" entre coisas do mesmo tipo.

Como eles testaram isso?

Eles misturaram esses "óculos" em diferentes camadas do cérebro do computador (como se fossem camadas de uma cebola ou andares de um prédio). Eles testaram 10 combinações diferentes:

• Usar só o óculo do tempo?
• Usar só o óculo do tipo?
• Usar os dois?
• Usar o óculo do tempo nos andares de baixo e o do tipo nos andares de cima?

O Resultado

O modelo vencedor, chamado STAR-Set, funcionou como um detetive muito esperto. Ele conseguiu prever três coisas críticas em pacientes de UTI melhor do que qualquer método anterior:

1. Risco de Parada Cardíaca (CPR): Identificou quem precisava de ajuda imediata com muito mais precisão.
2. Risco de Mortalidade: Previu quem tinha maior chance de não sobreviver.
3. Uso de Vasopressores: Antecipou quando o paciente precisaria de remédios para manter a pressão.

Por que isso importa?

A grande vantagem é que o STAR não precisa inventar dados para preencher buracos na grade. Ele entende a bagunça natural dos dados médicos. Além disso, os "óculos" que ele aprendeu são interpretáveis: os médicos podem olhar para o modelo e ver o que ele achou importante (ex: "Ah, o modelo percebeu que a pressão e a frequência cardíaca mudaram juntas 2 horas antes da crise").

Em resumo: O STAR ensina a inteligência artificial a olhar para os dados médicos desorganizados e ver padrões de tempo e de relacionamento entre variáveis, sem precisar forçá-los a se encaixar em um calendário rígido. É como dar a um médico um mapa que mostra não apenas onde os pacientes estão, mas como eles se conectam no tempo e no espaço.

Resumo técnico

Resumo Técnico: STAR-Set Transformer

1. O Problema: Representação de Séries Temporais Clínicas Irregulares

Os Registros Eletrônicos de Saúde (EHR) são séries temporais multivariadas, irregulares e assíncronas. Diferente de imagens ou texto, não existe uma discretização canônica do tempo para dados clínicos. A escolha da "layout" de entrada para modelos de rede neural é um desafio crítico:

• Grids Regulares: Discretizam o tempo em intervalos fixos (ex: horários). Exigem imputação de dados faltantes ou máscaras de ausência, o que pode introduzir erros ou levar o modelo a aprender "atalhos" baseados nas políticas de amostragem (ex: saber que um exame não foi feito não significa que o paciente está saudável).
• Grids Baseados em Eventos (Event-Time): Indexam o tempo pelos timestamps observados. Embora evitem a imputação densa, ainda sofrem com a assincronia e a necessidade de máscaras complexas.
• Representações em Conjuntos (Set Tokenization): Tratam cada evento observado como um token (valor, tempo, tipo de variável) em um conjunto, sem discretização. Embora evitem a imputação, essa abordagem perde duas estruturas indutivas cruciais presentes nos grids:
  1. Trajetórias dentro da mesma variável (evolução temporal de um único parâmetro).
  2. Relações contextuais entre variáveis no mesmo momento temporal (interações cruzadas).

O modelo proposto busca recuperar essas estruturas indutivas sem forçar uma discretização rígida do tempo.

2. Metodologia: STAR-Set Transformer

Os autores propõem o STAR-Set Transformer, uma arquitetura que combina a flexibilidade da tokenização em conjuntos com vieses de atenção (attention biases) suaves e eficientes em parâmetros. O objetivo é reintroduzir a estrutura "grid-like" (tempo × variável) dentro de um encoder baseado em conjuntos.

Componentes Principais:

1. Representação de Entrada: Um episódio de EHR é tratado como um conjunto de eventos irregulares $X \in \mathbb{R}^{B \times L \times 3}$, onde cada evento é um triplet $(tempo, valor, índice\_da\_variável)$. O modelo utiliza um Set Embedder (baseado em ITE) para gerar tokens iniciais, incluindo tokens especiais [CLS] e demográficos.
2. Encoder Transformer com Vieses Aditivos: O núcleo da inovação está na modificação dos logits de atenção (antes do softmax) para adicionar dois tipos de vieses aprendíveis:
   • Viés Temporal (Temporal Bias): Penaliza a atenção entre tokens distantes no tempo.
     • Fórmula: $b_{time} = -|\Delta t| / \tau$, onde $\tau$ é uma escala de tempo aprendível por camada e cabeça. Isso incentiva interações entre eventos temporalmente próximos, restaurando a prior de "localidade temporal".
   • Viés de Tipo de Variável (Variable-Type Bias): Promove interações entre tokens da mesma variável (ex: pressão arterial com pressão arterial) ou entre tipos compatíveis.
     • Fórmula: Utiliza uma matriz de afinidade de tipos $B \in \mathbb{R}^{F \times F}$ aprendida, onde $F$ é o número de tipos de variáveis.
3. Estratégias de Fusão por Camada (Layer-wise Fusion): Os autores investigaram sistematicamente onde injetar esses vieses dentro de um Transformer de 4 camadas. Eles testaram 10 agendamentos diferentes (ex: aplicar viés temporal apenas nas camadas iniciais, ou combinar ambos em todas as camadas). A configuração final proposta (vt-vt) aplica ambos os vieses em todas as camadas do encoder.

3. Contribuições Chave

1. STAR-Set Transformer: Introdução de um modelo que augmenta encoders de conjuntos para EHR com vieses de atenção simples, recuperando estrutura indutiva sem discretização.
2. Vieses Complementares e Eficientes: Combinação de um viés temporal (baseado em distância) e um viés de compatibilidade de variáveis, ambos com poucos parâmetros adicionais (escalas $\tau$ e matriz $B$).
3. Análise Sistemática de Ablação: Estudo detalhado de como a profundidade do Transformer afeta a eficácia dos vieses, identificando que a aplicação consistente de ambos os vieses em todas as camadas (vt-vt) oferece o melhor desempenho.
4. Interpretabilidade: Os parâmetros aprendidos ($\tau$ e $B$) fornecem resumos interpretáveis sobre os contextos temporais relevantes e as interações entre variáveis clínicas.

4. Resultados Experimentais

O modelo foi avaliado em três tarefas de predição em unidades de terapia intensiva (ICU) utilizando o conjunto de dados MIMIC-IV:

• Tarefas: Parada Cardiorrespiratória (CPR), Mortalidade e Uso de Vasopressores.
• Baselines Comparados: Modelos baseados em Grid Regular (SMART, DueTT), Grid de Tempo de Evento (PrimeNet) e Representação em Conjunto (STraTS).

Desempenho (AUC / APR):
O STAR-Set superou consistentemente todas as baselines em todas as tarefas:

• CPR: AUC 0.7158 (vs. 0.6478 do melhor baseline, DueTT).
• Mortalidade: AUC 0.9164 / APR 0.2033 (vs. 0.8778/0.1457 do STraTS).
• Vasopressores: AUC 0.8373 / APR 0.1258 (vs. 0.8255/0.1109).

Análise de Ablação:

• O viés temporal foi o principal motor de ganhos no AUC (especialmente para CPR).
• O viés de tipo de variável trouxe melhorias consistentes e foi crucial para o APR (Precisão Média).
• A combinação de ambos (vt-vt) resultou no melhor equilíbrio, superando modelos que usam apenas um tipo de viés ou nenhum viés.

5. Significado e Conclusão

O trabalho demonstra que é possível obter o melhor dos dois mundos: a flexibilidade de processar dados irregulares sem imputação (como em modelos de conjuntos) e a capacidade de capturar dependências temporais e relacionais estruturadas (como em grids).

Ao reintroduzir vieses indutivos suaves em modelos baseados em atenção, o STAR-Set oferece uma solução prática para modelagem de séries temporais clínicas, evitando os artefatos de imputação e as limitações de discretização rígida. Os resultados sugerem que a estrutura de "tempo × variável" é fundamental para a performance em tarefas críticas de saúde, mesmo quando os dados são inerentemente desordenados. O modelo serve como um "plug-in" eficaz para modelos de séries temporais baseados em contexto, melhorando a precisão e a interpretabilidade das previsões clínicas.