LassoFlexNet: Flexible Neural Architecture for Tabular Data
Organizations: RBC Borealis · RBC
Abstract
Despite their dominance in vision and language, deep neural networks often underperform relative to tree-based models on tabular data. To bridge this gap, we incorporate five key inductive biases into deep learning: robustness to irrelevant features, axis alignment, localized irregularities, feature heterogeneity, and training stability. We propose LassoFlexNet, an architecture that evaluates the linear and nonlinear marginal contribution of each input via Per-Feature Embeddings, and sparsely selects relevant variables using a Tied Group Lasso mechanism. Because these components introduce optimization challenges that destabilize standard proximal methods, we develop a Sequential Hierarchical Proximal Adaptive Gradient optimizer with exponential moving averages (EMA) to ensure stable convergence. Across datasets from three benchmarks, LassoFlexNet matches or outperforms leading tree-based models, achieving up to a % relative gain, while maintaining Lasso-like interpretability. We substantiate these empirical results with ablation studies and theoretical proofs confirming the architecture's enhanced expressivity and structural breaking of undesired rotational invariance.
Figures & tables
| Error | Inductive Bias | Solution |
|---|---|---|
| Generalization | Robustness to irrelevant features; axis alignment | Lasso Linear Skip Connection with Tied Group Lasso |
| Optimization | Stable training | Sequential Hierarchical Proximal Adaptive Gradient optimizer + EMA |
| Approximation | Axis alignment, localized irregularities; heterogeneity; complex interactions | PLE; PFE; MLP-Mixer Module |
| Ablation | Test RMSE |
|---|---|
| (Linear) LassoNet | |
| Baseline Linear LassoNet | 0.5263 0.0103 (0.5100, 0.5362) |
| AdamHGD | 0.5294 0.0088 (0.5073, 0.5352) |
| Adam | 0.5291 0.0082 (0.5090, 0.6044) |
| LassoFlexNet | |
| Baseline LassoFlexNet | 0.4047 0.0069 (0.3954, 0.4180) |
| Dataset \Model | CatBoost | LightGBM | RF | XGBoost | TabM-mini | LassoFlexNet |
|---|---|---|---|---|---|---|
| ecom-offers (class. AUC) | 0.6507 ± 0.0554 [6] | 0.6317 ± 0.0538 [6] | 0.6447 ± 0.0529 [6] | 0.6447 ± 0.0628 [6] | 0.6533 ± 0.0603 [6] | 0.6322 ± 0.0419 [6] |
| homecredit-default (class. AUC) | 0.8211 ± 0.0115 [6] | 0.8092 ± 0.0110 [6] | 0.7365 ± 0.0168 [6] | 0.8046 ± 0.0136 [6] | 0.8413 ± 0.0102 [6] | 0.8365 ± 0.0096 [6] |
| homesite-insurance (class. AUC) | 0.9559 ± 0.0088 [6] | 0.9576 ± 0.0075 [6] | 0.8414 ± 0.0168 [6] | 0.9568 ± 0.0088 [6] | 0.9557 ± 0.0083 [6] | 0.9605 ± 0.0071 [6] |
| cooking-time (reg. RMSE) | 0.4718 ± 0.0084 [6] | 0.4899 ± 0.0095 [6] | 0.5192 ± 0.0062 [6] | 0.4866 ± 0.0108 [6] | 0.5629 ± 0.0157 [6] | 0.4658 ± 0.0093 [6] |
| delivery-eta (reg. RMSE) | 0.5370 ± 0.0094 [6] | 0.5641 ± 0.0095 [6] | 0.5701 ± 0.0079 [6] | 0.5425 ± 0.0086 [6] | 0.6121 ± 0.0057 [5] | 0.5375 ± 0.0104 [6] |
| maps-routing (reg. RMSE) | 0.1665 ± 0.0015 [6] | 0.1641 ± 0.0015 [6] | 0.1746 ± 0.0012 [6] | 0.1709 ± 0.0018 [6] | 0.1656 ± 0.0021 [6] | 0.1619 ± 0.0013 [6] |
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
| Hyperparameter | Range | Step/Scale |
|---|---|---|
| Mini batch sizes | 512–2048 | scaled by 2 |
| Number of Mixer blocks | 2–5 | 1 |
| MLP1&2 hidden dimension | 64–512 | scaled by 2 |
| Feature embedding output dimension | 8–32 | 2 |
| Number of bins used in PLE | 4–32 | 2 |
| Tau (scalar multiplied with ) | 0.01–2.0 | Uniform |
| Hyperparameter | Range | Step/Scale |
|---|---|---|
| Mini batch sizes | 1024–4096 | scaled by 2 |
| Number of Mixer blocks | 2–6 | 1 |
| MLP1&2 hidden dimension | 64–512 | scaled by 2 |
| Feature embedding output dimension | 8–32 | 2 |
| Number of bins used in PLE | 16–64 | 2 |
| Tau (scalar multiplied with ) | 0.01–2.0 | Uniform |
Tabular ML
Tabular machine learning develops models and tools for learning from structured data arranged in rows and columns.
See all 115 papers