TabM：通过参数高效的集成推进表格深度学习》（ICLR 2025）

:scroll: arXiv   :books: 其他表格DL项目

这是论文《TabM：通过参数高效的集成推进表格深度学习》的官方仓库。它由两部分组成：

• 本文档中介绍的Python软件包。
• paper/README.md 中描述的与论文相关的内容（代码、指标、超参数等）。

TabM在Kaggle上的应用（截至2025年6月）

• TabM被用于UM举办的竞赛中的冠军方案。
• TabM也被用于CIBMTR举办的竞赛中的冠军方案，以及前三名、前四名、前五名和其他许多参赛方案中。后来发现，仅使用TabM而不与其他模型集成，也能在3300多个参赛者中获得第25名567863号讨论。

TabM在TabReD（一项具有挑战性的基准测试）上的表现

TabReD是一个基于真实世界工业数据集的基准测试，这些数据集存在与时间相关的分布漂移和数百个特征，因此比传统基准更具挑战性。下图显示，与先前的表格深度学习方法相比，TabM在TabReD（加上另一个真实世界数据集）上取得了更高的性能。

每个点代表一个数据集上的性能得分。对于给定的模型，菱形表示各数据集上的平均值。

训练与推理效率

TabM是一种简单且效率合理的模型，因此非常适合实际应用，包括处理大规模数据集。论文中使用的最大数据集包含1300万个样本，我们还了解到曾成功地在超过1亿个样本上进行过训练，尽管在这种情况下训练所需的时间会更长。

下图显示，TabM的速度相对MLP和GBDT较慢，但比之前的表格深度学习方法更快。请注意：(1) 推理吞吐量是在单线程CPU上测量的，且未进行任何优化，特别是未采用本文后续介绍的TabM专用加速技术；(2) 左侧图表采用了对数尺度。

每个点代表在一个数据集上的测量结果。对于给定的模型，菱形表示各数据集上的平均值。在左图中，$\mathrm{TabM_{mini}^{\dagger*}}$ 表示使用混合精度和torch.compile训练的$\mathrm{TabM_{mini}^{\dagger}}$。

TL;DR

TabM（即Tabular DL模型，能够进行Multiple预测）是一种简单而强大的表格深度学习架构，可高效模拟MLP集成的效果。与常规MLP集成相比，TabM的两个主要区别在于：

• 并行训练MLP。这使得可以在训练过程中监控集成的整体性能，并在对整个集成最优时停止训练，而不是分别针对每个MLP进行训练。
• 权重共享。实际上，整个TabM可以整合为一个类似MLP的单一模型。这不仅显著提升了运行时间和内存效率，还起到了有效的正则化作用，从而提升任务性能。

复现实验与浏览结果

[!IMPORTANT] 若要在实践中及未来工作中使用TabM，请使用下方介绍的tabm软件包。

与论文相关的内容（代码、指标、超参数等）位于paper/目录下，并在paper/README.md中进行了详细说明。

Python软件包

tabm是一个基于PyTorch的Python软件包，提供了TabM模型，以及用于构建自定义TabM-like架构（即高效MLP-like模型集成）的层和工具。

• 安装
• 基本用法
  • 创建TabM
  • 创建带有特征嵌入的TabM
  • 将TabM与自定义输入及输入模块结合使用
  • 训练
  • 推理
• 示例
• 高级用法
  • 直观理解
  • EnsembleView
  • MLP集成
  • 重要的实现细节
  • 示例：无权重共享的简单集成
  • 示例：MiniEnsemble
  • 示例：BatchEnsemble
  • 示例：自定义架构
  • 将现有模型转化为高效集成
• 超参数
  • 默认模型
  • 默认优化器
  • arch_type
  • k
  • num_embeddings
  • 初始化
  • 超参数调优
• 实用提示
  • 推理效率
• API

安装

pip install tabm

基本用法

本节展示了在典型应用场景下如何创建模型，并对训练和推理给出了高层次的说明。

创建TabM

以下示例展示了不带特征嵌入的基本版TabM。为了获得更好的性能，通常应按照下一节的说明传递num_embeddings参数。

[!NOTE] 下文使用的TabM.make(...)会根据提供的参数添加默认超参数。

import torch
from tabm import TabM

# >>> 所有后续章节的通用设置。
d_out = 1  # 例如，一个回归任务。
batch_size = 256

# 数据集包含24个数值型（连续）特征。
n_num_features = 24

