PaperlayerView exact revision on arXiv

Source-linked AI summary

TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023

Yury Gorishniy, Ivan Rubachev, Nikolay Kartashev, Daniil Shlenskii, Akim Kotelnikov, Artem Babenko

arXiv:2307.14338v2cs.LG

TL;DR

Tabular DL is attracting attention, but GBDT remains a longstanding go-to and prior retrieval-based tabular models offer only minor gains over tuned MLPs at greater complexity and cost. TabR adds a customized retrieval component to tabular DL and achieves strong benchmark results, including outperforming GBDT on a small-to-middle-scale benchmark.

  • Problem

    GBDT remains a strong go-to for tabular problems, while prior retrieval-based tabular DL models offer only minor gains over tuned MLPs at greater complexity and cost.

    Source: §1 · pp. 1–2

  • Method

    TabR is a tabular DL model with a customized retrieval component implemented as a single-head attention-like module.

    Source: §1 · p. 2, §2 · p. 3

  • Results

    TabR achieves the best average performance among tabular DL models, reaches state-of-the-art on several datasets, and outperforms GBDT on the small-to-middle-scale benchmark.

    Source: §1 · p. 2, The benchmark from · p. 7, §6 · p. 9

  • Takeaways & Limitations

    Similarity and value modules are important attention-mechanism details with significant impact on retrieval-component performance.

    Source: §6 · p. 9

    The paper focuses more on task performance, leaving other properties underexplored, and notes that greater efficiency is needed for applicability to tens or hundreds of millions of data points.

    Source: §6 · p. 9

Abstract

from arXiv · show

Deep learning (DL) models for tabular data problems (e.g. classification, regression) are currently receiving increasingly more attention from researchers. However, despite the recent efforts, the non-DL algorithms based on gradient-boosted decision trees (GBDT) remain a strong go-to solution for these problems. One of the research directions aimed at improving the position of tabular DL involves designing so-called retrieval-augmented models. For a target object, such models retrieve other objects (e.g. the nearest neighbors) from the available training data and use their features and labels to make a better prediction. In this work, we present TabR -- essentially, a feed-forward network with a custom k-Nearest-Neighbors-like component in the middle. On a set of public benchmarks with datasets up to several million objects, TabR marks a big step forward for tabular DL: it demonstrates the best average performance among tabular DL models, becomes the new state-of-the-art on several datasets, and even outperforms GBDT models on the recently proposed "GBDT-friendly" benchmark (see Figure 1). Among the important findings and technical details powering TabR, the main ones lie in the attention-like mechanism that is responsible for retrieving the nearest neighbors and extracting valuable signal from them. In addition to the much higher performance, TabR is simple and significantly more efficient compared to prior retrieval-based tabular DL models.

1 INTRODUCTION

Tabular DL is gaining attention while GBDT remains a longstanding go-to; the paper introduces TabR as a more effective and efficient retrieval-based alternative to prior tabular DL models.

1 INTRODUCTION

  • GBDT has long been a go-to for tabular ML, while tabular DL has recently attracted more attention and become more competitive.Source: §1 · p. 1
  • Retrieval-augmented models retrieve training objects, such as nearest neighbors, and use their features and labels to improve a target object's prediction.Source: §1 · pp. 1–2
  • Prior retrieval-based tabular DL models provide only minor benefits over a properly tuned MLP while being more complex and costly.The paper argues that overlooked design choices can yield a powerful, simpler, more efficient retrieval-based architecture.Source: §1 · p. 2
  • TabR achieves the best average performance among tabular DL models, reaches state-of-the-art on several datasets, and outperforms GBDT on a small-to-middle-scale benchmark.The paper also reports that tree-based models remain cheaper.Source: §1 · p. 2

2 RELATED WORK

The related work spans parametric tabular DL and retrieval methods; the paper positions TabR as a customized, efficient alternative to existing deep retrieval models.

