Paper: TabR: Tabular Deep Learning Meets Nearest Neighbors (ICLR-2024)

Reviewer comment (Theory):
The motivations behind the module designs are not entirely clear. It appears that the authors made meticulous module (equation) optimization based on its performance on some datasets empirically.

Author response:
Providing clear motivation is of high importance to us. Our reply consists of two parts.

**First**, *based on the content from the submission*, let us show how **the design steps behind TabR (Section 3.2) are motivated formally, informally and by prior work:**

- **Step-0 (baseline) is fully motivated by related work.**
    - (Section 3.2, Paragraph "Step-0") *"the self-attention operation was often used in prior work ... Then, instantiating retrieval module R as the vanilla self-attention ... is a reasonable baseline"*
- **Step-1 (using labels) is highly intuitive:** labels is an extremely valuable signal, so we try using them to improve the model.
    - (Section 3.2, Paragraph "Step-1") *"A natural attempt to improve the Step-0 configuration is to utilize labels of the context objects"*
- **Step-2 (modifying the similarity module $\mathcal{S})$ has a whole range of motivations** provided in the main text ("Step-1" and "Step-2" paragraphs of Section 3.2) and appendix (Section A.1 referenced from Section 5.3, Section A.3 referenced from Section 3.2):
    - **Intuitively,** it is motivated by the outcome of Step-1:  (Section 3.2, Paragraph "Step-1") *"Perhaps, the similarity module S taken from the vanilla attention does not allow benefiting from such a valuable signal as labels."*
    - Section A.1.1 provides **formal motivation** for removing the query representations.
    - Section A.1.1 provides **informal motivation** for using L2 instead of the dot product.
    - **Empirically**, it is motivated by the exhaustive (literally) ablation study in Section A.3, which is references from the main text.
- **Step-3 (modifying the value module $\mathcal{V})$**
    - **Formal motivation:** using $\tilde{x}$ in the value module makes it strictly more expressive: (Section 3.2 Paragraph "Step-3"): *"we make the value module V more expressive by taking the target object’s representation $\tilde{x}$ into account"*
    - **Related work:** (Section 3.2 Paragraph "Step-3") *"we take inspiration from DNNR -- the recently proposed generalization of the kNN algorithm"*
    - **Intuitive interpretation:** (Section 3.2 Paragraph "Step-3"): *"\<the decomposition into the "raw" and "correction" terms\>"*
    - **Quantifying the intuitive interpretation:** Section A.2 referenced from Section 5.3
- **Step-4** is the only purely technical step. However, it is not uncommon for a new architecture to require different technical defaults.

**Second,** we summarize how we further improve the communication and make the available content more visible in the new PDF:
- In Section 3.2 Paragraph "Step-2", we added a reference to Section A.1.
- In Section 3.2 Paragraph "Step-3", we added a reference to Section A.2.
- In Section A.2.1, we added extended comments on Section 3.2 Paragraph "Step-3".
- In Section A.8, we added a minor technical note on Section 3.2 Paragraph "Step-3".

Reviewer comment (Theory):
Why does employing the L2 distance, instead of the dot product, lead to improved performance (as shown in Eq. 3)?

Author response:
**The changes introduced in Equation 3 are discussed in the paper** in Section 3.2 Paragraph "Step-2" and Section A.1 referenced from Section 5.3 (and, in the new PDF, from Section 3.2):

- (Section 3.2 Paragraph "Step-2") Using L2 alone is not enough: *"removing any of the three ingredients (context labels, key-only representation, L2 distance) results in a performance drop back to the level of MLP"*.
- (Section A.1.1) Because of specific details of TabR's encoder and the nature of tabular data, L2 is a more reasonable default choice for TabR.
- (Section A.1.2) Using L2 leads to more diverse neighbor patterns, which, without any additional assumptions, may be a better default behavior for retrieval-based models.
- (Section A.1.3) For several datasets, we empirically observed that TabR with L2 was capable of uncovering much better neighbors than with the dot product.

Also, in Section 3.2 Paragraph "Step-2", we add that L2 is just a reasonable default choice: *"While the L2 distance is unlikely to be the universally best choice for problems (even within the tabular domain), it seems to be a reasonable default choice for tabular data problems."*

Reviewer comment (Theory):
I suggest that providing a theoretical analysis or intuitive motivation would enhance the reader's understanding of those module design details.

Author response:
Providing clear motivation is of high importance to us. Our reply consists of two parts.

**First**, *based on the content from the submission*, let us show how **the design steps behind TabR (Section 3.2) are motivated formally, informally and by prior work:**

- **Step-0 (baseline) is fully motivated by related work.**
    - (Section 3.2, Paragraph "Step-0") *"the self-attention operation was often used in prior work ... Then, instantiating retrieval module R as the vanilla self-attention ... is a reasonable baseline"*
- **Step-1 (using labels) is highly intuitive:** labels is an extremely valuable signal, so we try using them to improve the model.
    - (Section 3.2, Paragraph "Step-1") *"A natural attempt to improve the Step-0 configuration is to utilize labels of the context objects"*
