Distribution Shift Aware Neural Tabular Learning

2025-08-26

AI Summary

To address performance degradation in tabular learning caused by distribution shift between training and test data, this paper formally defines the novel task of Distribution-Shifted Tabular Learning (DSTL). We propose SAFT (Shift-Aware Feature Transformation), a differentiable, end-to-end framework that replaces discrete feature search with continuous, learnable feature generation. SAFT integrates four core mechanisms: embedding decorrelation, sample reweighting, suboptimal embedding averaging, and normalized distribution alignment. Extensive experiments across diverse real-world distribution shift scenarios—including covariate shift, concept drift, and domain shift—demonstrate that SAFT significantly enhances model robustness, generalization, and predictive accuracy. Compared to state-of-the-art tabular learning methods, SAFT achieves consistent and substantial improvements across multiple benchmarks and evaluation metrics, establishing a new foundation for distribution-robust tabular modeling.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Tabular learning transforms raw features into optimized spaces for downstream tasks, but its effectiveness deteriorates under distribution shifts between training and testing data. We formalize this challenge as the Distribution Shift Tabular Learning (DSTL) problem and propose a novel Shift-Aware Feature Transformation (SAFT) framework to address it. SAFT reframes tabular learning from a discrete search task into a continuous representation-generation paradigm, enabling differentiable optimization over transformed feature sets. SAFT integrates three mechanisms to ensure robustness: (i) shift-resistant representation via embedding decorrelation and sample reweighting, (ii) flatness-aware generation through suboptimal embedding averaging, and (iii) normalization-based alignment between training and test distributions. Extensive experiments show that SAFT consistently outperforms prior tabular learning methods in terms of robustness, effectiveness, and generalization ability under diverse real-world distribution shifts.

Problem

• Addressing performance deterioration in tabular learning under distribution shifts
• Transforming discrete feature search into continuous differentiable optimization
• Ensuring robustness through decorrelation, reweighting and distribution alignment

Innovation

• Shift-Aware Feature Transformation framework
• Differentiable optimization over feature sets
• Embedding decorrelation and sample reweighting