2 RELATED WORK

  • TabR demonstrates that DL can outperform GBDT on the small-to-middle-scale benchmark that had illustrated GBDT's superiority to parametric DL.The benchmark uses tasks with datasets of 50K objects or fewer.Source: §2 · p. 2
  • Tabular DL research seeks competitive performance on tabular data, with MLP-like backbones and continuous-feature embeddings among the approaches discussed.The paper presents retrieval as a way to boost tabular DL further.Source: §2 · p. 2
  • Retrieval-based models fetch relevant training samples and process them with the input; kNN is a simple example, and such models can offer interpretability and robustness.Source: §2 · p. 2
  • Tabular retrieval methods include classic neighbor and kernel methods as well as deep models; some deep models use all training points as retrieved instances.Source: §2 · pp. 2–3
  • TabR outperforms existing retrieval-based tabular DL models while being significantly more efficient, using one customized single-head attention-like module instead of often-used multi-layer, multi-head attention.Source: §2 · p. 3

3 TABR

The section builds TabR incrementally by adding a retrieval module to a feed-forward tabular model. It defines the supervised-data setup and develops retrieval choices that use neighbor labels and representations to improve predictions.

3.1 PRELIMINARIES

  • The setup covers binary and multiclass classification and regression, splits data into train, validation, and test sets, and uses training objects as shared retrieval candidates.The default benchmark uses datasets from prior literature; the additional middle-scale benchmark contains tasks with at most 50K objects.Source: §3.1 · p. 3

3.2 ARCHITECTURE

  • TabR adds a retrieval module after the encoder, selecting neighbors by similarity and aggregating their values before the predictor makes a prediction.The candidate objects and target share an encoder; the retrieval module enriches the target representation.Source: §3.2 · p. 4, Figure 2 · p. 4, Figure 4 · p. 5
  • The vanilla-attention-like retrieval baseline performs similarly to an MLP, and adding context labels alone yields no improvement.The label contribution is an embedding for classification and a linear layer for regression.Source: §3.2 · p. 5
  • Replacing query-based dot-product similarity with key-only L2 similarity improves performance on several datasets, while removing any of three ingredients returns performance to the MLP level.The three ingredients are context labels, a key-only representation, and L2 distance; the authors call L2 a reasonable default, not a universal best choice.Source: §3.2 · p. 5
  • Making neighbor values depend on the target representation further improves performance on several datasets, with a learned correction translating key-space differences into label-embedding differences.The raw contribution is the context object’s label embedding.Source: §3.2 · p. 5
  • Omitting similarity scaling and excluding the target from its own context improve TabR’s results on average; these are the model’s default recommendations.TabR’s value combines a context-label representation with a correction based on the difference between target and neighbor keys.Source: §3.2 · p. 6, Table 2 · p. 6

4 EXPERIMENTS ON PUBLIC BENCHMARKS

The experiments compare TabR with retrieval-based and parametric deep-learning models, then with GBDT ensembles. TabR improves over tuned GBDT on several datasets while remaining competitive on most others.

4.1 EVALUATING RETRIEVAL-AUGMENTED DEEP LEARNING MODELS FOR TABULAR DATA

  • TabR is the only retrieval-based model reported to significantly boost performance over MLP on many datasets, but its retrieval component is not universally beneficial.The full TabR outperforms MLP-PLR on CA, OT, BL, WE, and CO, performs on par on the rest except MI, and prior retrieval approaches faced technical limitations.Source: §4.1 · p. 6, Table 3 · p. 7
  • Retrieval and numerical-feature embeddings are powerful but not fully interchangeable: which to prefer, or whether to combine them, depends on the dataset.Source: §4.1 · p. 6

4.2 COMPARING TABR WITH GRADIENT-BOOSTED DECISION TREES

  • TabR improves on tuned GBDT on CH, CA, HO, HI, WE, and CO, and is competitive on the remaining datasets except MI.The default configuration is also reported as competitive.Source: §4.2 · p. 7

The benchmark from

On the small-to-middle-scale benchmark, TabR outperforms GBDT on average, although the benchmark was originally used to illustrate GBDT's superiority over parametric deep-learning models. The section also points to additional analyses of retrieval-augmented XGBoost and training times.

The benchmark from

  • TabR outperforms GBDT on average on the small-to-middle-scale benchmark originally used to illustrate GBDT's superiority over parametric deep learning.The comparison uses ensembles and the default TabR-S configuration against default GBDT configurations.Source: The benchmark from · p. 7, Table 5 · p. 8
  • Appendix analyses include adding retrieval to XGBoost and comparing TabR and GBDT training times.Source: The benchmark from · p. 7
  • The comparison confirms TabR as a strong tabular solution, with strong average performance and noticeable improvements over GBDT on some datasets.Source: The benchmark from · p. 7