- **Step-2 (modifying the similarity module $\mathcal{S})$ has a whole range of motivations** provided in the main text ("Step-1" and "Step-2" paragraphs of Section 3.2) and appendix (Section A.1 referenced from Section 5.3, Section A.3 referenced from Section 3.2):
    - **Intuitively,** it is motivated by the outcome of Step-1:  (Section 3.2, Paragraph "Step-1") *"Perhaps, the similarity module S taken from the vanilla attention does not allow benefiting from such a valuable signal as labels."*
    - Section A.1.1 provides **formal motivation** for removing the query representations.
    - Section A.1.1 provides **informal motivation** for using L2 instead of the dot product.
    - **Empirically**, it is motivated by the exhaustive (literally) ablation study in Section A.3, which is references from the main text.
- **Step-3 (modifying the value module $\mathcal{V})$**
    - **Formal motivation:** using $\tilde{x}$ in the value module makes it strictly more expressive: (Section 3.2 Paragraph "Step-3"): *"we make the value module V more expressive by taking the target object’s representation $\tilde{x}$ into account"*
    - **Related work:** (Section 3.2 Paragraph "Step-3") *"we take inspiration from DNNR -- the recently proposed generalization of the kNN algorithm"*
    - **Intuitive interpretation:** (Section 3.2 Paragraph "Step-3"): *"\<the decomposition into the "raw" and "correction" terms\>"*
    - **Quantifying the intuitive interpretation:** Section A.2 referenced from Section 5.3
- **Step-4** is the only purely technical step. However, it is not uncommon for a new architecture to require different technical defaults.

**Second,** we summarize how we further improve the communication and make the available content more visible in the new PDF:
- In Section 3.2 Paragraph "Step-2", we added a reference to Section A.1.
- In Section 3.2 Paragraph "Step-3", we added a reference to Section A.2.
- In Section A.2.1, we added extended comments on Section 3.2 Paragraph "Step-3".
- In Section A.8, we added a minor technical note on Section 3.2 Paragraph "Step-3".

Reviewer comment (Experiments):
Some sota DL approaches are not compared, such as T2G-Former (an improved version of FTT), TabPFN, and TANGOS. Especially, TabPFN is relatively similar to TabR. These papers are current SOTA and may outperform tree-based models.

Author response:
**First**, using the benchmark from `[1]` (43 classification and regression tasks; used in Figure 1 and Table 5 in our submission) we show that:
1. **Our model TabR significantly outperforms all the suggested baselines**.
2. **The baselines used in our paper are representative**.

The results:

| Model `[citation]`     | Average rank (+- std) | Wins/Ties/Losses vs TabR | Not applicable |
| ---------------------- | --------------------- | ------------------------ | -------------- |
| **Used in our paper:** |                       |                          |               |
| TabR `[ours]`          | 1.36 (± 0.65)        | 0/43/0                   | 0              |
| XGBoost `[11]`         | 1.93 (± 0.94)        | 7/13/23                     | 0              |
| MLP-PLR `[2]`          | 1.95 (± 0.92)        | 5/19/19                  | 0              |
| FT-Transformer `[4]`   | 2.45 (± 1.20)        | 2/14/27                  | 0              |
| MLP                    | 3.07 (± 1.55)        | 2/12/29                  | 0              |
| **Additional:**        |                       |                          |               |
| T2G-Former `[10]`      | 2.76 (± 1.56)        |  3/14/23                | 3           |
| MLP + TANGOS `[9]`     | 3.29 (± 1.69)        | 2/13/28                | 0           |
| TabPFN `[8]`           | 4.38 (± 1.98)        | 0/4/5                | 34           |

Technical details:
- "Not applicable" (the righmost column) means that the method is fundamentally not applicable to a given task or it does not fit into A100-80GB or (the case of T2G-Former) the hyperparameter tuning did not finish on time.
- If a method is not applicable to a task, it is assigned the largest rank.

**Second**, a closer look at the suggested baselines reveals things that are not consistent with being SOTA *in the scope and niche of our paper*:
- TANGOS
    - **Different niche:** TANGOS is not a model, it is a regularization method, which is completely orthogonal to our niche (model).
    - **Different scope**: the original work mostly focuses on small data with only two datasets having more than 1000 objects.
    - **Significant limitations:** the method seems to make training order(s) of magnitude slower (40x times slower on average on our datasets).
- TabPFN
    - **Different scope:** the work fully focuses on small (<1000 objects) classification tasks.
    - **Significant limitations:** not applicable to regression tasks, not applicable to more than 10 classes and more than 100 features, high memory consumption.
- T2G-Former
    - **Already tested** in its paper on many of the datasets used in our paper and in `[2]` (where embeddings for continuous features and, in particular, MLP-PLR were introduced). The results from `[10]` and `[2]` indicate that models from `[2]` perform at least on par (and often better) while being more efficient. And, the above table confirms that MLP-PLR is indeed a good baseline, while T2G-Former did not generalize well to the new benchmark.

**Third**, in the new PDF, in Section 2 Paragraph "Tabular deep learning", we improved the communication of the niche of our work and, in particular, cited TANGOS and TabPFN.