# 数据集包含2个分类特征。
# 第一个分类特征有3个唯一类别。
# 第二个分类特征有7个唯一类别。
cat_cardinalities = [3, 7]

# <<<

model = TabM.make(
    n_num_features=n_num_features,
    cat_cardinalities=cat_cardinalities,  # 将使用独热编码。
    d_out=d_out,
)
x_num = torch.randn(batch_size, n_num_features)
x_cat = torch.column_stack([
    # 第i个类别特征的取值范围为[0, cat_cardinalities[i])。
    torch.randint(0, c, (batch_size,)) for c in cat_cardinalities
])
y_pred = model(x_num, x_cat)

# TabM表示由k个模型组成的集成，因此每个样本会有k个预测。
assert y_pred.shape == (batch_size, model.k, d_out)

使用特征嵌入创建TabM

在典型的表格数据任务中，通常通过将特征嵌入模块作为num_embeddings传入来获得最佳性能（在论文中，带有嵌入的TabM被记作$\mathrm{TabM^\dagger}$）。TabM支持来自rtdl_num_embeddings包中的多种特征嵌入模块。下面的例子展示了最简单的嵌入模块LinearReLUEmbeddings。

[!TIP] 通常，更高级的嵌入方法，如PiecewiseLinearEmbeddings和PeriodicEmbeddings，能够带来更好的性能。它们的使用将在端到端示例中介绍（见#examples）。

from rtdl_num_embeddings import LinearReLUEmbeddings

model = TabM.make(
    n_num_features=n_num_features,
    num_embeddings=LinearReLUEmbeddings(n_num_features),
    d_out=d_out
)
x_num = torch.randn(batch_size, n_num_features)
y_pred = model(x_num)

assert y_pred.shape == (batch_size, model.k, d_out)

在自定义输入和输入模块中使用TabM

[!TIP] tabm.TabM的实现是定义基于TabM的模型中输入和输入模块的一个良好示例。

假设您希望改变TabM接收的输入内容或处理方式，但仍想以TabM作为骨干网络。那么典型的用法如下：

from tabm import EnsembleView, make_tabm_backbone, LinearEnsemble


class Model(nn.Module):
    def __init__(self, ...):
        # >>> 创建任何自定义模块。
        ...
        # <<<

        # 创建集成输入模块。
        self.ensemble_view = EnsembleView(...)
        # 创建骨干网络。
        self.backbone = make_tabm_backbone(...)
        # 创建预测头。
        self.output = LinearEnsemble(...)

    def forward(self, arg1, arg2, ...):
        # 根据需要将输入转换为一个张量。
        # 这一步骤可以包括特征嵌入和其他各种特征变换。
        # `handle_input`是一个假想的用户定义函数。
        x = handle_input(arg1, arg2, ...)  # -> (B, D) 或 (B, k, D)

        # 与传统模型唯一的区别在于调用了self.ensemble_view。
        x = self.ensemble_view(x)  # -> (B, k, D)
        x = self.backbone(x)
        x = self.output(x)
        return x  # -> (B, k, d_out)

[!NOTE] 关于x = handle_input(...)这行中x的形状：

• TabM可以用作传统的类似MLP的骨干网络，此时x在训练和推理阶段都具有标准形状(B, D)。由于其简单性和更高的效率，这种做法是默认推荐的。
• 另一种高级训练策略是，在训练时x的形状为(B, k, D)，而在推理时为(B, D)。

端到端示例（见#examples）涵盖了这两种方法。

训练

至关重要的是要独立训练TabM的k个预测结果，而不要对它们进行平均。 换句话说，必须优化平均损失，而不是平均预测结果的损失。端到端示例（见#examples）提供了关于如何训练TabM的完整参考。

推理

在推理阶段，为了得到某个样本的预测结果，需要对k个预测结果进行平均。具体的平均策略取决于任务和损失函数。例如，在分类任务中，通常应该对概率进行平均，而不是对logit进行平均。端到端示例（见#examples）展示了如何使用TabM进行预测。

示例

example.ipynb提供了一个训练TabM的端到端示例：

• 在GitHub上查看
• 在Colab中打开

高级用法

[!TIP] 在构建自定义模型之前，请先尝试并调整TabM。尽管TabM使用的是普通的MLP作为基础模型，但它的简单性可能会让人误以为它不够强大，实际上它是一个非常强的基线。

