Research map / Label-efficient and robust learning / Recognition and long-tail
To Predict or Not to Predict? Proportionally Masked Autoencoders for Tabular Data Imputation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PMAE, a masked autoencoder for tabular imputation that masks columns in proportion to their observed rates via a logit-based masking function, plus MLP-Mixer token mixing, and reports gains over ReMasker and other baselines.
Problem
Uniform random masking in masked autoencoders disrupts the column-wise missingness distribution of heterogeneous tabular data, leading to suboptimal imputation, especially for general (non-monotone) missing patterns.
Approach
Compute column-wise observed proportions and apply a logit-transformed masking function so that columns with fewer observations are masked (prioritizing prediction) less or more as designed, with loss balanced by inverse observed proportions; replace self-attention token mixing with MLP-Mixer blocks.
Key results
- PMAE-mix reaches 44.2 average imputation accuracy versus 34.7 for ReMasker (27.3% gain for PMAE-trf over ReMasker) across 9 datasets and 3 patterns under MNAR.
- PMAE outperforms ReMasker by up to 34.1% on the General missing pattern.
- Ablation: replacing uniform 0.5 masking with the logit masking function raises PMAE-trf from 34.7-level baseline to 43.1 (+20.0% relative to ReMasker-style loss setup).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use proportional logit masking instead of constant masking (ablation) | PMAE · 9 datasets, 3 missing patterns, MNAR | Imputation accuracy: 34.7 -> 43.1 | Helped |
| Replace the ReMasker loss with the proposed loss (ablation) | ReMasker · Shoppers, Wine, Diabetes | Imputation accuracy: 34.7 -> 36.2 | Helped |
| Replace Transformer with MLP-Mixer (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 43.1 -> 44.2 | Helped |
| Use proportional logit masking rather than constant masking (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 33.4 -> 35.4 | Helped |
| Remove reconstruction from the training objective (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 35.4 -> 7.7 | Hurt |
| Remove prediction from the training objective (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 35.4 -> 8.6 | Hurt |
| Use linear masking instead of logit masking (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 35.4 -> 34.8 | Hurt |
| Reverse masking direction to mask in proportion to observed proportion (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 35.4 -> 31.9 | Hurt |
| Use piece-wise masking instead of logit masking (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 35.4 -> 33.6 | Hurt |
| Use sigmoid-like masking with reversed concavities (ablation) | PMAE-trf · Shoppers, Wine, Diabetes | Imputation accuracy: 35.4 -> 33.7 | Hurt |
| Use MLP-Mixer instead of Transformer | PMAE · 9 datasets, 3 missing patterns, MNAR | DT (reg): 25.2 -> 24.8 | Hurt |
Surprising findings
- MLP-Mixer token mixing (PMAE-mix) outperforms Transformer-based PMAE-trf, contrary to the common assumption that attention is superior.
- The paper reports the rank of ReMasker shifts when using RMSE, showing metric choice changes method ordering.
Limitations
- Does not analyze the relationship between covariance structure and the proportional masking scheme.
- Evaluated only with semi-synthetic missingness on nine tabular datasets.
- Hyperparameters a and b are fixed from a grid search on validation datasets rather than learned.
Experimental setup
Datasets: Diabetes Wine Obesity Bike Shoppers Letter Default News Adult
Benchmarks: Diabetes Wine Obesity Bike Shoppers Letter Default News Adult
Metrics: Imputation Accuracy R^2 Accuracy RMSE downstream task AUROC/R^2
Baselines: Naive KNN EM MissForest MiceForest MIWAE GAIN MIRACLE HyperImpute TDM ReMasker IGRM
Builds on: ReMasker MLP-Mixer inverse propensity score weighting
Closest papers by meaning
- ReMasker: Imputing Tabular Data with Masked Autoencoding (ICLR 2024 · similarity 0.93)
- CACTI: Leveraging Copy Masking and Contextual Information to Improve Tabular Data Imputation (ICML 2025 · similarity 0.90)
- Hybrid Autoencoders for Tabular Data: Leveraging Model-Based Augmentation in Low-Label Settings (NeurIPS 2025 · similarity 0.86)
- Discretizing Continuous Time Series for Imputation with Masked Diffusion Training (NeurIPS 2026 · similarity 0.85)
- Self-Guided Masked Autoencoder (NeurIPS 2024 · similarity 0.85)
- Learning What Not to Impute: An Uncertainty-Aware Diffusion Framework for Meaningful Missingness (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.