Reviewer comment (Experiments):
The major comparison lies among middle-scale datasets, accompanied with some results on few other datasets shown in Table 3. In scenarios involving sparse, medium, and dense data-distributed datasets (which typically occur in small, medium-sized, and large-sized datasets, respectively), there may exist a variance in the nearest neighbor retrieval pattern. Hence, conducting tests solely on medium-sized datasets may not suffice. Furthermore, the issue of inefficiency when dealing with large-scaled datasets appears to have hindered the authors from proving the method's effectiveness in large-scaled datasets.

Author response:
Our understanding is that the above quoted comments express uncertainty about our benchmark and whether it is comprehensive enough to support our story.

We did our best to build a **comprehensive, challenging and representative benchmark** to support our claims. Our reply consists of several points.

**First**, let us show that **our experimental setup is competitive** by comparing it to other works on tabular DL:

| Model `[citation]`   | The number of datasets used in the paper                                | Max. train set size               |
| -------------------- | ----------------------------------------------------------------------- |:------------------------------- |
| **TabR** `[ours]`    | **~50 datasets (roughly, `[1]` + some datasets from `[2]` and `[3]`)**  | **2.9M** (+ 10M in the rebuttal) |
| NODE `[7]`           | 6 datasets                                                              | 10M                            |
| TabNet  `[5]`        | <10 datasets                                                            | 10M                            |
| FT-Transformer `[4]` | 11 datasets                                                             | 723K                            |
| SAINT `[6]`          | 30 datasets                                                             | 321K                            |
| \<survey\> `[12]`    | 5 datasets                                                              | 10M                            |
| MLP-PLR `[2]`        | 11 datasets                                                             | 723K                            |
| TabPFN `[8]`         | 67 small datasets (different scope)                                     | 1K                            |
| \<benchmark\> `[1]`  | <50 *unique* datasets (>50 with multiple versions of the same datasets) | 50K                            |
| TANGOS `[9]`         | 20 datasets                                                             | <100K                            |
| T2G-FORMER `[10]`    | 12 datasets                                                             | 320K                            |

**Second,** to illustrate that **our study is based on many diverse datasets**, we provide the structure of our benchmark:
- The "default" benchmark based on prior work (mainly `[2]`):
    - 11 datasets (binary classification, multiclass classification and regression)
    - from 6.4K to 723K training objects
    - from 8 to 136 features
- The widely known and cited benchmark from `[1]`:
    - 43 datasets (binary classification and regression)
    - from 1.7K to 50K training objects
    - from 3 to 613 features
- One weather prediction regression task from `[3]`: 2.9M training objects and 119 features.

**Third, crucially**, to make the benchmark challenging for TabR, we ensured that **there are many challenging tasks that favour GBDT over prior neural networks.** In particular:
- On the benchmark from `[1]`, GBDT was shown to be superior to neural networks.
- `[2]` also contains multiple GBDT-friendly tasks.
- The big weather prediction dataset `[3]` favors GBDT over simple neural networks.

**To sum up**, we truly hope that the above points demonstrate why we consider our benchmark to be of reasonable size, diversity and complexity to support the claims that we make in the abstract and in the list of contribution in Section 1.

Reviewer comment (Experiments):
The method proposed by the authors appears to have achieved slight performance advantages on certain datasets (although some SOTA are not compared). However, due to the lack of explanation for the model details that are designed empirically, it seems unnecessary and risky to apply this method in real-world scenarios. For example, it's unclear whether L2 distance may fail when uninformative features are present; or when a table has a feature with values [f_1, f_2, f_3, ..., f_n], and these values are transformed (e.g., logarithm [log f_1, log f_2, log f_3, ..., log f_n] or reciprocals), the method may perform poorly.

Author response:
Below, we provide results indicating that **the presense of L2 does not make TabR any worse than alternatives**, (dot-product-based TabR or simple MLP) in the presence of uninformative features or after the logarithmic transformation.
Moreover, the alternatives continue lagging significantly behind TabR.

Technical setup:
- Datasets (see Table 1 in the paper): CA, HO, AD, DI, HI.
- We use TabR with the default hyperparameters.
- The uniformative (noisy) features are added as new features.
- For the logarithmic transformation, original features are replaced with the transformed ones.

The table below reports the average (over datasets) difference in `%` compared to the metric of the corresponding model trained on the original unmodified datasets.
 
| Ratio of new uninformative features: | 0.1    | 0.25   | 0.5    |
|:------------------------------------ |:------ |:------ |:------ |
| TabR                                 | -0.6%  | -1.9%  | -2.87% |
| TabR (dot product)                   | -1.27% | -2.26% | -3.77% |
| MLP                                  | -1.32% | -2.63% | -4.37% |

| Ratio of log-transformed features | 0.1    | 0.25   | 0.5    |
|:--------------------------------- |:------ |:------ |:------ |
| TabR                              | -0.15% | -0.77% | -1.07% |
| TabR (dot product)                | -0.23% | -0.18% | -0.39% |
| MLP                               | -0.05% | -0.55% | -0.45% |

