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

TL;DR

This paper introduces Structure-Aware Set Transformers (STAR) that incorporate temporal and variable-type biases into attention mechanisms, improving asynchronous clinical time series modeling without discretizing time.

Contribution

The paper proposes a novel set transformer architecture with parameter-efficient attention biases for better modeling of irregular EHR data, outperforming existing baselines.

Findings

STAR achieves higher AUC/APR on ICU prediction tasks.

Learned biases provide interpretable insights into temporal and variable interactions.

The model effectively handles irregular, asynchronous time series without imputation.

Abstract

Electronic health records (EHR) are irregular, asynchronous multivariate time series. As time-series foundation models increasingly tokenize events rather than discretizing time, the input layout becomes a key design choice. Grids expose time$\times$variable structure but require imputation or missingness masks, risking error or sampling-policy shortcuts. Point-set tokenization avoids discretization but loses within-variable trajectories and time-local cross-variable context (Fig.1). We restore these priors in STructure-AwaRe (STAR) Set Transformer by adding parameter-efficient soft attention biases: a temporal locality penalty $-|\Delta t|/\tau$ with learnable timescales and a variable-type affinity $B_{s_i,s_j}$ from a learned feature-compatibility matrix. We benchmark 10 depth-wise fusion schedules (Fig.2). On three ICU prediction tasks, STAR-Set achieves AUC/APR of 0.7158/0.0026 (CPR), 0.9164/0.2033 (mortality), and 0.8373/0.1258 (vasopressor use), outperforming regular-grid, event-time grid, and prior set baselines. Learned $\tau$ and $B$ provide interpretable summaries of temporal context and variable interactions, offering a practical plug-in for context-informed time-series models.