SEMF: Supervised Expectation-Maximization Framework for Predicting Intervals

Abstract: This work introduces the Supervised Expectation-Maximization Framework (SEMF), a versatile and model-agnostic approach for generating prediction intervals with any ML model. SEMF extends the Expectation-Maximization algorithm, traditionally used in unsupervised learning, to a supervised context, leveraging latent variable modeling for uncertainty estimation. Through extensive empirical evaluation of diverse simulated distributions and 11 real-world tabular datasets, SEMF consistently produces narrower prediction intervals while maintaining the desired coverage probability, outperforming traditional quantile regression methods. Furthermore, without using the quantile (pinball) loss, SEMF allows point predictors, including gradient-boosted trees and neural networks, to be calibrated with conformal quantile regression. The results indicate that SEMF enhances uncertainty quantification under diverse data distributions and is particularly effective for models that otherwise struggle with inherent uncertainty representation.

• The paper introduces SEMF, extending the EM algorithm to supervised learning to produce robust prediction intervals despite incomplete data.
• It integrates an encoder-decoder architecture with Monte Carlo sampling to optimize latent representations for improved uncertainty quantification.
• Experimental results demonstrate that SEMF outperforms traditional methods, achieving higher coverage and narrower intervals across diverse datasets.

Overview of the Supervised Expectation-Maximization Framework (SEMF)

The paper presents the Supervised Expectation-Maximization Framework (SEMF), extending the foundational Expectation-Maximization (EM) algorithm into a supervised learning context. It aims to generate both point predictions and prediction intervals (PIs) for datasets that may contain missing values. This framework offers a novel, model-agnostic approach to uncertainty quantification.

Methodology

SEMF leverages the EM algorithm to derive latent representations optimal for supervised prediction tasks. The traditional use of the EM algorithm in clustering and unsupervised learning has been adapted here to incorporate supervised learning, particularly benefiting the generation of robust PIs. The framework encompasses an encoder-decoder structure, and its architecture is designed to handle incomplete data effectively.

Central to SEMF is the decomposition of the likelihood function and the use of Monte Carlo (MC) sampling to estimate the integral terms. These components enable the EM algorithm to function iteratively in the supervised domain, thus converging to parameter estimates that maximize the log-likelihood.

The algorithm processes datasets with multiple inputs, recognizing missing data patterns by simulating the missing values via empirical models. This often involves the use of models like neural networks for the encoder $g_{\phi}$, decoder $p_{\theta}$, and missing data handling $p_{\xi}$. For training, the framework iterates over batched data, optimizing the objective function using weighted samples of the latent space.

Experimental Validation

The paper implements SEMF across 11 tabular datasets from the OpenML-CTR23 suite, selected for their varying sizes and feature complexity. The experiments cover both complete and 50% artificially missing data scenarios. Baseline comparisons include traditional models like XGBoost, Extremely Randomized Trees (ET), and neural networks (MLPs), highlighting the versatility of SEMF in integrating with any machine learning algorithm.

The results show SEMF's potential in generating intervals with more reliable coverage probabilities and narrower widths compared to traditional quantile regression methods. For example, MultiXGB (SEMF with XGBoost) achieves higher Coverage-Width Ratio (CWR) and better PICP across datasets such as "cpu_activity" and "naval_propulsion_plant," outperforming baseline estimations by significant margins. The framework adapts well to normal and latent representations, despite constraints imposed by non-normal data distributions in some datasets.

Implications and Future Directions

The introduction of SEMF marks a step forward in supervised learning by understanding the uncertainties inherent in predictions. This is crucial for high-stakes domains like healthcare and finance, where decisions hinge on the model's certainty. SEMF's robustness against missing data makes it practically relevant for real-world applications, ensuring reliable predictions even when complete data is unavailable.

However, this research highlights some of the framework's limitations. The inherent complexity of the EM algorithm and its reliance on normality assumptions may constrain the application scope. Future research should look into extending SEMF to handle diverse distributions, potentially integrating conformal prediction techniques to enhance interval coverage reliability further.

Additionally, future work could explore applying SEMF in multi-modal settings, extending its flexibility to heterogeneous data types, including image and text data. Optimizing computational efficiency remains a critical area, specifically for large-scale applications.

Conclusion

The SEMF framework provides a comprehensive approach to generating prediction intervals in supervised learning, especially in the presence of missing data. By adapting the EM algorithm for supervised learning tasks and leveraging MC sampling, SEMF succeeds in offering narrower intervals with better coverage. The extensive empirical validation across multiple datasets demonstrates its efficacy and potential to significantly enhance uncertainty quantification in various domains. While further refinements and optimizations are necessary, SEMF lays the groundwork for advancing robust predictive modeling in machine learning.