[论文解读] TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023

Yury Gorishniy, Ivan Rubachev+4 more|arXiv (Cornell University)|Jul 26, 2023

Machine Learning and Data Classification被引用 7

一句话总结

TabR 提出了一种简单且高效的检索增强深度学习模型，用于表格数据，其在前馈网络中集成了 k-最近邻（k-NN）机制。该模型在表格基准测试中取得了最先进性能，在近期提出的‘GBDT 友好’基准上，优于先前的深度学习模型和梯度提升决策树（GBDT），同时相比以往的检索增强方法显著提升了效率。

ABSTRACT

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.

研究动机与目标

• 解决中等规模表格数据集中，表格深度学习模型与梯度提升决策树（GBDT）之间长期存在的性能差距。
• 通过在标准前馈架构中引入一种新颖且轻量级的 k-NN 机制，提升检索增强表格深度学习的效率与效果。
• 证明深度学习模型可以在表格数据上超越 GBDT，尤其是在为偏好基于树的模型而设计的近期基准上。
• 识别并利用注意力机制类组件中的关键设计选择，以增强表格设置下的检索性能。
• 提供一种简单、高性能且计算高效的替代方案，以取代复杂且计算开销大的检索增强表格深度学习模型。

提出的方法

• TabR 采用标准多层感知机（MLP）主干网络，并在中间插入一个自定义的 k-NN 类组件，取代典型的注意力机制。
• 检索组件使用可学习的、类似注意力的机制，计算输入样本与所有训练样本之间的相似度得分，随后检索最相似的 k 个邻居。
• 将检索到的邻居的特征和标签进行聚合，并与输入嵌入拼接，以在最终预测前增强表示。
• 模型通过标准反向传播进行端到端训练，检索机制通过邻居的软分配实现可微分。
• 方法对连续特征使用学习嵌入，对正则化应用标准归一化和 dropout。
• 该架构设计为轻量化且高效，避免了完整注意力或内存密集型检索模块的计算开销。

实验结果

研究问题

• RQ1 检索增强的深度学习模型是否能在中等规模表格基准上超越 GBDT，尤其是在为偏好基于树的模型而设计的基准上？
• RQ2 在检索机制中，特别是注意力类组件的设计选择，哪些最显著地影响了表格深度学习中的性能与效率？
• RQ3 将 k-NN 检索集成到简单前馈网络中，与更复杂的检索增强架构相比，在准确率和推理成本方面有何差异？
• RQ4 一个简单且可微分的 k-NN 机制在多大程度上能提升表格数据任务中的泛化能力和鲁棒性？
• RQ5 所提出的模型是否在多样化的表格数据集上实现了最先进性能，同时保持计算效率？

主要发现

• 在包含 43 个中等规模表格任务的基准上，TabR 在所有表格深度学习模型中实现了最佳平均性能。
• 在多个独立数据集上，包括 'adult'、'german'、'proteins' 和 'sulfur'，TabR 创造了新的最先进记录。
• 在近期提出的 'GBDT-friendly' 基准（Grinsztajn 等，2022）上，TabR 在单模型和集成设置中均优于 XGBoost、LightGBM 和 CatBoost。
• 在 'adult' 数据集上，TabR 的 AUC 达到 0.871（单模型）和 0.876（集成），甚至超过了调优后的 XGBoost 和 CatBoost。
• 该模型相比以往的检索增强表格模型显著更高效，其精简架构避免了昂贵的注意力或内存模块。
• 消融研究证实，类似注意力的检索机制对性能至关重要，且相似度计算与邻居聚合的合理设计可带来一致的性能提升。