5 ANALYSIS

The analysis examines context freezing to speed training, adding unseen training data as retrieval candidates without retraining, and further technical properties of TabR.

5.1 FREEZING CONTEXTS FOR FASTER TRAINING OF TABR

  • On Weather prediction, context freezing cuts training time from 18h9min to 3h15min—nearly sevenfold—while maintaining competitive RMSE.After a fixed number of epochs, TabR computes contexts for training objects once more and reuses them for the rest of training; validation and test contexts are not frozen.Source: §5.1 · p. 8, Table 6 · p. 8

5.2 UPDATING TABR WITH NEW TRAINING DATA WITHOUT RETRAINING

  • Adding unseen training data to the retrieval candidates can incorporate new data without retraining, but the authors characterize this experiment as preliminary.They also propose training on a subset and retrieving from the full dataset as a way to scale TabR.Source: §5.2 · p. 8, Figure 6 · p. 9

5.3 FURTHER ANALYSIS

  • Further analyses examine the key-only L2 similarity module, the value module's correction semantics on regression, and training-time efficiency versus prior retrieval models.The appendix also reports that TabR's absolute training times are affordable for most practical scenarios within the considered dataset-size scope.Source: §5.3 · p. 9

6 CONCLUSION & FUTURE WORK

The conclusion presents TabR as a retrieval-augmented tabular deep-learning architecture with strong average performance and state-of-the-art results on several datasets. It identifies efficiency and underexplored properties as directions for future work.

6 CONCLUSION & FUTURE WORK

  • TabR delivers strong average performance and achieves new state-of-the-art results on several datasets; the similarity and value modules significantly affect attention-based retrieval performance.Source: §6 · p. 9
  • Future work includes making retrieval-augmented models applicable to tens or hundreds of millions of data points, while other properties of TabR remain underexplored.Examples include prediction interpretation through context-object influence, continual learning, improved similarity and value modules, and multiple retrieval rounds.Source: §6 · p. 9

SUPPLEMENTARY MATERIAL

The supplementary analyses motivate and examine TabR’s key-only L2 similarity module, show how it uses dataset-specific neighbor signals, and test the value module and component choices. They also report that the label-aligned value correction matters for regression, while the tested hypotheses yielded no interesting classification results.

A.1.1 INTUITIVE MOTIVATION

  • The authors motivate key-only L2 similarity by avoiding query-key representation alignment and using L2, which they consider more suitable than dot product for tabular features.The argument is especially relevant when the encoder is linear and should remain lightweight because it is applied to many candidates.Source: §A.1.1 · p. 13

A.1.2 ANALYZING ATTENTION PATTERNS OVER CANDIDATES

  • Step-2’s key-only L2 similarity module produces more diverse average candidate distributions than vanilla attention.Each input’s distribution has exactly m + 1 non-zero weights, including the target object added to its context.Source: Table 7 · p. 13, §A.1.2 · p. 13

A.1.3 CASE STUDIES

  • Across the three case studies, the new similarity module finds and exploits natural, dataset-specific neighbor signals that provide hints for prediction.The examples include geographic neighbors and near-copies in the inherited Facebook comments split.Source: §A.1.3 · p. 14
  • For California housing, Step-2 assigns about 77% of attention to context objects within 10 km, versus 1% for Step-1, and boosts performance.Step-2 also assigns about 7% of attention to the target itself, versus 0.07% for Step-1.Source: §A.1.3 · p. 14
  • For weather prediction, Step-2 assigns about 25% of attention to context objects within 200 km, versus 0.5% for Step-1, despite coordinates being absent from the features.The analysis found that the model learned the relevant notion of neighbors from other features.Source: §A.1.3 · p. 14
  • Facebook comments data contained leaked near-copies; Step-2 assigned them about 35% of attention mass, versus 0.09% for Step-1.The split made the task dramatically simpler for TabR, but not for other models.Source: §A.1.3 · p. 14