本部分的内容超出了TabM论文的范围，提供了用于构建自定义TabM类架构的组件，包括：

• MLP、线性层和归一化层的高效集成；
• 其他对高效集成有用的层，如EnsembleView和ElementwiseAffine；
• 将单个模型就地转换为高效集成的函数；
• 以及其他工具。

部分内容将在专门章节中讨论，其余则在下面的示例中展示。

[!IMPORTANT] 理解TabM的实现细节对于构建正确且有效的TabM类模型至关重要。其中一些内容将在本节稍后讨论。其他重要参考资料包括：

• 本包的源代码，特别是TabM模型和make_tabm_backbone函数；
• TabM论文，尤其是第3.3节中的$\mathrm{TabM_{mini}}$段落（arXiv v3版本）。

直觉

回想一下，在传统的类似MLP的模型中，一个典型的模块代表一层作用于形状为(B, D)的张量上，其中B是批次大小，D是潜在表示的维度。相比之下，本包中的典型模块：

• 代表k个层的集合，这些层并行地作用于k个输入上；
• 操作的对象是形状为(B, k, D)的张量，该张量表示k个输入（每个层对应一个输入）。

示例：

• LinearEnsemble是k个独立线性层的集合；
• LinearBatchEnsemble是k个共享大部分权重的线性层的集合。需要注意的是，权重共享并不会改变模块的应用方式：它仍然代表k个层并行作用于k个输入上。

EnsembleView

EnsembleView是一个特殊的轻量级模块，只做一件简单的事情：

• 将形状为(B, D)的张量转换为形状为(B, k, D)的张量，其中存储了原始张量的k个完全相同的视图。这是一个廉价且无需复制的操作。
• 形状为(B, k, D)的张量会原样传递，不做任何更改。

MLP 集成模型

该包提供了以下高效的 MLP 集成模型：

• MLPBackboneBatchEnsemble（由 $\mathrm{TabM}$ 使用）
• MLPBackboneMiniEnsemble（由 $\mathrm{TabM_{mini}}$ 使用）
• MLPBackboneEnsemble（由 $\mathrm{TabM_{packed}}$ 使用）

[!NOTE] 直接创建上述集成模型与使用 tabm.make_tabm_backbone 函数之间的区别在于，make_tabm_backbone 中的某些细节是为 TabM 优化的。这些细节在 TabM 之外也可能有用，但目前尚未深入探讨。

与 TabM 不同，它们仅接受形状为 (batch_size, k, d) 的三维输入。因此，用户需要负责将输入转换为一个张量（例如使用嵌入、独热编码等），其中存储的是同一对象的 k 个视图，或者 k 个完整的批次。一个基本的使用示例：

import tabm
import torch
import torch.nn as nn

d_in = 24
d_out = 1
k = 32
model = nn.Sequential(
    tabm.EnsembleView(k=k),
    tabm.MLPBackboneBatchEnsemble(
        d_in=d_in,
        n_blocks=3,
        d_block=512,
        dropout=0.1,
        k=k,
        tabm_init=True,
        scaling_init='normal',
        start_scaling_init_chunks=None,
    ),
    tabm.LinearEnsemble(512, d_out, k=k)
)
x = torch.randn(batch_size, d_in)
y_pred = model(x)
assert y_pred.shape == (batch_size, k, d_out)

重要实现细节

本节介绍构建自定义 TabM 类模型时需要注意的实现细节。

层的顺序。 最重要的指导原则是：在用线性层混合表格特征之前，应先生成 $k$ 种不同的对象表示。这直接源自论文（arXiv V3）第 3.3 节中关于 $\mathrm{TabM_{mini}}$ 的说明。

d_in = 24
d = 512
k = 16

# 好的：前两个模块共同在第一个线性层之前生成 k 种不同的对象表示。
scaling_init = "normal"  # 或 "random-signs"
good_model = nn.Sequential(
    tabm.EnsembleView(k=k),
    tabm.ElementwiseAffine((k, d_in), bias=False, scaling_init=scaling_init),
    nn.Linear(d_in, d),
    ...
)

# 好的：LinearBatchEnsemble 内部首先进行非共享的逐元素缩放，从而在进行线性变换之前使 k 种对象表示多样化。
scaling_init = "normal"  # 或 "random-signs"
                         # 或 ("normal", "ones")
                         # 或 ("random-signs", "ones")
