KAI: A Scalable Kalman-Attention Imputation Method for Robust Inference in Probabilistic Data
摘要
The lack of data in multivariate time series poses significant challenges for statistical inference and predictive modeling, especially in high-dimensional nonlinear dynamics and irregular patterns of missing data. Recent deep learning approaches, such as BRITS and SAITS, improve accuracy by leveraging recurrent and attention-based architectures, but they remain computationally demanding and lack explicit probabilistic calibration. To address these limitations, we propose KAI (Kalman Attention Imputation), a fully differentiable hybrid model that integrates Kalman smoothing with a self-attention mechanism with mask recognition. KAI combines local linear-Gaussian filtering for uncertainty propagation with global attention for long-range dependencies, employing uncertainty-driven gate activation and covariance consistency regularization. Through extensive simulation studies with various data generation processes (AR, ARMA, GARCH, VAR) and missing data mechanisms (MCAR, MAR, MNAR), KAI demonstrates competitive performance in point accuracy (RMSE, MAE), restoration of dynamic properties (autocorrelation, stationarity), and interval coverage under multiple imputation. These results show KAI as a scalable and robust solution for probabilistic imputation in complex time series.