Sparse Autoencoder Decomposition of Clinical Sequence Model Representations: Feature Complexity, Task Specialisation, and Mortality Prediction
Authors: Chris Sainsbury, Feng Dong, Andreas Karwath
Organizations: School of Cardiovascular & Metabolic Health, University of Glasgow, UK · School of Medicine, University of Dundee, UK · NHS Greater Glasgow and Clyde, UK · Department of Computer and Information Sciences, University of Strathclyde, UK · Cancer and Genomic Sciences, School of Medical Sciences, College of Medicine and Health, University of Birmingham, UK
Sparse autoencoders (SAEs) have been applied to large language models and protein language models, but not systematically to electronic health record (EHR) foundation models. We train TopK SAEs on FlatASCEND, a 14.5-million-parameter autoregressive clinical sequence model, at all 10 residual stream extraction points on INSPECT (outpatient) and MIMIC-IV (ICU). SAE decomposition reveals progressive abstraction across transformer depth: layer-0 features are near-perfect token detectors (45.7% singleton), while layer-6 features span approximately 30 token types across multiple clinical categories (0.5% singleton). Under full-sequence simple linear probes, SAE features outperform dense representations for discrete event prediction (mortality) while dense representations outperform for continuous magnitude prediction (length of stay) - a probe-level representational phenomenon that does not extend to clinically relevant leakage-safe windows, where dense representations match or exceed SAE features across all tested settings (eICU-CRD 48-hour AUC: SAE 0.871 versus dense 0.880; base model zero-shot, SAE dictionaries trained on eICU activations; MIMIC-IV: 0.836 versus 0.914; INSPECT 1-year/3-year: 0.697 versus 0.800). A delta-mode intervention method reduces SAE perturbation noise by 86x, enabling cleaner feature-level experiments, though the resulting perturbation effects are larger than random controls in 3 of 4 conditions but not formally significant. Feature reproducibility across random seeds is 21%, and individual features should be interpreted as illustrative rather than stable.
Figures & tables
Figure 1: Explained variance U-curve across transformer layers. Sparse autoencoder reconstruction quality (explained variance) at each of 10 extraction points for INSPECT (outpatient) and MIMIC-IV (ICU). Both datasets show a characteristic U-curve: near-perfect reconstruction at the input layer (layer 0), a trough at layer 6 corresponding to the most distributed representations, and partial recovery toward the output.
Layer
INSPECT
MIMIC-IV
0 (input)
1.000
1.000
3
0.891
0.919
6 (trough)
0.848
0.895
9 (post-ln_f)
0.860
0.884
Table 1: Explained variance across selected layers (TopK SAE, expansion 8 × , K = 16). The U-curve is consistent across INSPECT and MIMIC-IV, with the trough at layer 6 corresponding to the most distributed representations under this SAE dictionary. The development dataset sweep is shown in Supplementary Table S2.
Figure 2: Concept emergence from layer 0 to layer 6. Four metrics quantifying the transition from token-level to concept-level representation in INSPECT SAEs. Layer-0 features are near-perfect token detectors (45.7% singleton, 2.9 mean tokens per feature, 85.3% single-category). Layer-6 features span approximately 30 tokens each across multiple clinical categories (0.5% singleton, 29.9 mean tokens, 4.4% single-category), representing a ten-fold increase in representational complexity.
Metric
Layer 0
Layer 6
Singleton features (%)
45.7
0.5
Mean tokens per feature
2.9
29.9
Single-category features (%)
85.3
4.4
Category entropy (nats)
0.06
0.43
Table 2: Feature complexity evolution (INSPECT SAEs). By layer 6, features have progressed from near-perfect token detectors to distributed clinical concepts spanning multiple categories.
Dataset
Mortality Δ AUC
LoS ΔR2
INSPECT
+0.118
− 0.540
MIMIC-IV
+0.022
− 0.026
eICU-CRD (zero-shot)
+0.031
− 0.025
Dev (diabetes)
+0.075
− 0.908
Table 3: Task-dependent specialisation at layer 6. Δ = SAE minus dense representation performance. Positive values indicate SAE superiority. Full-sequence values used for consistency across datasets.
Figure 3: Task-dependent specialisation with K -sensitivity analysis. Mortality AUC (blue) and length-of-stay R2 (red) for SAE features at three sparsity levels ( K = 16, 64, 128) versus K -independent dense layer representation performance (dashed lines). INSPECT, layer 6, full-sequence. The mortality advantage of SAE features increases with K (0.929 to 0.960) while the length-of-stay disadvantage widens ( R2 from 0.133 to − 0.152).
K
Mortality AUC (SAE)
LoS R2 (SAE)
LoS R2 (Dense)
16
0.929
0.133
0.673
64
0.938
0.069
0.673
128
0.960
− 0.152
0.673
Table 4: K -sensitivity analysis (INSPECT, layer 6, full-sequence). Mortality advantage of SAE features increases with K , while length-of-stay disadvantage widens . Dense representation performance is K -independent.
Figure 4: Time-windowed mortality prediction across ICU and outpatient settings. AUC for in-hospital mortality prediction at multiple observation windows on eICU-CRD (48h, full), MIMIC-IV (48h, full), and INSPECT (48h, 365d, 1y/3y, full). Full-sequence AUCs are inflated by end-of-stay events that encode the outcome (marked with asterisks). At leakage-safe windows, dense layer representations (orange) match or exceed SAE features (blue), with both outperforming bag-of-tokens baselines (grey) on most datasets.
Dataset
Window
SAE AUC [95% CI]
Dense AUC [95% CI]
BoT AUC
eICU-CRD
48h
0.871 [0.843, 0.903]
0.880 [0.857, 0.903]
0.827
eICU-CRD
Full*
0.998
0.967
0.855
MIMIC-IV
48h
0.836 c
0.914 c
0.864
MIMIC-IV
Full*
0.997
0.975
0.945
Table 5: Time-windowed mortality prediction on ICU datasets. SAE = sparse autoencoder features (layer 6). Dense = dense layer representations. BoT = bag-of-tokens baseline ([DEATH] excluded). eICU-CRD: base model zero-shot transferred from MIMIC-IV; SAE dictionary and downstream probes trained on eICU activations/data. Bootstrap 95% CIs shown for the primary eICU 48h comparison. c MIMIC-IV 48h CIs not available from stored results; the 0.078 gap between SAE and dense is large relative to typical bootstrap width at this sample size. *Full-sequence results contaminated by end-of-stay information leakage—reported for transparency only.
Window
SAE AUC
Dense AUC
BoT AUC
Deaths
Primary leakage-safe evaluation
1y obs / 3y outcome
0.697
0.800
0.785
41
Underpowered exploratory windows
48h
0.587
0.630
0.625
∼ 3
365d (1 year)
0.675
0.694
0.615
∼ 20
Full*
0.929
0.811
0.767
53
Table 6: Mortality prediction on INSPECT (outpatient). SAE = sparse autoencoder features (layer 6). Dense = dense layer representations. BoT = bag-of-tokens. The 1-year observation / 3-year outcome window is the primary leakage-safe evaluation. *Full-sequence result includes the complete clinical trajectory.
Experiment
Mode
Targeted Δ DiD
Random mean Δ DiD
Ratio
Layer 6 reconstruct
Full SAE
+0.034
+0.011
3.1 ×
Layer 6 delta
Delta-mode
− 0.010
+0.037
opposite
Layer 9 delta
Delta-mode
− 0.039
− 0.022
1.8 ×
Layer 9 attenuation × 5
Scaled
− 0.140
− 0.055
2.5 ×
Table 7: Feature perturbation experiments on the warfarin → INR association. Δ DiD = change in warfarin effect size relative to unperturbed model; negative = attenuation (expected for ablation of INR-encoding features). Targeted features encode high-INR concepts; random features are size-matched controls. Ratio = ∣ targeted ∣ / ∣ random mean ∣ .
Dataset
Setting
Hospitals
N
Vocabulary
Mortality
Access
Dev (diabetes)
Outpatient
1 system
61K patients
76 tokens
—
Proprietary
INSPECT
Outpatient
1
19K patients
220 tokens
2.8%
PhysioNet
MIMIC-IV
ICU
1
425K admissions
239 tokens ‡
2.6%
PhysioNet
eICU-CRD
ICU
208
200K stays
239 tokens †
5.3%
PhysioNet
Table 8: Dataset characteristics. Dev = proprietary development dataset used for architecture design. ‡ The SAE and mortality analyses on MIMIC-IV used the drug-level 239-token vocabulary (6 medication classes expanded to 33 individual drugs), which is why individual-drug tokens appear in the feature case studies; the companion paper’s generation and pharmacological testing used the 220-token class-level variant. † eICU-CRD uses the same 239-token vocabulary with MIMIC-derived quintile boundaries; the base model is zero-shot transfer (no adaptation), but SAE dictionaries were trained on eICU activations from this transferred model.