good_model = nn.Sequential(
    tabm.EnsembleView(k=k),
    tabm.LinearBatchEnsemble(d_in, d, k=k, scaling_init=scaling_init),
    ...
)

# 不好的：表格特征在线性层之前就被混合了。
bad_model = nn.Sequential(
    nn.Linear(d_in, d),
    tabm.EnsembleView(k=k),
    nn.ReLU(),
    ...
)

# 不好的：虽然在第一次线性变换之前生成了 k 种表示，但这些表示并不不同。从数学上讲，下面的代码与上一个例子等价。
bad_model = nn.Sequential(
    tabm.EnsembleView(k=k),
    nn.Linear(d_in, d),
    nn.ReLU(),
    ...
)

权重共享。 在选择 torch.nn.Linear（在集成成员之间完全共享线性层）、tabm.LinearBatchEnsemble（共享大部分权重）和 tabm.LinearEnsemble（不共享权重）时，请注意基于权重共享的参数高效集成策略（如 BatchEnsemble 和 MiniEnsemble）不仅能显著提高 TabM 的效率，还能提升其任务性能。因此，权重共享似乎是一种有效的正则化方法。然而，这种正则化的“最佳程度”以及它如何依赖于具体任务，仍有待进一步研究。

示例：无权重共享的简单集成模型

以下代码是对 tabm.MLPBackboneEnsemble 的重新实现：

k = 32
d_in = 24
d = 512
d_out = 1
dropout = 0.1

model = nn.Sequential(
    tabm.EnsembleView(k=k),

    # >>> MLPBackboneEnsemble(n_blocks=2)
    tabm.LinearEnsemble(d_in, d, k=k),
    nn.ReLU(),
    nn.Dropout(dropout),

    tabm.LinearEnsemble(d, d, k=k),
    nn.ReLU(),
    nn.Dropout(dropout),
    # <<<

    tabm.LinearEnsemble(d, d_out, k=k),
)

示例：MiniEnsemble

MiniEnsemble 是一种简单的参数高效集成策略：

1. 通过将对象分别经过 $k$ 个非共享且随机初始化的仿射变换，生成 $k$ 种不同的表示。
2. 将这 $k$ 种表示并行地送入一个共享的主干网络。可以使用任何主干网络。
3. 使用非共享的输出头进行预测。

以下代码是对 tabm.MLPBackboneMiniEnsemble 的重新实现。实际上，backbone 可以是任何类似 MLP 的模型。对 backbone 的唯一要求是能够支持任意数量的批量维度，因为 EnsembleView 会增加一个新的维度。另一种方法是在主干网络前后对表示进行重塑。

d_in = 24
d = 512
d_out = 1
k = 32

# 可以使用任何类似 MLP 的主干网络。
backbone = tabm.MLPBackbone(
    d_in=d_in, n_blocks=2, d_block=d, dropout=0.1
)
model = nn.Sequential(
    tabm.EnsembleView(k=k),

    # >>> MLPBackboneMiniEnsemble
    tabm.ElementwiseAffine((k, d_in), bias=False, scaling_init='normal'),
    backbone,
    # <<<

    tabm.LinearEnsemble(d, d_out, k=k),
)

示例：BatchEnsemble

以下代码是对 tabm.MLPBackboneBatchEnsemble 在 n_blocks=2 情况下的重新实现：

k = 32
d_in = 24
d = 512
dropout = 0.1
tabm_init = True  # TabM 风格的初始化
scaling_init = 'normal'  # 或 'random-signs'

model = nn.Sequential(
    tabm.EnsembleView(k=k),

    # >>> MLPBackboneBatchEnsemble(n_blocks=2)
    tabm.LinearBatchEnsemble(
        d_in, d, k=k,
        scaling_init=(scaling_init, 'ones') if tabm_init else scaling_init,
    ),
    nn.ReLU(),
    nn.Dropout(dropout),

    tabm.LinearBatchEnsemble(
        d, d, k=k,
        scaling_init='ones' if tabm_init else scaling_init
    ),
    nn.ReLU(),
    nn.Dropout(dropout),
    # <<<

    tabm.LinearEnsemble(d, d_out, k=k),
)

示例：自定义架构

以下是一个随机且很可能糟糕的架构，展示了该包中可用的各种层：

d_in = 24
d = 512
d_out = 1
k = 16
dropout = 0.1