**Importantly**, in both of the above regimes, TabR continues to *significantly* outperform alternatives.
For example, this is how the second of the above tables would look like if we reported the differences relative to the *same* baseline, namely, to TabR trained on unmodified original data:

| Ratio of log-transformed features | 0.1    | 0.25   | 0.5    |
|:--------------------------------- |:------ |:------ |:------ |
| TabR                              | -0.15% | -0.77% | -1.07% |
| TabR (dot product)                | -6.95% | -6.95% | -7.17% |
| MLP                               | -6.65% | -7.23% | -7.15% |

Note that, currently, TabR uses a simple MLP-like encoder, and naturally inherits all properties related to robustness to various data challenges  such as uniformative features, irregularly distributed features, etc. `[1]`, `[13]`, `[14]`. In particular:
- These properties explain why, in the above table, there is a performance loss for all DL models regardless of the usage of L2.
- It is also known how to avoid/alleviate these problems: by using transformer-like encoders `[1]`, `[13]` or by using embeddings for continuous features `[2]`.

Reviewer comment (Experiments):
What's the reason for choosing m to be 96? How does m affect the performance of TabR?

Author response:
> What's the reason for choosing m to be 96? How does m affect the performance of TabR?

**In the new PDF, this is addressed in Section A.6.** Below, we provide a summary for convenience and an explanation for the obtained results.

**First,** to evaluate how performance depends on $m$, we consider TabR-S with default hyperparameters and, for each value of $m$, we report the average rank of TabR-S on the default benchmark (Table 1 of the main text). For each `(dataset, m)` pair, the performance is computed as the average over five random seeds.

The choice of $m$  must be made based on the *validation* metrics (not on the *test* metrics), so we provide the ranks on the *validation* sets first:

|     |  m=1 | m=2  | m=4  | m=8  | m=16 | m=32 | m=64 | m=96  | m=128 | m=256 |
|-----|------|------|------|------|------|------|------|-------|-------| ------|
| avg | 5    | 4.5  | 4    | 3.25 | 2.75 | 2.12 | 2.12 | 1.88  | 1.88  | 1.88 |
| std | 2.45 | 1.87 | 1.73 | 1.71 | 1.56 | 1.45 | 1.45 | 0.93  | 1.05  | 1.36 |

It seems that the m=96 turns out to be a reasonable default choice. 

Additionally, here are the ranks on the *test* sets:

|     | m=1 | m=2  | m=4  | m=8 | m=16 | m=32 | m=64 | m=96 | m=128 | m=256 |
|-----|------|------|-----|------|------|------|------|-------|-------|-----|
| avg | 4.12 | 3.62 | 3.5 | 3.12 | 2.38 | 2.12 | 2    | 1.62  | 1.75  | 1.75 |
| std | 2.2  | 1.8  | 1.5 | 1.05 | 1.41 | 1.27 | 1.41 | 0.99  | 0.97  | 1.09 |

**Second**, we suggest an explanation for the obtained results. The presence of the softmax function in the retrieval module of TabR gives a hope that the only requirement for $m$ is to be large enough, and softmax will "automatically" choose the optimal value for each sample. That said, in this work, we don't analyse extreme cases like $m = |I_{train}|$ and recommend values like 96 as a starting point.

Reviewer comment (Evaluation):
What's the inference efficiency of TabR and how does it compare with other baselines (e.g., GBDT)?

Author response:
> What's the inference efficiency of TabR and how does it compare with other baselines (e.g., GBDT)?

**In the new PDF, this is addressed in Section A.4.2.** For convenience, here, we provide a summary.

Below, we report the inference throughput of TabR and XGBoost.

(The technical setup:
- XGBoost and TabR-S with tuned hyperparameters as in Table 4 of the main text
- Computation is performed on NVIDIA 2080 Ti.
- For both models, objects are passed by batches of 4096 objects.)

**The key observations:**
- On the considered datasets, the throughputs of TabR and XGBoost are mostly comparable.
- **Important**: our implementation of TabR is naive and lacks even basic optimizations.

|                                | CH    | CA   | HO   | AD   | DI   | OT   | HI  | BL   | WE   | CO   | MI   |
|--------------------------------|------|------|------|------|------|-----|------|------|------|------|------|
| `#trainingObjects`                | 6400 | 13209 | 14581 | 26048 | 34521 | 39601 | 62751 | 106764 | 296554 | 371847 | 723412 |
| `#features`                      | 11    | 8    | 16   | 14   | 9    | 93  | 28   | 9    |  119 | 54   | 136 |
| `n_estimators` in XGBoost         | 121   | 3997 | 1328 | 988  | 802  | 524 | 1040 | 1751 | 3999 | 1258 | 3814 |
| `max_depth` in XGBoost            | 5     | 9    | 7    | 10   | 13   | 13  | 11   | 8    | 13   | 12   | 12   |
| XGBoost throughput (obj./sec.)  | 2197k | 33k  | 179k | 131k | 417k | 19k | 72k  | 84k  | 15k  | 10k  | 14k  |
| TabR-S throughput (obj./sec.)  | 35k   | 35k  | 55k  | 33k  | 43k  | 40k | 37k  | 27k  | 34k  | 23k  | 11k  |
| Overhead                        | 62.3  | 0.9  | 3.3  | 3.9  | 9.6  | 0.5 | 1.9  | 3.1  | 0.5  | 0.4  | 1.2  |

