Uncertainty-Aware Self-Attention Model for Time Series Prediction with Missing Values

Missing values in time series data present a significant challenge, often degrading the performance of downstream tasks such as classification and forecasting. Traditional approaches address this issue by first imputing the missing values and then independently solving the predictive tasks. Recent methods have leveraged self-attention models to enhance imputation quality and accelerate inference. These models, however, predict values based on all input observations—including the missing values—thereby potentially compromising the fidelity of the imputed data. In this paper, we propose the Uncertainty-Aware Self-Attention (UASA) model to overcome these limitations. Our approach introduces two novel techniques: (i) A self-attention mechanism with a partially observed diagonal that effectively captures complex non-local dependencies in time series data—a characteristic also observed in fractional-order systems. This approach draws inspiration from fractional calculus, where non-integer-order derivatives better characterize complex dynamical systems with long-memory effects, providing a more comprehensive mathematical framework for handling temporal data. And (ii) uncertainty quantification in data imputation to better inform downstream tasks. The UASA model comprises an upstream component for data imputation and a downstream component for time series prediction, trained jointly in an end-to-end fashion to optimize both imputation accuracy and task-specific objectives simultaneously. For classification tasks, the UASA model demonstrates remarkable performance even under high missing data rates, achieving a ROC-AUC of 99.5%, a PR-AUC of 58.5%, and an F1-SCORE of 49.3%. For forecasting tasks on the AUST-Gait dataset, the UASA model achieves a Mean Squared Error (MSE) of 0.72 under 0% missing data conditions (i.e., complete data input). Under the end-to-end training strategy evaluated across all missing data rates, the model achieves an average MSE of 0.74, showcasing its adaptability and robustness across diverse missing data scenarios.