model = nn.Sequential(
                                         #    (B, d_in) 或 (B, k, d_in)
    tabm.EnsembleView(k=k),              # -> (B, k, d_in)

    // 大部分权重共享
    tabm.LinearBatchEnsemble(            # -> (B, k, d)
        d_in, d, k=k, scaling_init=('random-signs', 'ones')
    ),
    nn.ReLU(),                           # -> (B, k, d)
    nn.Dropout(dropout),                 # -> (B, k, d)

    // 不共享权重
    tabm.BatchNorm1dEnsemble(d, k=k),    # -> (B, k, d)

    // 不共享权重
    tabm.LinearEnsemble(                 # -> (B, k, d)
        d, d, k=k
    ),
    nn.ReLU(),                           # -> (B, k, d)
    nn.Dropout(dropout),                 # -> (B, k, d)

    // 权重完全共享
    nn.Linear(                           # -> (B, k, d)
        d, d,
    ),
    nn.ReLU(),                           # -> (B, k, d)
    nn.Dropout(dropout),                 # -> (B, k, d)

    // 不共享权重
    tabm.ElementwiseAffine(              # -> (B, k, d)
        (k, d), bias=True, scaling_init='normal'
    ),

    // 权重完全共享
    tabm.MLPBackbone(                    # -> (B, k, d)
        d_in=d, n_blocks=2, d_block=d, dropout=0.1
    ),

    // 几乎所有权重共享
    tabm.MLPBackboneMiniEnsemble(        # -> (B, k, d)
        d_in=d, n_blocks=2, d_block=d, dropout=0.1,
        k=k, affine_bias=False, affine_scaling_init='normal',
    ),

    // 不共享权重
    tabm.LinearEnsemble(d, d_out, k=k),  # -> (B, k, d_out)
)

x = torch.randn(batch_size, d_in)
y_pred = model(x)
assert y_pred.shape == (batch_size, k, d_out)

将现有模型转换为高效集成模型

[!WARNING] 下文讨论的方法需要充分理解原始模型和高效集成的所有实现细节，因为它无法保证所得到模型的正确性和性能。

假设您有一个类似 MLP 的模型，并希望在不修改模型源代码的情况下快速评估高效集成对您的模型的潜力。那么，您可以创建一个单模型实例，并通过将其层替换为对应的集成层来将其转换为高效集成模型。为此，该包提供了以下函数：

• tabm.batchensemble_linear_layers_
• tabm.ensemble_linear_layers_
• tabm.ensemble_batchnorm1d_layers_
• tabm.ensemble_layernorm_layers_

例如，以下代码创建了一个由 $k$ 个独立 MLP 组成的完整集成模型：

d_in = 24
d = 512
d_out = 1

// 创建一个标准的 MLP 主干。
backbone = tabm.MLPBackbone(
    d_in=d_in, n_blocks=3, d_block=d, dropout=0.1
)

// 将该主干转换为高效集成。
k = 32
tabm.ensemble_linear_layers_(backbone, k=k)

// 组装最终模型。
model = nn.Sequential(
    tabm.EnsembleView(k=k),
    backbone,
    tabm.LinearEnsemble(d, d_out, k=k),
)
assert model(torch.randn(batch_size, d_in)).shape == (batch_size, k, d_out)

超参数

默认模型

TabM.make 允许用户使用默认超参数创建 TabM，并根据需要进行覆盖：

[!NOTE] 默认超参数并非“常量”，即它们会根据提供的参数而变化。

// 使用默认超参数的 TabM。
model = TabM.make(n_num_features=16, d_out=1)

// 自定义 n_blocks，其余超参数保持默认值的 TabM。
model = TabM.make(n_num_features=16, d_out=1, n_blocks=2)

默认优化器

目前，对于默认的 TabM，其默认优化器是 torch.optim.AdamW(..., lr=0.002, weight_decay=0.0003)。

arch_type

TL;DR：默认情况下使用 TabM。

• 'tabm' 是默认值，在大多数情况下应能提供最佳性能。
• 'tabm-mini' 可能在训练和/或推理速度上更快，同时性能下降不明显，不过此选项可能需要更精确地选择适合给定 k 值的 d_block 和 n_blocks。在某些情况下，如果更高的正则化程度对特定任务有益，'tabm-mini' 甚至可能带来略微更好的性能。
• 'tabm-packed' 仅出于完整性而实现。在大多数情况下，它会导致模型更慢、更重且性能较差。

k

[!TIP] 除第一点外，以下关于 k 的内容均受论文（arXiv v3）图 7 的启发。