**In more detail:**
- On "non-simple" tasks (CA, OT, WE, CO), TabR is faster ("non-simple" = XGBoost needs many trees AND/OR XGBoost needs high depth AND/OR a dataset has more features).
- On "simple" tasks, XGBoost is faster ("simple" = XGBoost is shallow AND/OR dataset has few features).
- With the growth of training size (MI), TabR may become slower because of the retrieval, however, there is *a lot* of room for optimizations (caching candidate representations instead of recomputing them on each forward pass; using float16 instead of the current float32; doing approximate search instead of the current brute force; using only a subset of the training data as candidates, etc.)

(P.S. in the paper, we also analyze *training* efficiency in Table 10)

Dear reviewer,

In our "Rebuttal (part 1/2)" reply, we provided some analysis on the inference efficiency.

Now, we are excited to share an **extended study on the inference throughput covering SEVEN models in TWO modes.** We did not update the PDF document with the new content, but we are committed to do that based on feedback.

**The main observations**:
- TabR is faster than prior retrieval-based tabular DL (SAINT).
- TabR is slower than simple retrieval-free models.
- Notably, TabR is only moderately slower than FT-Transformer (a non-retrieval model).
- *(not reflected in the tables below)* On large dataset, where the retrieval becomes the bottleneck, TabR can be made *significantly* faster (e.g. by order of magnitude) by using approximate nearest neighbor search instead of the current brute force.

Technical details:
- We did *not* use DL-specific optimizations like `torch.compile`, mixed precision computation, pruning and other techniques.
- Models with tuned hyperparameters are used (that is, from Table 3 and Table 4). For FT-Transformer, the tuned hyperparameters were taken from `[1]` and `[2]` for most datasets, and the default hyperparemeters were used on two datasets.

**The first mode: "Online predictions"**
- Device: CPU Intel Core i7-7800X 3.50GHz
- Thread count: 1
- Batch size: 1

*Notation*
- *Values are aggregated over the 11 datasets from Table 1 (the "default" benchmark)*
- *(A) ~ Absolute throughput (objects per second)*
- *(R) ~ Relative throughput w.r.t. TabR-S*

| [CPU & Batch size = 1]   | (A) Min   | (A) Median   | (A) Max   | (R) Min   |   (R) Median |   (R) Max |
|------------------------|:-----------|:--------------|:-----------|:-----------|:--------------|:-----------|
| TabR-S                 | 31        | 470          | 2.0K      | 1       |         1    |       1   |
| TabR                   | 26        | 236          | 1.2K      | 0.2       |         0.5  |       1.3 |
| SAINT                  | < 1       | 15           | 59        | < 0.01    |         0.02 |       0.1 |
| XGBoost                | 270       | 2.2K         | 10K       | 0.6       |         5.5  |      17.5 |
| MLP                    | 2.6K      | 4.8K         | 15K       | 4.0       |        13.9  |     387.1 |
| MLP-PLR                | 532       | 1.6K         | 5.1K      | 0.9       |         5.2  |      47.2 |
| FT-Transformer         | 219       | 694          | 1.7K      | 0.4       |         1.6  |      11.6 |

**The second mode: "Offline batch processing"**
- Device: GPU NVIDIA 2080Ti 12GB
- Batch size: the largest possible batch size for a given model (but no larger then `2 ** 17 ~= 128_000`)

It turns out that XGBoost can achieve better numbers in this mode than with the fixed batch size 4096.

*Notation*
- *Values are aggregated over the 11 datasets from Table 1 (the "default" benchmark)*
- *(A) ~ Absolute throughput (objects per second)*
- *(R) ~ Relative throughput w.r.t. TabR-S*

| [CUDA & Batch size = max]   | (A) Min   | (A) Median   | (A) Max   |   (R) Min |   (R) Median |   (R) Max |
|---------------------------|:-----------|:--------------|:-----------|:-----------|:--------------|:-----------|
| TabR-S                    | 24K       | 123K         | 319K      |      1    |          1   |       1   |
| TabR                      | 28K       | 91K          | 308K      |      0.4  |          0.8 |       2.5 |
| SAINT                     | 2.1K      | 32K          | 133K      |      0.05 |          0.3 |       0.8 |
| XGBoost                   | 63K       | 1.4M         | 16.6M     |      0.9  |         11.4 |     126.8 |
| MLP                       | 3.2M      | 7.7M         | 47.8M     |     23.9  |         67.4 |     512   |
| MLP-PLR                   | 63K       | 1.1M         | 8.5M      |      0.5  |          6.7 |     101.7 |
| FT-Transformer            | 21K       | 198K         | 752K      |      0.2  |          2   |       4.5 |

**References**

- `[1]` "Revisiting Deep Learning Models for Tabular Data", Gorishniy et al.
- `[2]` "On Embeddings for Numerical Features in Tabular Deep Learning", Gorishniy et al.