A.2 ANALYZING THE VALUE MODULE OF TABR

  • Ignoring the value module’s correction along the label-embedding space breaks TabR-S on regression, while ignoring a random one-dimensional projection has little effect.For classification, the tested hypotheses yielded no interesting results; the authors suggest the value module may need different design for those tasks.Source: §A.2 · pp. 14–15, Table 8 · p. 15

A.3 ABLATION STUDY

  • Step-2’s configuration—using labels, omitting queries, and using L2 instead of dot product—cannot be trivially simplified without metric loss on CH, CA, BL, and WE.Omitting queries is equivalent to sharing WQ and WK.Source: §A.3 · p. 15, Table 9 · p. 15

A.4 COMPARING TRAINING TIMES

TabR offers a better performance–training-time trade-off than prior retrieval-based tabular models, though it is slower than simple models. The supplementary material also reports an unsuccessful XGBoost retrieval augmentation, extended context-freeze results, and technical efficiency properties.

A.4 COMPARING TRAINING TIMES

  • TabR has a significantly better performance–training-time trade-off than prior retrieval-augmented tabular models, including NPT, but remains relatively slower than GBDT and parametric DL models.The authors report that training times are affordable within the considered dataset-size scope, despite using a naive implementation lacking basic optimizations.Source: §A.4 · pp. 15–16

A.5 AUGMENTING XGBOOST WITH A RETRIEVAL COMPONENT

  • Adding averaged features and labels from 96 nearest training objects to XGBoost produced no noticeable profit, and varying the neighbor count brought no significant improvements.The neighbor search used the original feature space.Source: §A.5 · p. 16

A.6 ADDITIONAL RESULTS FOR THE “CONTEXT FREEZE” TECHNIQUE

  • The supplementary material extends the context-freeze results with Figure 7 and Tables 12–13, including training-time fractions relative to training without context freeze.The formal definition of the Δ-context metric is given in subsection D.2.Source: §A.6 · p. 17, Figure 7 · p. 17, Table 12 · p. 17

A.7 ADDITIONAL TECHNICAL NOTES ON TABR

  • TabR’s candidate value representations depend on the target and must be computed at runtime, roughly halving deployment memory; its retrieval attention avoids quadratic candidate complexity by attending only to the target.Only key representations and labels for training objects need to be deployed.Source: §A.7 · p. 17

B LIMITATIONS & PRACTICAL CONSIDERATIONS

Retrieval-based predictions require application-specific decisions about which training objects may be used and whether training retrieval matches deployment. Retrieval also adds overhead that may limit TabR on truly large datasets as-is.

B LIMITATIONS & PRACTICAL CONSIDERATIONS

C.1 THE DEFAULT BENCHMARK

The appendices describe the default benchmark's datasets, the protocol used to evaluate the Grinsztajn et al. benchmark, and the hardware used for experiments.

C.1 THE DEFAULT BENCHMARK

  • Table 14 summarizes the default benchmark’s datasets, feature counts, and default batch sizes for training deep-learning models.Source: §C.1 · p. 18, Table 14 · p. 19

C.2 THE BENCHMARK FROM GRINSZTAJN ET AL. (2022)

  • For the Grinsztajn et al. benchmark, the authors retain its splits, exclude the leaky Eye movements dataset, and tune GBDT models with an extensive 200-iteration budget.For datasets with multiple splits, they tune and evaluate each separately, then average metrics over splits.Source: §C.2 · p. 19

D.1 HARDWARE

  • Most GPU experiments used one Nvidia A100; the remaining small portion used one Nvidia 2080 Ti, and some runs used CPUs.All LightGBM experiments were among the CPU runs.Source: §D.1 · p. 19

D.2 IMPLEMENTATION DETAILS OF SUBSECTION 5.1

This section documents implementation and evaluation protocols for TabR and its baselines, including preprocessing, feature embeddings, dataset restrictions, attention analysis, and SAINT context handling. It also describes the expanded tuning used for LightGBM and CatBoost comparisons and CatBoost’s fixed evaluation parameters.

