# Distribution Shift Aware Neural Tabular Learning

arXiv: https://arxiv.org/abs/2508.19486 · HTML: https://mlsanity.com/paper/2508.19486/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 45 references that are in the map.
- First posted on arXiv 27 Aug 2025.

## Builds on (most cited first, 28 of 28)

- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [OpenML: networked science in machine learning](https://mlsanity.com/paper/1407.7722/) (2014)
- [Data-centric Artificial Intelligence: A Survey](https://mlsanity.com/paper/2303.10158/) (2023)
- [Feature Engineering for Predictive Modeling using Reinforcement Learning](https://mlsanity.com/paper/1709.07150/) (2017)
- [The autofeat Python Library for Automated Feature Engineering and Selection](https://mlsanity.com/paper/1901.07329/) (2019)
- [Stable Learning via Sample Reweighting](https://mlsanity.com/paper/1911.12580/) (2019)
- [Continual Learning on Dynamic Graphs via Parameter Isolation](https://mlsanity.com/paper/2305.13825/) (2023)
- [Group-wise Reinforcement Feature Generation for Optimal and Explainable Representation Space Reconstruction](https://mlsanity.com/paper/2205.14526/) (2022)
- [Self-optimizing Feature Generation via Categorical Hashing Representation and Hierarchical Reinforcement Crossing](https://mlsanity.com/paper/2309.04612/) (2023)
- [Evolutionary Large Language Model for Automated Feature Transformation](https://mlsanity.com/paper/2405.16203/) (2024)
- [Distribution Shift Inversion for Out-of-Distribution Prediction](https://mlsanity.com/paper/2306.08328/) (2023)
- [Reinforcement-Enhanced Autoregressive Feature Transformation: Gradient-steered Search in Continuous Space for Postfix Expressions](https://mlsanity.com/paper/2309.13618/) (2023)
- [Feature Selection as Deep Sequential Generative Learning](https://mlsanity.com/paper/2403.03838/) (2024)
- [CT-PatchTST: Channel-Time Patch Time-Series Transformer for Long-Term Renewable Energy Forecasting](https://mlsanity.com/paper/2501.08620/) (2025)
- [Knockoff-Guided Feature Selection via A Single Pre-trained Reinforced Agent](https://mlsanity.com/paper/2403.04015/) (2024)
- [Neuro-Symbolic Embedding for Short and Effective Feature Selection via Autoregressive Generation](https://mlsanity.com/paper/2404.17157/) (2024)
- [A Survey on Data-Centric AI: Tabular Learning from Reinforcement Learning and Generative AI Perspective](https://mlsanity.com/paper/2502.08828/) (2025)
- [Reinforcement Feature Transformation for Polymer Property Performance Prediction](https://mlsanity.com/paper/2409.15616/) (2024)
- [Sculpting Features from Noise: Reward-Guided Hierarchical Diffusion for Task-Optimal Feature Transformation](https://mlsanity.com/paper/2505.15152/) (2025)
- [Topology-aware Reinforcement Feature Space Reconstruction for Graph Data](https://mlsanity.com/paper/2411.05742/) (2024)
- [Unsupervised Feature Transformation via In-context Generation, Generator-critic LLM Agents, and Duet-play Teaming](https://mlsanity.com/paper/2504.21304/) (2025)
- [AdaER: An Adaptive Experience Replay Approach for Continual Lifelong Learning](https://mlsanity.com/paper/2308.03810/) (2023)
- [Agentic Feature Augmentation: Unifying Selection and Generation with Teaming, Planning, and Memories](https://mlsanity.com/paper/2505.15076/) (2025)
- [Enhancing Customer Contact Efficiency with Graph Neural Networks in Credit Card Fraud Detection Workflow](https://mlsanity.com/paper/2504.02275/) (2025)
- [LLM-ML Teaming: Integrated Symbolic Decoding and Gradient Search for Valid and Stable Generative Feature Transformation](https://mlsanity.com/paper/2506.09085/) (2025)
- [Decorrelate Irrelevant, Purify Relevant: Overcome Textual Spurious Correlations from a Feature Perspective](https://mlsanity.com/paper/2202.08048/) (2022)

## Built on by (most cited first, 4 of 4)

- [DELTA: Variational Disentangled Learning for Privacy-Preserving Data Reprogramming](https://mlsanity.com/paper/2509.00693/) (2025)
- [Permutation-Invariant Representation Learning for Robust and Privacy-Preserving Feature Selection](https://mlsanity.com/paper/2510.05535/) (2025)
- [Dataforge: Agentic Platform for Autonomous Data Engineering](https://mlsanity.com/paper/2511.06185/) (2025)
- [Evolving Demonstration Optimization for Chain-of-Thought Feature Transformation](https://mlsanity.com/paper/2603.09987/) (2026)

## Abstract (arXiv)

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.