Reviewer comment (Experiments):
Is TabR applicable to categorical features? It seems like the paper only considers continuous features.

Author response:
> It seems like the paper only considers continuous features

**In fact, the used datasets include both continuous and categorical features**:
- For the default benchmark, Table 1 in the main text provides information on how many features of each type is presented in each dataset.
- The benchmark from `[1]` also includes many datasets with categorical features, please, see the original paper for details.

> Is TabR applicable to categorical features?

**Yes, TabR is applicable to all kinds of features:** its input encoder consists of conventional common modules, as explained in the caption of Figure 3. In particular, categorical features are encoded with the one-hot encoding, and, if needed, one can use other encoding/embedding schemes.

**References**

- `[1]` "Why do tree-based models still outperform deep learning on tabular data?", NeurIPS 2022 Datasets and Benchmarks, Grinsztajn et al.

Reviewer comment (Novelty):
Paper doesn't go into detail describing differences with prior deep learning-based tabular methods. What might explain the performance differences? Ex. "prior work, where several layers with multi-head attention between objects and features are often used" but was this what led to retrieval's low benefit in the past?

Author response:
> What might explain the performance differences? ... but was this what led to retrieval's low benefit in the past?

(1) **TL;DR: yes, the main reason behind the difference in the performance lies in the retrieval module designed in Section 3.2.** A quick summary of the main differences in the retrieval-related functionality:
- Prior work: *multiple multi-head vanilla* attention modules, where the attention is performed *between objects and features*.
- TabR: *one single-head custom* attention-like module, where the attention is performed *only between objects*.

In the new PDF, in Section 2 on Page 3, we improved the wording and the formatting to better contrast TabR against prior solutions.

(2) **The relevant parts from the original PDF** (here, we uppercased the key pieces for convenience):

- (Paragraph "Step 2" Page 5, where we improved the similarity module commonly used in prior work) *This change is a turning point in our story, WHICH WAS OVERLOOKED IN PRIOR WORK.*
- (abstract) *"NOVEL FINDINGS ... LIE IN THE ATTENTION-LIKE MECHANISM that is responsible for retrieving the nearest neighbors and extracting valuable signal from them."*
- (Section 2 Page 3) *"Compared to prior work, where SEVERAL LAYERS WITH MULTI-HEAD ATTENTION ... TabR implements its retrieval component with just ONE SINGLE-HEAD ATTENTION-LIKE MODULE. Importantly ... module of TabR is CUSTOMIZED in a way that makes it BETTER SUITED FOR TABULAR DATA PROBLEMS"*
- etc.

Reviewer comment (Evaluation):
Insufficient discussion of categorical variables. Is accuracy or training time particularly affected by their relative abundance relative to numerical features?

Author response:
> Insufficient discussion of categorical variables.

However, as indicated by Figure 3 and its caption, **TabR is equally capable of handling all feature types, there is no conceptual difference between them.** In particular:
- Categorical features are encoded with the one-hot encoding (as mentioned in the caption of Figure 3 and in Section D6). This is a simple and efficient operation, and one can switch to any other encoding scheme if needed.
- Continuous features are normalized and, optionally, transformed with embeddings from `[2]` (which, in fact, makes continuous features more demanding than categorical features in terms of compute).

Just in case, we will mention that **the used datasets include both continuous and categorical features**:
- The default benchmark: Table 1 provides information on how many features of each type is presented in each dataset.
- The benchmark from `[1]` also includes many datasets with categorical features.

> Is accuracy or training time particularly affected by their relative abundance relative to numerical features?

We did not notice any different behaviour caused by the presence of categorical features, which is consistent with the explanation provided above (*"... for TabR, there is no conceptual difference between feature types ..."*).

Reviewer comment (Experiments):
The steps of Section 3.2 seem rather arbitrary. Some of the detail could be compressed to make room for more intuition why the final architecture makes more sense (content from A.1.1). Description of architectural changes that didn't work would also be very insightful.

Author response:
> The steps of Section 3.2 seem rather arbitrary.

**First**, just in case, we provide a quick informal recap of the story behind the design steps in Section 3.2:
- "Step-0. The vanilla-attention-like baseline" is motivated by prior work.
- "Step-1. Adding context labels" is a natural attempt to use the available labels of the neighbors. The important outcome: *"the similarity module $\mathcal{S}$ taken from the vanilla attention does not allow benefiting from such a valuable signal as labels"*.
- "Step-2. Improving the similarity module $\mathcal{S}$" is fully motivated by the quoted outcome of Step-1.
- "Step-3. Improving the value module $\mathcal{V}$" is inspired by a simple observation that $\mathcal{V}$ can be made more expressive by adding the dependency on the target object's representation $\tilde{x}$.
- Step-4 is the only purely technical step where we identified better technical defaults for the new architecture.

**Second**, how we improved the communication and made the available content more visible in the new PDF:
- In Section A.2.1, we added extended motivation behind the design of the value module of TabR.
- In the "Step-2" and "Step-3" paragraphs, we added references to the extended motivation and analysis available in the appendix.
- In the "Step-3" paragraph, we improved the wording to highlight the transition of our focus from the similarity module to the value module.

