Temporal Island Sparse Autoencoders for Interpreting Clinical Time-Series Models

AI summary

Overview

Research area: Interpretable machine learning for clinical time-series — specifically, mechanistic interpretability of deep EHR (electronic health record) models via sparse dictionary learning.

Technical level: Intermediate. The paper assumes familiarity with attention-based sequence models, post-hoc attribution baselines (Integrated Gradients, SHAP, Dynamask, TimeX), and the sparse autoencoder (SAE) line of work from language-model interpretability. The mathematics is presented compactly but is not the main barrier; the barrier is background context.

Scope: The paper introduces TI-SAE, a two-stage sparse autoencoder that decomposes a frozen clinical time-series model's internal representations into a reusable dictionary of temporally localized, variable-sparse "clinical concepts," and evaluates it across five backbone architectures on MIMIC-III and MIMIC-IV.

What This Paper Is About

Deep models trained on ICU event sequences predict mortality, readmission, and long stays well, but they do not tell clinicians what clinical patterns they are relying on. Existing explanation methods output per-cell importance scores over a time × variable grid — they answer "which individual inputs mattered" rather than "which clinical concepts the model uses," and those scores lack temporal extent, cross-variable structure, and any notion of reuse across patients.

TI-SAE addresses this by learning a finite, patient-shared dictionary of clinical concepts directly from a frozen backbone's hidden representations. Each dictionary entry encodes when a concept is active (via learned temporal Gaussian "islands") and which variables it involves (via a sparse composition vector), and a separately trained selector then identifies which of those concepts drive an individual patient's prediction.

Key Contributions

1. A temporally structured SAE architecture. TI-SAE replaces flat per-token SAE latents with latents defined by M learned Gaussian islands (center μ and width σ, estimated partly via EM with σ clamped to [0.5, 12.0]) plus a sparse, row-normalized, non-negative variable-routing vector shared between encoder and decoder (κ_enc = 5, κ_dec = 3). This encodes spatiotemporal composition directly in the latent space rather than hoping it emerges.

2. Decoupled dictionary learning and attribution. Stage A learns the concept dictionary purely by reconstructing frozen backbone representations (per-cell normalized MSE on observed grid cells, with dead-latent revival and four regularizers); Stage B trains a small (~38K-parameter) prediction-aware selector on the frozen dictionary. The authors isolate this design: applying the same selection objective directly to raw hidden states collapses to neighborhood-based AOPC of .038.

3. A recognizable gap in the explanation-method landscape. Table 1 positions existing families (attention, gradient-based, perturbation-based, learned selection, concept bottlenecks) against four structural properties — temporal structure, decomposability, post-hoc operation, and no concept-label supervision — and identifies TI-SAE as the only method satisfying all four.

4. Multi-backbone, multi-dataset, multi-rater evaluation. Five backbone configurations (discrete transformer with and without 2-hour value-bin forecast pretraining, STraTS, HyMaTE, BAT) on MIMIC-III (N = 28,626 episodes, 14.1% mortality) and MIMIC-IV (N = 46,140 episodes, 14.4% mortality), three tasks each, with faithfulness, SOFA-based clinical plausibility, LLM-judge verbalization, and a four-clinician panel.

Main Findings

• Highest faithfulness on four of five backbones. Under neighborhood-based perturbation, TI-SAE (soft) reaches AOPC of .086 (Discrete), .076 (Disc.+PT), .097 (HyMaTE), and .043 (BAT), versus e.g. .053/.019/.051/– for attention, .060/.051/.049/.031 for Integrated Gradients, and .063/.054/.062/.037 for TSR. The hard-selection variant scores .079/.068/.090/.041.

• STraTS is the exception, and the authors explain why. On STraTS, TI-SAE's AOPC is .016 while IG reaches .067 and TSR .071. The paper attributes this to STraTS's residual-averaging architecture, which "does not preserve the per-variable locality that TI-SAE's routing requires."

• Competitive SOFA-based plausibility. TI-SAE (soft) achieves AUROC/AUPRC of .593/.286 (Discrete), .589/.255 (Disc.+PT), .614/.276 (HyMaTE), .588/.251 (BAT), .536/.227 (STraTS). The paper states TI-SAE has the highest AUROC and AUPRC on the discrete backbone; on other backbones it is described as competitive — TSR reaches .625/.272 on HyMaTE and WinIT reaches .621/.280 on STraTS, which are higher than TI-SAE's corresponding values.

• Clinician-rated dictionary labels. Four independent clinicians rated a sample of 30 dictionary labels on a 5-point Likert scale, yielding mean ratings of 4.07–4.40. High-mortality latents corresponded to vasopressor-dependent shock and respiratory failure; protective latents to hemodynamic stability and normal acid-base status.

