MIRRAMS: Learning Robust Tabular Models under Unseen Missingness Shifts

TL;DR

MIRRAMS is a deep learning framework that enhances robustness of tabular models against unseen missingness shifts by leveraging mutual information conditions, applicable across various missingness scenarios and semi-supervised settings.

Contribution

The paper introduces MIRRAMS, a novel loss-based framework that enforces mutual information robustness conditions, improving tabular model performance under unseen missingness shifts without relying on specific missingness assumptions.

Findings

Outperforms existing baselines across multiple datasets.

Maintains stable performance under diverse missingness conditions.

Effective even in fully observed data scenarios.

Abstract

The presence of missing values often reflects variations in data collection policies, which may shift across time or locations, even when the underlying feature distribution remains stable. Such shifts in the missingness distribution between training and test inputs pose a significant challenge to achieving robust predictive performance. In this study, we propose a novel deep learning framework designed to address this challenge, particularly in the common yet challenging scenario where the test-time dataset is unseen. We begin by introducing a set of mutual information-based conditions, called MI robustness conditions, which guide the prediction model to extract label-relevant information. This promotes robustness against distributional shifts in missingness at test-time. To enforce these conditions, we design simple yet effective loss terms that collectively define our final objective, called MIRRAMS. Importantly, our method does not rely on any specific missingness assumption such as MCAR, MAR, or MNAR, making it applicable to a broad range of scenarios. Furthermore, it can naturally extend to cases where labels are also missing in training data, by generalizing the framework to a semi-supervised learning setting. Extensive experiments across multiple benchmark tabular datasets demonstrate that MIRRAMS consistently outperforms existing state-of-the-art baselines and maintains stable performance under diverse missingness conditions. Moreover, it achieves superior performance even in fully observed settings, highlighting MIRRAMS as a powerful, off-the-shelf framework for general-purpose tabular learning.