Evaluating Bayesian deep learning for robust biomedical time-series classification under distribution shift: A sleep scoring case study
Organizations: Technical University of Denmark, Denmark · H. Lundbeck A/S, Denmark
Abstract
In biomedical time-series classification, machine learning models often suffer from deteriorating accuracy and calibration due to distribution shift. We investigate to what degree Bayesian-inspired deep learning improves robustness through Bayesian model averaging (BMA). Using rodent sleep scoring as a case study, we evaluate BMA under distribution shifts in genotype, laboratory, and species. We systematically compare nine approximate inference methods spanning single- and multi-basin approaches across two neural network backbones on in-distribution (ID) data and five out-of-distribution (OOD) datasets. In our results, multi-basin methods perform best, outperforming baselines and single-basin methods. Compared with point estimates, BMA generally improves OOD accuracy and calibration across both architectures, while yielding only modest ID gains. It also often improves selective prediction and OOD and artifact detection, but, interestingly, not inter-rater disagreement detection.
Figures & tables
| Method | |
|---|---|
| MAP | |
| LLLA | |
| SWA | |
| SWAG | |
| VGD | |
| MFVI |
| Cohort | Use | Species | Primary shift | |
|---|---|---|---|---|
| WT | Train | Mouse | 169 | — |
| WT | Val. | Mouse | 36 | — |
| WT | ID test | Mouse | 36 | None |
| NT1 | OOD test | Mouse | 28 | Genotype |
| spindle -A | OOD test | Mouse | 4 | Lab |
| spindle -D | OOD test | Mouse | 6 | Lab |
| Model | MF1 | LPD | Brier | |
|---|---|---|---|---|
| U-Sleep | ||||
| MAP | ||||
| Temp. | ||||
| SWA | ||||
| SWAG | ||||
| Laplace | ||||
| Model | MF1 | LPD | Brier | |
|---|---|---|---|---|
| U-Sleep | ||||
| MAP | ||||
| Temp. | ||||
| SWA | ||||
| SWAG | ||||
| Laplace | ||||
| Single solution/trajectory | Multi-basin | Other baselines | |||||||||||
| Scope | MAP | Temp. | SWA | SWAG | Laplace | VGD | MFVI | SGLD | DE | MultiSWA | MultiSWAG | Somnotate | Spindle |
| ID | 17.5 | 16.5 | 6.5 | 6.5 | 18.5 | 15.0 | 19.5 | 12.0 | 9.0 | 2.5 | 3.0 | 23.0 | 24.0 |
| OOD | 18.1 | 16.0 | 13.0 | 9.1 | 18.5 | 12.6 | 13.5 | 9.9 | 6.2 | 6.8 | 4.7 | 23.6 | 19.6 |
| ID + OOD | 18.0 | 16.1 | 11.9 | 8.7 | 18.5 | 13.0 | 14.5 | 10.2 | 6.7 | 6.1 | 4.4 | 23.5 | 20.3 |
| AUROC | AURC | |||
| Model | OOD | Artifact | Disagree | Reject |
| U-Sleep | ||||
| MAP | ||||
| Temp. | ||||
| SWA | ||||
| SWAG | ||||
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
| Dataset | Species/genotype | Source | Channels | Role | Shift w.r.t. training | |
|---|---|---|---|---|---|---|
| Training and validation | ||||||
| WT | Mouse/wild-type | Lundbeck | ECoG, LFP, EMG | 169 | Train | — |
| WT | Mouse/wild-type | Lundbeck | ECoG, LFP, EMG | 36 | Validation | None |
| Test | ||||||
| WT | Mouse/wild-type | Lundbeck | ECoG, LFP, EMG | 36 | Test (ID) | None |
| NT1 | Mouse/NT1 | Lundbeck | ECoG, LFP, EMG | 28 | Test (OOD) | Genotype |
| Model | Architecture | Input | Context |
|---|---|---|---|
| Somnotate | LDA+HMM | Spectrogram | Full recording |
| Spindle | CNN+HMM | Spectrogram | Full recording |
| U-Sleep | Fully convolutional | Raw signals | 35 windows |
| CNN-GRU | CNN+GRU | Raw signals | 35 windows |
| Model | Architecture | Parameters | ||
|---|---|---|---|---|
| U-Sleep | 9 blocks, | M | ||
| CNN-GRU | Encoder , GRU | M |
| Method | CNN-GRU | U-Sleep | Training/collection budget | Shared settings |
|---|---|---|---|---|
| MAP | 2000 blocks | Final checkpoint; deterministic prediction | ||
| SWA / SWAG | 5 transition 200 collection blocks | Same trajectory for SWA and SWAG; one iterate per collection block; ; SWAG ; BN statistics recomputed using of the training data | ||
| LLLA | Full covariance | Full covariance | – | MAP prior variance; GLM predictive with Monte Carlo samples |
| VGD | 3000 blocks | Random initialization; MAP prior retained; , , final ; learning rate | ||
| MFVI | 3000 blocks | MAP initialization of posterior means; MAP prior retained; , , final ; learning rate | ||
| SGLD | 3000 blocks | MAP initialization and prior; burn-in 1000 blocks; states stored every 10 blocks; selected uniformly from 200 post-burn-in states |
| Model | Acc. | MF1 | LPD | Brier | ECE | Entropy | |
|---|---|---|---|---|---|---|---|
| U-Sleep | |||||||
| MAP | |||||||
| Temp. | |||||||
| SWA | |||||||
| SWAG | |||||||
| Laplace | |||||||
| Model | Acc. | MF1 | LPD | Brier | ECE | Entropy | |
|---|---|---|---|---|---|---|---|
| U-Sleep | |||||||
| MAP | |||||||
| Temp. | |||||||
| SWA | |||||||
| SWAG | |||||||
| Laplace | |||||||
| Single solution/trajectory | Multi-basin | Other baselines | |||||||||||
| Scope | MAP | Temp. | SWA | SWAG | Laplace | VGD | MFVI | SGLD | DE | MultiSWA | MultiSWAG | Somnotate | Spindle |
| ID | 19.5 | 19.5 | 8.0 | 7.0 | 21.0 | 13.0 | 15.0 | 10.5 | 9.5 | 3.5 | 1.5 | 20.0 | 24.0 |
| OOD | 16.6 | 16.6 | 12.9 | 10.1 | 17.2 | 13.6 | 13.6 | 10.9 | 7.4 | 8.4 | 7.0 | 16.8 | 14.6 |
| ID + OOD | 17.1 | 17.1 | 12.1 | 9.6 | 17.8 | 13.5 | 13.8 | 10.8 | 7.8 | 7.6 | 6.1 | 17.3 | 16.2 |
| AUROC | AURC | |||||||||||
| OOD | Artifact | Disagree | Reject | |||||||||
| Model | TU | EU | AU | TU | EU | AU | TU | EU | AU | TU | EU | AU |
| U-Sleep | ||||||||||||
| MAP | ||||||||||||
| Temp. | ||||||||||||
| SWA | ||||||||||||
| Entropy ( ) | Detection AUROC | AURC | ||||||||
| Model | ID | OOD | Agree | Disagree | No artifact | Artifact | OOD | Disagree | Artifact | OOD |
| U-Sleep | ||||||||||
| MAP | ||||||||||
| Temp. | ||||||||||
| SWA | ||||||||||
| SWAG | ||||||||||
Distribution Shift Robustness
Distribution shift robustness studies how to keep machine-learning predictions reliable when deployment data differ from training data.
See all 285 papers