Adaptive wavelet base selection for deep learning-based ECG diagnosis: A reinforcement learning approach

Abstract

Electrocardiogram (ECG) signals are crucial in diagnosing cardiovascular diseases (CVDs). While wavelet-based feature extraction has demonstrated effectiveness in deep learning (DL)-based ECG diagnosis, selecting the optimal wavelet base poses a significant challenge, as it directly influences feature quality and diagnostic accuracy. Traditional methods typically rely on fixed wavelet bases chosen heuristically or through trial-and-error, which can fail to cover the distinct characteristics of individual ECG signals, leading to suboptimal performance. To address this limitation, we propose a reinforcement learning-based wavelet base selection (RLWBS) framework that dynamically customizes the wavelet base for each ECG signal. In this framework, a reinforcement learning (RL) agent iteratively optimizes its wavelet base selection (WBS) strategy based on successive feedback of classification performance, aiming to achieve progressively optimized feature extraction. Experiments conducted on the clinically collected PTBXL dataset for ECG abnormality classification show that the proposed RLWBS framework could obtain more detailed time-frequency representation of ECG signals, yielding enhanced diagnostic performance compared to traditional WBS approaches.