Cross-Attention Alignment and Spatiotemporal Augmentation for Cross-Site Screening of Cognitive Impairment from Polysomnography

Leyu Zhou1, Donglin Xie2, Guangkun Nie3, Jinshuai Gu4, Qingshuo Zhao5, Shijie Chen6, Deyun Zhang7, Yuxi Zhou3, Shenda Hong8
1School of Intelligence Science and Technology University of Science and Technology Beijing, 2National Institute of Health Data Science, Peking University, 3Peking University, 4National Institute of Health Data Science, Peking University, Beijing, China Department of Computer Science, Tianjin University of Technology, Tianjin, China., 5School of Computer Science and Engineering, Tianjin University of Technology, 6Anhui Heartvoice Medical Technology Co., 7Heartvoice Medical Technology, 8Georgia Institute of Technology

Abstract

Sleep polysomnography (PSG) provides rich neurophysiological and cardiopulmonary information and may support early screening for future cognitive impairment. However, this task is challenging because PSG recordings vary substantially across sites in channel configuration, signal characteristics, recording duration, and annotation quality. We developed a cross-site screening framework that combines pretrained sleep representations, annotation-aware alignment, and spatiotemporal augmentation within a downstream classification model.

Sequential PSG embeddings are first extracted from pretrained sleep foundation features, while algorithmic sleep annotations are converted into temporally aligned annotation features. A cross-attention module is then introduced to align PSG representation tokens with annotation tokens, allowing the model to emphasize task-relevant sleep dynamics and incorporate stage-related contextual information into the fused representation. The aligned features are further aggregated and passed to a record-level classifier for prediction of future cognitive impairment. This design explicitly couples physiological representations with sleep-structure information, improving the model's ability to focus on clinically informative temporal patterns.To enhance robustness under heterogeneous real-world conditions, we further apply spatiotemporal augmentation during training. Temporal augmentation improves tolerance to local shifts and variability in informative sleep segments, whereas spatial augmentation improves robustness to missing channels and cross-site differences in channel configurations. Together, these components improve both representation alignment and generalization across heterogeneous PSG acquisitions.

On the validation set, our team(Better PSG) achieved an AUROC of 0.584, the official Challenge metric, with accuracy 0.926. These results suggest that cross-attention-based alignment and spatiotemporal augmentation provide a promising strategy for cognitive-impairment screening from routine PSG collected across heterogeneous sites.