Reviewer comment (Evaluation):
Paper describes training times in A.4, but I believe a summary of this is important enough to warrant inclusion in the main paper. Something like a mention of the geometric mean (over the datasets) of the ratio between TabR's training time to a gradient boosted methods, described in the conclusion, would be sufficient. While the ratio is likely >1, it is better to acknowledge this weakness than to hide it.

Author response:
> Paper describes training times in A.4, but I believe a summary of this is important enough to warrant inclusion in the main paper. Something like a mention of the geometric mean (over the datasets) of the ratio between TabR's training time to a gradient boosted methods, described in the conclusion, would be sufficient.

We are excited to discuss the aspect of efficiency in more detail!

Regarding the suggestion to announce Table 10 (training times), currently, we achieve that and prepare readers for Table 10 by:

- repeatedly mentioning two things throughout the paper:
    - TabR is more efficient than prior retrieval-based tabular DL.
    - TabR is less efficient than simple models and may require special considerations on larger datasets.
- directly mentioning Section A.4 and Section A.4.1 multiple times.

**Importantly**, we believe that Table 10 (training times) needs a full-fledged presentation and nuanced discussion, and it may suffer from simplifications. In particular, announcing Table 10 with a single number can be misleading: for example, the suggested geometric mean of the ratios will completely hide absolute training times, which is what actually matters in practice. To avoid such effects, we report specific numbers for both "positive" (e.g. order(s) of magnitude improvements over prior work) and "negative" efficiency-related stories only in Table 10.

>  it is better to acknowledge this weakness than to hide it.

Based on the content from the original submission, let us illustrate that, in fact, **we care deeply about efficiency, invest heavily in this aspect of our work, and strive for complete transparency in this matter.** We truly consider the storyline about efficiency an important part of our paper. In particular, throughout the work, we repeatedly mention two things:
- TabR is more efficient than prior retrieval-based tabular DL.
- TabR is less efficient than simple models and may require special considerations on larger datasets.
    - (The end of Introduction) *"Tree-based models, in turn, remain a cheaper solution"*
    - **The whole Section 5.1 is dedicated to improving training times of TabR**.
    - Even the seemingly unrelated Section 5.2 mentions that the continual updates can be used to make TabR train faster: *"Additionally, this approach can be used to scale TabR to large datasets by training the model on a subset of data and retrieving from the full data."*
    - (Conclusion) *"An important direction for future work is improving the efficiency of retrieval-augmented models to make them faster in general and in particular applicable to tens and hundreds of millions of data points"*
    - (Limitations on Page 6 refer to Section B, where we write:) *"the retrieval module R still causes overhead compared to purely parametric models, so TabR may not scale to truly large datasets as-is"*

Reviewer comment (Writing):
What is $I_{cand}$? Is it all rows of the table that labels have been provided for? It's mentioned in page 3 that "we use the same set of candidates for all input objects" but what is the set of candidates exactly?

Author response:
> Also, what is $I_{cand}$? Is it all rows of the table that labels have been provided for?

1. **In most of the experiments, yes**: the neighbors are retrieved from all training objects, which formally means $I_{cand} = I_{train}$.
2. Strictly speaking, we implicitly rely on the fact that  $I_{cand}$ can be different from $I_{train}$ in the following places:
    - In Section 5.1, for *training* objects,  $I_{cand} = I_{train}$ is true only until the context freeze, after which the retrieval for *training* objects is not performed (instead, the same neighbors are reused until the end of training).
    - In Section 5.2, for *test* objects, $I_{cand} = I_{train}$ is true only until new labeled objects are added to the set of candidates. 
    - The *"Second, depending on an application ..."* paragraph of Section B.

> It's mentioned in page 3 that "we use the same set of candidates for all input objects"

In the new PDF, we improved the wording, now it says: *"In this work, unless otherwise noted, we use the same set of candidates for all input objects and set $I_{cand} = I_{train}$ (which means retrieving from all training objects)"*.

Reviewer comment (Evaluation):
A comparison of the inference and query complexity between the methods is lacking.

Author response:
We thank the reviewer for the positive feedback!

> A comparison of the inference and query complexity between the methods is lacking.

**In the new PDF, this is addressed in Section A.4.2.** For convenience, here, we provide a summary.

Below, we report the inference throughput of TabR and XGBoost.

(The technical setup:
- XGBoost and TabR-S with tuned hyperparameters as in Table 4 of the main text
- Computation is performed on NVIDIA 2080 Ti.
- For both models, objects are passed by batches of 4096 objects.)

**The key observations:**
- On the considered datasets, the throughputs of TabR and XGBoost are mostly comparable.
- **Important**: our implementation of TabR is naive and lacks even basic optimizations.