• Reproducible across independent runs. Full-pipeline retraining gives decoder cosine .724 and variable overlap 77%; training the dictionary alone gives decoder cosine .735 and variable overlap 82%; training the selector on a fixed dictionary gives top-64 Jaccard .708 and rank correlation .87.

• Zero-ablating concepts degrades predictions monotonically, and more than random-order removal, indicating individual concepts carry non-redundant predictive contributions.

• Cross-dataset transfer without retraining. A MIMIC-III-trained model applied to MIMIC-IV (restricted to 2014+ admissions to remove patient and temporal overlap) holds AOPC at .130 out-of-distribution versus .149 in-distribution, with AUROC/AUPRC of .756/.370 versus .756/.386. Backbone AUROC moves only .861 → .864. By contrast, IG drops from .086 to .045 and TSR from .088 to .046.

• Better verbalized explanations. Against Integrated Gradients, the LLM judge (N = 50 per baseline) gives TI-SAE 33–15 wins on concept-level reasoning and 32–18 on actionability; against attention, 30–17 and 28–22; against TSR, 33–16 and 27–23. The four-clinician panel on the IG comparison (N = 50) votes 39–4 on concept-level reasoning and 23–9 on actionability, with faithfulness near-tied at 4–4 (42 ties).

• The dictionary labels carry much of the verbalization gain. An ablation shows that removing the cross-patient concept labels eliminates TI-SAE's clear advantage over per-cell attribution — but only TI-SAE's shared dictionary supplies those labels with spatiotemporal grounding.

• Structural ablations separate faithfulness from readability. Sparse variable routing is the critical mechanism for faithfulness (removing it costs 58–75% of AOPC), while temporal islands and channel organization barely affect faithfulness but drive readability: channel coherence falls from 91.9% to 36.9% without them, and temporal localization collapses. Reducing to one island per concept (M = 1) degrades both faithfulness and plausibility; replacing the symmetric Gaussian with a learned split-normal window changes neither metric.

Methodology in Plain English

The setup. Each patient is a sparse, irregular list of events — a timestamp, a variable identity, and a measured value — over a 48-hour ICU window, with V = 135 variables for MIMIC-III and V = 132 for MIMIC-IV. A pretrained backbone (already trained on the prediction task and then frozen) turns these events into contextual embeddings. The authors "scatter" those embeddings into a dense time × variable grid with an observation mask marking which cells contain real measurements. Everything downstream operates on this grid.

Stage A — learning the dictionary. TI-SAE has K = 1280 latents, each with M = 4 candidate temporal islands. The encoder produces an activation trace per latent by attending to a sparse subset of variables (chosen by a routing vector) and projecting onto a learned direction. A small convolutional network (IslandParamNet, hidden dimension 32, kernel size 25) then reads each trace and locates where in time the concept is active: it predicts island centers by softmax-weighted averaging over time positions, and estimates widths via three EM iterations. Each island is a Gaussian window acting as a temporal support gate — it says when, while the encoder trace carries the actual activation within that window. From the K × M = 5120 candidate islands per sample, a budget of n_global = 128 is selected through differentiable top-k with straight-through estimation, and the selected islands' Gaussians compose into a soft temporal mask. The decoder reconstructs the grid from the masked activations using the same routing weights plus a per-latent atom.

Keeping the dictionary honest. The training loss is a normalized reconstruction error over observed cells only, supplemented by dead-latent revival (an auxiliary loss, k_aux = 256, that retrains unused latents on the residual) and four regularizers: center repulsion to stop islands collapsing to the same time, decoder diversity to prevent latents duplicating variable compositions, mask coverage to prevent gaming reconstruction by widening windows, and channel purity to concentrate each latent on one clinical system. No prediction loss is used here.

Stage B — finding what mattered for this patient. With the whole SAE frozen, a small selector scores each island using features (temporal parameters, activation statistics, routing weights, observation coverage) combined with a patient-level query drawn from the backbone's attention pooling. It keeps n_global^pred = 64 islands per patient via softmax-weighted selection. It is trained with two losses: match the teacher logit when using only the selected islands, and a hinge term ensuring that dropping the selected islands hurts the prediction in the teacher's predicted class direction.

Making concepts human-readable. After training, each latent is summarized from its top exemplar episodes — its variable composition and temporal extents — and an LLM labeler produces a clinical concept label and description. These labels are validation artifacts, not supervision.

