MARFI: Missingness-Aware Rectified-Flow Imputation for MNAR Data in Predictive Maintenance

Abstract

Missing data frequently occur for various reasons and can significantly degrade the performance and reliability of data-driven models. In predictive maintenance environments, missing data may be associated with equipment operating conditions, sensor anomalies, and data-collection processes. These characteristics make Missing Not at Random (MNAR) an important missingness mechanism to consider when evaluating imputation methods in predictive maintenance. Nevertheless, existing missing data imputation studies have not sufficiently leveraged the structural information that the missingness mask can provide in such environments. To address this limitation, this paper proposes Missingness-Aware Rectified-Flow Imputation (MARFI). MARFI incorporates missingness-mask information into a rectified-flow-based restoration process by conditioning velocity prediction on both observed feature values and the missingness mask. It further introduces a Mask Prediction module that provides an auxiliary consistency objective during training. Experimental results on six predictive maintenance datasets using controlled missingness simulations show that MARFI achieved the best average MAE under MAR, MBOV, and MBUV settings. In particular, MARFI achieved relative MAE reductions of approximately 12.27% under the MBOV setting and 6.30% under the MBUV setting compared with the second-best models. Furthermore, downstream supervised classification experiments show that the data restored by MARFI can preserve information useful for subsequent predictive maintenance tasks. MARFI achieved the highest overall average F1-score in both the controlled MNAR downstream evaluation and the naturally occurring missing-value evaluation, with relative improvements of 0.74% and 1.19%, respectively, over the corresponding second-best methods. These results indicate that explicitly using the missingness mask in the restoration process can improve numerical imputation performance while also preserving information useful for downstream prediction under missing-data settings relevant to predictive maintenance.