|                                | CH    | CA   | HO   | AD   | DI   | OT   | HI  | BL   | WE   | CO   | MI   |
|--------------------------------|------|------|------|------|------|-----|------|------|------|------|------|
| `#trainingObjects`                | 6400 | 13209 | 14581 | 26048 | 34521 | 39601 | 62751 | 106764 | 296554 | 371847 | 723412 |
| `#features`                      | 11    | 8    | 16   | 14   | 9    | 93  | 28   | 9    |  119 | 54   | 136 |
| `n_estimators` in XGBoost         | 121   | 3997 | 1328 | 988  | 802  | 524 | 1040 | 1751 | 3999 | 1258 | 3814 |
| `max_depth` in XGBoost            | 5     | 9    | 7    | 10   | 13   | 13  | 11   | 8    | 13   | 12   | 12   |
| XGBoost throughput (obj./sec.)  | 2197k | 33k  | 179k | 131k | 417k | 19k | 72k  | 84k  | 15k  | 10k  | 14k  |
| TabR-S throughput (obj./sec.)  | 35k   | 35k  | 55k  | 33k  | 43k  | 40k | 37k  | 27k  | 34k  | 23k  | 11k  |
| Overhead                        | 62.3  | 0.9  | 3.3  | 3.9  | 9.6  | 0.5 | 1.9  | 3.1  | 0.5  | 0.4  | 1.2  |

**In more detail:**
- On "non-simple" tasks (CA, OT, WE, CO), TabR is faster ("non-simple" = XGBoost needs many trees AND/OR XGBoost needs high depth AND/OR a dataset has more features).
- On "simple" tasks, XGBoost is faster ("simple" = XGBoost is shallow AND/OR dataset has few features).
- With the growth of training size (MI), TabR may become slower because of the retrieval, however, there is *a lot* of room for optimizations (caching candidate representations instead of recomputing them on each forward pass; using float16 instead of the current float32; doing approximate search instead of the current brute force; using only a subset of the training data as candidates, etc.)

Reviewer comment (Evaluation):
Inference time and complexity -- are the studies based on normalized inference time between models? If not, could you comment more? How does the inference complexity depend on the size of the table data?

Author response:
> Inference time and compexity -- are the studies based on normalized inference time between models? If not, could you comment more?

Strictly speaking, the results reported in all main tables are obtained without fixing the inference time budget: all methods can spend any time they need to make predictions. *However,* generally, the efficiency of TabR vs. prior work is an important storyline of our paper with significant wins over prior work, and, in the considered scope of datasets, the inference time of TabR is always reasonable (as indicated by the above table).

> How does the inference complexity depend on the size of the table data?

As indicated by the above table, for datasets of sizes up to (roughly) 500K objects, for TabR-S, the dataset size is not a major factor defining the inference throughput. For larger datasets, the dataset size becomes a more important factor (formally, for the current brute force search, the search complexity grows linearly with the dataset size). Luckily, the current implementation of TabR has a lot of room for simple optimizations:
- caching candidate key representations instead of recomputing them on each forward pass.
- performing the search in float16 instead of the current float32.
- performing approximate similarity search instead of the current brute force.
- using only a subset of the training data as candidates.
- etc.

Reviewer comment (Experiments):
Could a different selection of datasets prove that the tabR is not superior to GBDT? In other words, are these datasets highly representative?

Author response:
> are these datasets highly representative?

We did our best to build a representative benchmark using **~50 datasets from prior work**.
In particular, we ensured that there are **many challenging tasks that favour GBDT over traditional neural networks.** Examples:
- We use the widely cited `[1]` where GBDT was shown to be superior to neural networks.
- We use some datasets from `[2]` including multiple GBDT-friendly tasks.
- We also use the real world large scale dataset for weather prediction from `[3]`.

To put our work in the context of prior work, here are the main properties of the benchmarks used in some prior work on tabular deep learning:

| Model `[citation]`   | The number of datasets used in the paper                                |
| -------------------- | ----------------------------------------------------------------------- |
| **TabR** `[ours]`    | **~50 datasets (roughly, `[1]` + some datasets from `[2]` and `[3]`)**  |
| NODE `[7]`           | 6 datasets                                                              |
| TabNet  `[5]`        | <10 datasets                                                            |
| FT-Transformer `[4]` | 11 datasets                                                             |
| SAINT `[6]`          | 30 datasets                                                             |
| \<survey\> `[11]`    | 5 datasets                                                              |
| MLP-PLR `[2]`        | 11 datasets                                                             |
| TabPFN `[8]`         | 67 small datasets (different scope)                                     |
| \<benchmark\> `[1]`  | <50 *unique* datasets (>50 with multiple versions of the same datasets) |
| TANGOS `[9]`         | 20 datasets                                                             |
| T2G-FORMER `[10]`    | 12 datasets                                                             |

Reviewer comment (Evaluation):
Is it not surprising that Step-1 (adding context labels) did not help that much? One would guess that this is a big component of signal in retrieval augmentation.

Author response:
> Is it not surprising that Step-1 (adding context labels) did not help that much?

This is discussed in the "Step-1" paragraph on Page 5:

*> "Table 2 shows no improvements from using labels, which is counter-intuitive. Perhaps, the similarity module taken from the vanilla attention does not allow benefiting from such a valuable signal as labels."*

In other words, yes, this is indeed surprising! However, after the subsequent Step-2 in Section 3.2 (and, in particular, after the ablation study in Section A.3), we can confidently say that *to benefit from a valuable signal (such as labels), good similarity module is required.*