• 如果您打算调整 k，建议分别运行多个独立的超参数调优实验，每次使用不同的但固定的 k 值。原因是改变 k 可能会影响其他超参数的最佳取值，若与其他超参数同时调整，则可能使超参数调优过程更加复杂。
• 在深度和宽度一定的情况下，适当增加 k 可以提升性能；但超过某个阈值后，性能可能不再提升，甚至有所下降。这种现象在 arch_type='tabm-mini' 时更为明显（未在图中显示）。
• 为了探索性实验，可以尝试使用较小的 k 值（如 24 或 16），通常仍能获得具有竞争力的结果（尽管更改 k 后可能需要重新调整其他超参数）。
• 一般而言，当增大 k 时，也应考虑相应增加 d_block 或 n_block（或两者）。其背后的直觉是，由于权重共享机制的存在，较大的集成规模可能需要更大的基础架构，才能有效容纳 k 个子模型。
• 如果您的数据集规模与论文中使用的数据集相近，除非您有充足的预算进行超参数调优，否则通常应避免使用 n_blocks=1。

num_embeddings

• 从历史来看，分段线性嵌入通常是用户的更常见选择。然而，在某些任务上，周期性嵌入可能是更好的选择。
• rtdl_num_embeddings 包的文档提供了关于嵌入超参数调优的建议。

初始化

本节概述了在使用 TabM.make 之外的 API 时，需要指定的初始化相关选项。

[!NOTE] 初始化相关设置对任务性能的影响程度取决于具体任务。通常，这些设置旨在充分发挥 TabM 的潜力，而非避免某些失败模式。

TabM 风格的初始化。tabm_init 是一个标志，用于触发基于 BatchEnsemble 模型的 TabM 风格初始化。简而言之，这是一种保守的初始化策略，使得 $k$ 个子模型仅在第一个集成层中有所不同，而在其他所有层中则完全相同（初始化时）。在基准测试中，tabm_init=True 表现出更好的默认策略。如果训练过程中 $k$ 个子模型坍缩为同一个模型，请尝试将 tabm_init=False。

缩放参数的初始化。start_scaling_init、scaling_init 和 affine_scaling_init 这些参数在略微不同的上下文中指代同一内容。TabM 使用一种非正式的经验法则，大致可以总结为：如果 TabM 主干之前存在“非平凡”的模块（例如 num_embeddings），则使用 "normal"；否则使用 "random-signs"。这仍然是 TabM 中尚未充分探索的一个方面。

初始化分块。类似 start_scaling_init_chunks 或 scaling_init_chunks 的参数会触发缩放参数的启发式分块初始化，其中各分块的大小之和必须精确等于主干输入大小（即 d_in）。默认情况下，在 TabM 中，“分块”被简单定义为特征表示的大小（参见 tabm.TabM.__init__ 中的 d_features 变量），这意味着在初始化过程中，每个特征恰好会采样一个随机标量。在其他场景下，如何定义分块可能并不明确。在这种情况下，一种可行的方法是先从 None 开始，必要时再通过试错找到更优的方案。

超参数调优

[!TIP] 上文提供的通用注意事项同样适用于自动超参数调优。

在论文中，为了调优 TabM 的超参数，使用了 Optuna 的 TPE 采样器（Optuna 文档），在较小的数据集上进行了 100 次迭代，在较大的数据集上进行了 50 次迭代。如果对达到最高性能并非至关重要，则进行 30–50 次迭代即可得到一个较为合理的配置。

下表提供了论文中使用的超参数分布。 请根据您的实际设置并结合前文内容调整这些分布，尤其是在迭代次数较少的情况下。

超参数	TabM	带有 PiecewiseLinearEmbeddings 的 TabM
k	Const[32]（不调优）	与 TabM 相同
n_blocks	UniformInt[1, 5]	UniformInt[1, 4]
d_block	UniformInt[64, 1024, step=16]	与 TabM 相同
lr	LogUniform[1e-4, 5e-3]	与 TabM 相同
weight_decay	{0, LogUniform[1e-4, 1e-1]}	与 TabM 相同
n_bins	无	UniformInt[2, 128]
d_embedding	无	UniformInt[8, 32, step=4]

实用提示

推理效率