How it was judged. Faithfulness was measured via comprehensiveness and sufficiency under three perturbation strategies (removal, marginal replacement with the training-set median, and neighborhood-based replacement with the patient's temporally nearest value of the same variable), summarized as AOPC, with neighborhood-based as the primary. Clinical plausibility used SOFA organ-system severity as soft ground truth (AUROC/AUPRC on whether top-attributed cells correspond to severe observations). Explanation quality was assessed by pairwise LLM judge and by four tertiary-hospital clinicians on concept-level reasoning, clinical actionability, and faithfulness, comparing attribution-augmented versus TI-SAE-augmented verbalizations at matched information budget.

Why This Matters

For research. The paper makes a structural argument rather than only an empirical one: point attributions are the wrong output format for clinical reasoning, and the fix is architectural — build temporal extent and variable composition into the latent space. It also imports the SAE/dictionary-learning toolkit from language-model interpretability into clinical time series, where it had not previously been applied with explicit spatiotemporal structure, and it is candid that the winning ingredient is the shared dictionary (the source of cross-patient concept labels), not the reconstruction objective alone.

Real-world applications:

• Clinical model review and sign-off. A hospital deploying a sepsis or mortality model could inspect a finite, labeled concept vocabulary (e.g. a vasopressor-dependent shock latent, a hemodynamic-stability latent) instead of reviewing heatmaps per patient.
• Chart-level explanation for clinicians. Concept-level, temporally grounded explanations that beat cell-level attribution on concept reasoning and actionability in both LLM-judge and clinician votes are more plausible material for a bedside or tumor-board style discussion.
• Registry-level pattern auditing. Because the dictionary is patient-shared, it supports population-level questions — which concept families recur, and which drive predictions across a cohort — that per-patient heatmaps cannot answer.
• Model comparison and debugging. Because TI-SAE is backbone-specific and applied identically across five architectures, differences in what concepts a model has learned become inspectable — the paper's STraTS finding is exactly this kind of diagnostic.

Industry relevance. For EHR vendors, regulators, and hospital AI-governance teams, the appeal is operational: the method is fully post-hoc on a frozen backbone, requires no architectural modification or retraining, and requires no external concept-label annotation. That makes it a plausible bolt-on to models already in deployment, and the reported cross-dataset transfer (MIMIC-III → MIMIC-IV, AOPC .149 → .130) is a first indication of the robustness such deployment would demand.

Future Directions

• Backbones beyond single-token-per-event representations. TI-SAE assumes each event maps to one grid cell; the authors flag that multi-token strategies (e.g. ETHOS-style serialization, time-interval tokens) may need architectural adaptation, and that routing requires hidden layers balancing per-variable locality with contextual concept formation — not all layers provide this, as the STraTS result shows.
• Beyond fixed 48-hour grids. Variable-length sequences and observation windows other than the 48-hour setting are explicitly unvalidated.
• Stronger ground truth than SOFA. The authors describe SOFA-derived severity as a coarse, organ-independent proxy that cannot capture full clinical reasoning; more faithful clinical reference standards are an open need.
• Disentangling the LLM verbalizer's contribution. The evaluation is mediated by a verbalizer with strong clinical priors that can produce fluent explanations from any input. The authors note that the concept-based condition constrains the verbalizer more tightly but that this asymmetry is not directly measured — a controlled study of verbalizer influence is a natural next step.

Target Audience

Researchers working on interpretability for time-series and clinical ML; practitioners building or auditing deployed EHR prediction models who need explanations that map onto clinical vocabulary rather than cell-level scores; and mechanistic-interpretability researchers interested in whether SAE/dictionary-learning methods transfer from language models to structured, irregular, temporally extended data. Clinicians and clinical informaticists involved in model governance will find the framing and the clinician-rated results accessible, though the method sections assume ML background. Readers should note the evaluation is retrospective on two public MIMIC databases, and the explanation-quality assessment is a verbalization study rather than a measure of whether explanations faithfully reflect model computation.

Authors’ abstract

Deep learning models for electronic health records achieve strong predictive performance but provide little insight into what clinical concepts they represent internally. Post-hoc attribution methods produce per-cell importance scores that lack temporal coherence, cross-variable structure, and reusability across patients. We introduce the Temporal Island Sparse Autoencoder (TI-SAE), which learns a backbone-specific, patient-shared clinical concept dictionary from a frozen predictive model's representations. Each dictionary entry (latent) is defined by learned temporal Gaussian islands (specifying when the concept is active) and a sparse variable-composition vector (specifying which variables are involved). A decoupled selector then identifies which dictionary entries drive each patient's prediction, separating representation learning from attribution. We evaluate TI-SAE across five backbone architectures spanning different representation strategies on MIMIC-III and MIMIC-IV (three prediction tasks each). TI-SAE produces structured clinical concepts that go beyond transient attribution heatmaps while achieving strong performance on standard faithfulness benchmarks.