D.2 IMPLEMENTATION DETAILS OF SUBSECTION 5.1

  • Δ-context tracks epoch-to-epoch attention updates as the sum of novel context attention mass and nonnegative increases in attention mass on shared context objects.The authors collect training-object context distributions after each epoch; the formal definition uses contexts of 96 objects.Source: §D.2 · pp. 19–20

D.4 IMPLEMENTATION DETAILS OF SUBSUBSECTION A.1.2

  • For entropy analysis, the authors average test-object candidate distributions after top-m selection and reassign self-attention mass to a shared virtual context object.This avoids misleadingly high entropy when objects attend uniformly or exclusively to themselves.Source: §D.4 · p. 20

D.6 EXPERIMENT SETUP

  • DL models share dataset preprocessing, use AdamW without learning-rate schedules or augmentation, and stop after 17 consecutive epochs without validation improvement.Numerical features generally receive quantile normalization, categorical features one-hot encoding, and binary features mapping to {0, 1}; classification uses cross-entropy and regression mean squared error. Algorithms are counted among the best when their mean score is within one preliminary-best standard deviation.Source: §D.6 · pp. 20–21

D.7 EMBEDDINGS FOR NUMERICAL FEATURES

  • The study uses LR and PLR numerical embeddings and introduces PLR(lite), which shares its linear layer across features for lower weight without critical performance loss.Each scalar numerical feature maps to a vector, and the resulting embeddings are concatenated; all numerical features use the same embedding architecture.Source: §D.7 · p. 21

D.8 TABR

  • TabR-S defaults average tuned hyperparameters across 13 datasets; leave-one-out defaults showed no meaningful changes, and the common configuration performed decently beyond those datasets.TabR-S was used with default hyperparameters for subsections 5.1 and 5.2 and for Table 8.Source: §D.2 · p. 19, §D.3 · p. 20, §D.5 · p. 20

D.9 MLP; D.10 FT-TRANSFORMER; D.11 KNN

  • The MLP and FT-Transformer use established implementations with tuned hyperparameters, while KNN uses the same preprocessing as DL models and tunes its neighbor count from 1 to 128.For FT-Transformer, the implementation uses rtdl version 0.0.13 and its make_baseline method.Source: §D.9 · p. 22, §D.10 · p. 23, §D.11 · p. 23

D.12 DNNR; D.13 DKL; D.14 ANP

  • DNNR, DKL, and ANP are evaluated only on regression; scaling issues exclude DKL on WE and MI and ANP on MI, while DNNR’s efficiency-improved implementation leaves its underlying method unchanged.ANP tuning on WE took more than four days for 20 iterations; DKL tuning alone takes 1 day and 17 hours, compared with 3 hours for TabR on the medium DI dataset.Source: §D.12 · p. 24, §D.13 · p. 24, §D.14 · p. 25

D.15 NPT; D.16 SAINT; D.17 XGBOOST; D.18 LIGHTGBM

  • SAINT is modified so each sample attends only to itself and training-set context, while XGBoost receives more extensive tuning than in prior work to strengthen comparisons.NPT uses its official implementation with model and training code unchanged; LightGBM implementation details are deferred to source code.Source: §D.15 · p. 25, §D.16 · p. 26, §D.17 · p. 27

D.18 LIGHTGBM

  • LightGBM tuning doubled iterations and trees, increased maximum depth, and quadrupled early-stopping rounds versus prior work (Gorishniy et al., 2021; 2022), using TPE.The tuning space is listed in Table 23.Source: §D.18 · p. 27, Table 23 · p. 28

D.19 CATBOOST

  • CatBoost tuning doubled iterations and trees, increased maximum depth, and quadrupled early-stopping rounds versus prior work (Gorishniy et al., 2021; 2022), using TPE.The tuning space is listed in Table 24.Source: §D.19 · p. 28, Table 24 · p. 28
  • CatBoost fixed n_estimators = 4000, early_stopping_rounds = 200, and od_pval = 0.001 rather than tuning them.Source: §D.19 · p. 28

E EXTENDED RESULTS WITH STANDARD DEVIATIONS

The appendix provides extended results with standard deviations for both the default benchmark and the Grinsztajn et al. benchmark; the reported metrics include RMSE and accuracy.

E EXTENDED RESULTS WITH STANDARD DEVIATIONS

Related papers

Loading 2307.14338v2…