如论文第 5.2 节所示，可以在训练完成后大幅剪枝 TabM 的子模型数量（即降低 k），代价是性能略有下降。理论上，还有更先进的算法可用于选择最佳子模型子集，例如 Caruana 等人提出的算法（AMLTK 文档），但这些算法并未在论文中进行分析。此外，请注意，先以 k=32 训练后再选择 k=8 子模型，其效果会优于直接以 k=8 训练。

API

要浏览包的 API 和文档字符串，可采取以下任一方法：

• 克隆此仓库并运行 make docs。
• 在 GitHub 上打开源代码并使用符号面板。
• 在 VSCode 中打开源代码并使用大纲视图。

若想在不克隆或安装任何东西的情况下列出所有可用项，可运行以下代码片段：

uv run --no-project --with tabm python -c """
import tabm

for x in sorted(
    x
    for x in dir(tabm)
    if getattr(getattr(tabm, x), '__module__', None) == 'tabm'
    and not x.startswith('_')
):
    print(x)
"""

TabM 快速上手指南

TabM (Tabular Deep Learning with Parameter-Efficient Ensembling) 是一个基于 PyTorch 的强大表格深度学习模型。它通过参数高效的方式模拟多个 MLP（多层感知机）的集成，在保持高预测性能的同时，显著提升了训练和推理的效率。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python 版本: 3.8 或更高
• 核心依赖:
  • torch (PyTorch)
  • tabm (本工具包)
  • rtdl-num-embeddings (可选但推荐，用于提升数值特征的处理性能)

提示：国内用户建议使用清华源或阿里源加速 Python 包的安装。

安装步骤

使用 pip 进行安装。为了获得最佳下载速度，推荐使用国内镜像源：

pip install tabm -i https://pypi.tuna.tsinghua.edu.cn/simple

如果需要支持更高级的特征嵌入（推荐用于生产环境），请同时安装辅助库：

pip install rtdl-num-embeddings -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下示例展示了如何构建一个基础的 TabM 模型并进行前向推理。

1. 基础模型构建（无特征嵌入）

适用于快速验证或简单场景。模型会自动处理数值特征和类别特征（采用 One-hot 编码）。

import torch
from tabm import TabM

# --- 数据集配置 ---
d_out = 1  # 输出维度，例如回归任务为 1
batch_size = 256
n_num_features = 24  # 数值特征数量
cat_cardinalities = [3, 7]  # 类别特征的基数列表（每个类别特征的唯一值数量）

# --- 创建模型 ---
# TabM.make 会根据参数自动填充默认的超参数
model = TabM.make(
    n_num_features=n_num_features,
    cat_cardinalities=cat_cardinalities, 
    d_out=d_out,
)

# --- 准备输入数据 ---
x_num = torch.randn(batch_size, n_num_features)
x_cat = torch.column_stack([
    # 第 i 个类别特征的取值范围必须在 [0, cat_cardinalities[i]) 之间
    torch.randint(0, c, (batch_size,)) for c in cat_cardinalities
])

# --- 前向推理 ---
y_pred = model(x_num, x_cat)

# TabM 输出形状为 (Batch_Size, k, d_out)，其中 k 是集成模型的数量
assert y_pred.shape == (batch_size, model.k, d_out)
print(f"预测结果形状：{y_pred.shape}")

2. 进阶模型构建（带特征嵌入）

在实际表格任务中，使用特征嵌入模块（如 LinearReLUEmbeddings）通常能获得更好的性能。

import torch
from tabm import TabM
from rtdl_num_embeddings import LinearReLUEmbeddings

# 配置
d_out = 1
batch_size = 256
n_num_features = 24

# --- 创建带嵌入的模型 ---
model = TabM.make(
    n_num_features=n_num_features,
    num_embeddings=LinearReLUEmbeddings(n_num_features), # 启用数值特征嵌入
    d_out=d_out
)

# 准备输入 (仅需数值特征，类别特征可根据需要额外添加)
x_num = torch.randn(batch_size, n_num_features)

# 前向推理
y_pred = model(x_num)

assert y_pred.shape == (batch_size, model.k, d_out)

关键使用说明

• 训练策略：训练时必须独立优化 k 个预测结果的损失，即优化 平均损失 (mean loss)，而不是先平均预测结果再计算损失。
• 推理策略：在推理阶段，通常需要对 k 个预测结果进行平均以获得最终预测值。
  • 回归任务：直接对输出值求平均。
  • 分类任务：通常先对 Softmax 后的概率求平均，而不是对 Logits 求平均。