LimiX-2：表格数据上的Masked Modeling自监督预训练实战

说实话第一次认真把 masked modeling 搬到表格模型上做预训练是在 LimiX-2 这个项目里。之前我一直在折腾表格数据上的深度学习模型试过各种结构也写过不少自监督方案但效果总是不如 XGBoost 这类树模型稳。表格数据和图片、文本最大的区别就是它没有天然的空间顺序和语义顺序特征之间是异构的、稀疏的、带有大量缺失和噪声的。所以当 LimiX-2 决定把 masked modeling 作为核心预训练任务时团队内部争论了很久——有人说这不就是把 BERT 的掩码语言建模硬搬到表格上吗其实做下去才发现问题远没有那么简单但一旦把掩码单元重建目标特征关系这三件事想清楚收益是实打实的。这篇文章就把 LimiX-2 里关于 masked modeling 的完整设计思路、实现细节、实验结论和踩过的坑整理出来。适合已经在做表格深度模型、想引入自监督预训练但不知道从哪下手的团队参考也适合想理解表格领域掩码建模到底在学什么的算法工程师。内容偏实践不只是讲原理。1. 为什么表格模型需要单独的 masked modeling 方案1.1 表格数据的预训练为什么这么难图像有像素网格文本有 token 序列这两类数据都自带一种局部结构相邻像素大概率相关相邻 token 组成词组和句子。表格数据完全不一样每一行是一条独立样本每一列是一个有业务含义的字段比如年龄、收入、职业、历史逾期次数。列与列之间的相关性没有固定模式同一列内部不同取值之间也不存在顺序邻近性。这导致一个直接后果你在 NLP 里用一句BERT 式掩码预训练的思路在表格上根本没法直接抄。因为你不知道掩码的单元应该是什么不知道用什么信号来重建被掩码的部分更不清楚预训练阶段学到的表示在下游任务中是否真的有用。更尴尬的是表格数据本身信息密度不低。XGBoost 等树模型靠特征分裂就能抓住大部分非线性关系而且训练快、可解释性也好。深度模型如果只是堆 MLP优势不大如果引入 Transformer 结构又非常依赖特征工程和预处理稍有不慎效果反而不如树模型。所以很多团队干脆放弃深度模型直接走 GBDT 路线。但问题是树模型本质上不做特征表示学习它是在人工特征已经给定的空间里做分箱切分。对于高基数类别特征、稀疏特征、跨特征组合这类问题树模型的表达能力是有上限的。这正是深度模型的机会——前提是你能找到一种有效的自监督预训练方法让模型在没有标签的情况下先把表格的内在结构学出来。LimiX-2 选择 masked modeling就是想在结构学习这件事上做一个真正的尝试。1.2 自编码器和对比学习为什么不够在 LimiX-2 之前我们试过两种常见的自监督路线。第一种是自编码器式的重建。把整个表格输入压缩成低维向量再重建出完整表格。听起来合理但问题在于重建任务太宽松了模型只要记住每个特征的边缘分布甚至直接把输入恒等复制loss 就能很低。换句话说自编码器学到的主要是拷贝能力而不是特征关系理解能力。第二种是对比学习。构造一张表格的两个扰动视图拉近同一行样本的表示。这个思路在图像上很有效因为图像的随机裁剪、色彩抖动是天然的、语义保持的增广方式。表格数据没有这种便宜又可靠的增广策略你不可能随便把年龄从 30 岁扰动到 40 岁还指望语义不变。我们试过按列扰动、按行采样、特征 dropout 等方式构造正样本效果都很不稳定有时候增广强度稍微大一点预训练收益就变成了噪声。对比学习的另一个麻烦是它依赖 batch 内的负样本设计表格数据的类别分布通常很不均匀负样本采样稍有偏差模型学到的东西就会有偏。masked modeling 和上面两条路线最大的不同是它不要求模型去做全局压缩或视图对齐而是让模型在局部信息缺失的情况下去猜测缺失内容。这个任务天然地迫使模型去捕捉特征之间的条件依赖关系要知道收入被掩码后的大致取值就必须学会用职业年龄城市等级这些特征去推断。这恰好就是表格数据里最有价值的知识——特征与特征之间的依赖结构。1.3 LimiX-2 的基本假设LimiX-2 这个代号来自我们内部对表格模型能力边界的探索系列LimiX 可以理解成 Limit Exploration 的意思。第一版 LimiX 是一个常规的深度表格基线模型效果中规中矩。到 LimiX-2 时我们做了一个关键假设表格数据的核心知识不在于单个特征的取值而在于特征之间的条件分布。如果模型能通过掩码重建学会这些条件分布下游任务上的表现就会有本质提升。这个假设并不新奇很多研究也在提但真正的问题是怎么把它落地。LimiX-2 的 masked modeling 设计本质上就是在回答三个问题掩码什么、怎么掩码、用什么目标去重建。2. 掩码单元与掩码策略从掩码什么开始设计2.1 掩码单元的选择在文本里掩码单元是 token在图像里掩码单元是 patch。表格里的单元最自然的选择是单元格cell也就是第 i 行第 j 列的值。但只掩码单元格是不够的。我们做了一个很直观的实验如果每次只随机掩码单个单元格模型会很轻松地通过其他列的强相关特征重建出来学到的东西偏局部。比如年龄被掩码时通过出生年份几乎可以无损还原那模型根本不需要理解更复杂的跨列关系。所以 LimiX-2 的掩码单元思路是混合的单元格级掩码随机挑一部分单元格强制模型做细粒度重建。整行掩码随机选中部分行把该行的所有特征都掩掉模型只能利用整个表格的列分布信息来生成该行的合理取值。这相当于强制模型学习列的联合分布。整列掩码随机选中部分列把整列都掩码掉模型需要从其他列推断这一列的分布。这比单元格掩码难很多但对跨列依赖的表征能力要求更高。局部块掩码在行列方向上做一个矩形块把连续的一片区域掩码掉模拟现实中一批样本的某几个字段同时缺失的场景。这种混合设计的好处是不同掩码单元对应不同粒度的特征依赖学习。单元格掩码学局部相关性整列掩码学全局跨列关系整行掩码学联合分布局部块掩码则提高模型对缺失模式多样性的鲁棒性。2.2 五种掩码策略的优劣势对比我把 LimiX-2 里实际用过的掩码策略整理成一个表方便你对比掩码策略操作方式学习目标优点缺点Random Cell Masking均匀随机掩码单元格局部特征依赖简单、稳定、收敛快容易被强相关特征作弊式重建Row Masking整行掩码列联合分布迫使模型学习行内特征的结构关系难度过高早期训练不稳Column Masking整列掩码跨列全局依赖对特征工程的依赖性强列数少时不适用信息不足Block Masking连续矩形块掩码缺失模式鲁棒性贴近真实缺失场景掩码形状要调参Value-aware Masking按特征重要度/分布加权掩码难重建特征的重点表示提高困难特征的表示质量需要额外计算特征重要性最初我们只用了 Random Cell Masking简单但效果一般。后来加入 Block Masking 和 Column Masking预训练后的微调表现明显提升。特别是 Column Masking它让模型不再依赖偷看某几个强特征来蒙混过关逼着模型建立多特征融合的表示。2.3 掩码率应该怎么定NLP 里 BERT 的掩码率是 15%这个值在表格里不能照搬。原因很简单文本的 token 之间有很强的序列依赖掩码 15% 已经能产生足够的预测难度但表格里特征之间的相关性通常更稀疏、更非线性15% 的掩码率可能只让模型学到了微不足道的简单重建。我们在 LimiX-2 上做了掩码率扫描实验范围从 10% 到 60%。结果发现25% 到 40% 之间是一个比较稳的区间。掩码率太低预训练 loss 很快就降到底下游微调收益很小掩码率太高输入被破坏得太严重模型要用太多精力去猜测完全缺失的行结构反而学不到干净的表示。有趣的是最佳掩码率和特征维度有关。特征是 20 维左右时30% 掩码率效果最好特征是 100 维以上时25% 反而更稳。我们推测原因是高维表格里有很多冗余特征模型可以通过冗余信息重建被掩码的部分所以不需要太高的掩码率低维表格里每个特征都更珍贵需要更高的掩码率来制造足够的学习信号。后面训练 LimiX-2 时我们干脆用了一个动态掩码率前 20% 的训练轮次用较低的 20%慢慢升到 35%最后 10% 轮次再降到 15%。有点类似课程学习效果比固定掩码率略好而且收敛更稳定。2.4 掩码 token 的表示细节这个细节很容易被忽略但影响非常大。BERT 里有一个专门的 [MASK] token 来表示被掩码的位置LimiX-2 也需要类似的掩码占位符。问题是表格里不同列的类型完全不同数值列的掩码占位不能和类别列的掩码占位混用。我们的做法是每一列单独维护一个可学习的掩码嵌入向量。数值列给它一个初始化为 0 的向量类别列给它一个随机初始化的向量。这样模型在看输入序列时能够明确区分这一列被完全掩码了还是这一列本身是缺失值。这一点我们在实验里对比过用所有列共享一个掩码嵌入和每列独立掩码嵌入下游任务表现差了大概 0.5 到 1 个点。原因很好理解不同列被掩码后需要依赖的信息源完全不同如果掩码占位符长得都一样模型就要额外花精力去区分现在缺的是哪个字段。3. 模型架构与训练目标LimiX-2 的具体实现3.1 特征编码从异构原始输入到统一序列LimiX-2 的底层编码器是 Transformer Encoder。要进 Transformer先要把异构的表格特征变成一组向量序列这一步我称之为特征 Tokenizer。具体做法是数值特征先做分位数归一化把取值映射到 [0,1] 区间然后通过一个线性层映射到 d_model 维。低基数类别特征用 nn.Embedding 查表得到向量。高基数类别特征不直接 embedding而是先做 target encoding在预训练阶段用训练集整体统计再经过线性层。这一步是为了避免高基数类别 embedding 在预训练时陷入记忆稀疏 ID而不是学习特征关系。每个特征向量加上一个可学习的特征类型嵌入Feature Type Embedding用来区分这是一个数值列还是这是一个类别列。还有一个关键处理位置编码。表格特征没有天然的顺序所以我们不是用正弦位置编码而是用一个可学习的列位置嵌入Column Position Embedding。它的作用不是表达顺序而是让 Transformer 区分当前 attention 权重应该分配给哪一列相当于给每个特征一个身份标识。3.2 编码器与预测头编码器就是标准的 Transformer Encoder我们用的是 6 层、hidden size 256、8 个 attention head。参数量不大因为表格数据不像文本和图像那么海量模型太大会在预训练阶段就过拟合后面微调反而吃亏。编码器输出的每个 token 向量对应输入序列里的一列特征。预测头是一个分层的 MLP数值列输出一个标量训练目标用 MSE。类别列输出一个 logits 向量训练目标用交叉熵。和 BERT 类似我们只对被掩码的位置计算 loss。但有一个小改动不同的特征类型loss 权重不一样。默认设置是数值列 loss 权重为 1.0类别列权重为 0.8。为什么类别列权重低一些因为类别特征的重建难度通常更高如果权重太大会让模型过度关注猜中类别忽略数值特征之间的关系建模。3.3 预训练和微调的一致性设计LimiX-2 的预训练和微调不是完全脱节的。我们做了一件很重要的事在微调阶段也会以很小的概率比如 5%随机掩码掉部分输入特征让模型尽量适应特征缺失的情况。这个设计的初衷来自一个实际观察线上预测时经常会出现某个字段临时拿不到的情况。如果模型只在完整数据上训练遇到这种输入就会产生奇怪的预测。我们在微调时加入轻量掩码既能让模型复用预训练阶段学会的从其他特征推断缺失特征的能力又能提高部署时的鲁棒性。当然微调的损失主力还是下游任务本身的监督 loss掩码重建 loss 只是一个小小的辅助约束。我们试过把两个 loss 直接相加权重比是 10:1效果最好。辅助 loss 权重再大一些反而会影响下游任务的拟合。3.4 训练策略细节这里写一些 LimiX-2 实测下来比较关键的超参数和训练细节优化器用 AdamW初始学习率 1e-3warmup 比例 10%使用余弦退火。Batch size 我们用了 4096。表格数据的一行样本通常很短batch 太小的话学习信号不够稳定。预训练轮次一般在 300 epoch 以内。LimiX-2 在 200 个 epoch 之后指标基本平稳再加轮次容易过拟合到掩码重建任务上。梯度裁剪设为 1.0这个很重要。masked modeling 的 loss 偶尔会出现突变如果不裁剪训练直接震荡甚至发散。数值列的 MSE loss 和类别列的交叉熵 loss 量级差异很大我们把数值列做了 z-score 标准化让两类 loss 的初始量级差不多否则模型训练初期会完全忽略数值列的重建。3.5 代码层面的掩码矩阵实现这里给出一个简化的掩码生成逻辑方便复现。核心是一次性生成整个 batch 的掩码矩阵而不是在 Python 里逐行循环不然训练速度会慢得离谱。import torch def generate_mask_matrix(batch_size, num_features, mask_ratio): # 返回 shape 为 (batch_size, num_features) 的 0/1 掩码矩阵1 表示被掩码 mask torch.rand(batch_size, num_features) mask_ratio return mask.float() def apply_mask(features, mask, mask_embeddings): # features: (batch, num_features, d_model) # mask: (batch, num_features, 1) # mask_embeddings: (num_features, d_model) masked_features features * (1 - mask) mask * mask_embeddings.unsqueeze(0) return masked_features这里 mask_embeddings 是每个特征列专用的掩码嵌入矩阵和模型一起训练。实际代码里还可以加一个保证每行至少掩码一个特征的逻辑避免出现完全未掩码的行导致该行的梯度为空。4. 实验验证与消融masked modeling 到底带来了什么4.1 实验设置与数据集我们在三个数据集上做了完整的验证Adult Income收入预测15 维特征、Forest Cover Type植被类型分类54 维特征、IEEE-CIS Fraud欺诈检测高维稀疏特征。前两个是公开数据集方便复现第三个是相对接近真实业务的数据集带有大量缺失值和噪声。对比基线选了四种XGBoost作为树模型代表。普通 MLP不预训练直接监督训练。FT-Transformer用 Transformer 但不做 masked modeling 预训练。LimiX-1上一代模型用自编码器式预训练。所有模型都控制了大致的参数量和训练预算避免不公平对比。4.2 主实验结果先看一组数据以 AUC / Accuracy 为指标数值为多次实验均值模型Adult (Acc)Forest (Acc)IEEE-CIS (AUC)XGBoost0.8630.9350.908MLP0.8480.9110.872FT-Transformer0.8570.9240.890LimiX-1 (AE)0.8550.9200.885LimiX-2 (masked modeling)0.8720.9410.916从结果看LimiX-2 在三个数据集上都拿到了最好的成绩。但我们更关心的是两个现象第一LimiX-2 在 Adult 这种小数据集上的提升幅度最大比 FT-Transformer 高了 1.5 个点比 XGBoost 高 0.9 个点。这说明 masked modeling 在小数据场景下的预训练收益非常明显模型通过自监督学到了先验结构减少了对大规模标注数据的依赖。第二在 IEEE-CIS 这种高维稀疏特征数据集上LimiX-2 的 AUC 也超过了 XGBoost。这在表格深度模型里不太常见因为我们没有对它做极其精细的特征工程。masked modeling 让模型学会自动捕捉高维特征之间的关联这比人工做特征组合要省力得多。4.3 消融实验哪种掩码策略贡献最大我们对掩码策略做了详细的消融结论是Column Masking 的贡献最大Block Masking 次之Random Cell Masking 最小。预训练时使用的掩码策略Adult (Acc)Forest (Acc)随机单元格掩码0.8600.928随机单元格 整列掩码0.8680.936全部策略混合0.8720.941列掩码的贡献大本质原因是它迫使模型面对一个特征完全不可见的情况这时候模型必须综合多个其他列的信息才能做出合理重建。这个能力和下游任务高度相关——下游预测时往往也不是单个特征决定结果而是多个特征共同作用。另外我们还测了一个有意思的变体如果把掩码后的 reconstruction loss 换成对比学习 loss来预训练同一个 Transformer效果会变成什么样结果是对比学习版本的预训练收益几乎为零甚至在某些情况下还会掉点。这进一步说明对表格数据来说masked modeling 这种条件依赖重建的预训练范式比视图一致性拉近的范式更匹配数据本身的特性。4.4 从嵌入可视化看模型学到了什么我们画了预训练后特征嵌入的相似度矩阵。发现了一个很明显的模式在业务上本来就有强关联的特征比如年龄和退休状态、收入和职业等级它们的特征嵌入在学习之后距离变得非常接近。没有预训练的模型这些特征嵌入之间的关系基本是随机的。这说明 masked modeling 的确在让模型自发地去发现特征之间的依赖结构而不是简单地记忆每个特征的分布。这也是为什么微调阶段只需要少量标注数据模型就能快速拟合下游任务——因为它已经提前学到了一张特征关系地图。5. 踩坑实录与工程化建议5.1 坑一标准化和归一化导致的信息泄漏数值特征在进模型之前通常要做分位数归一化或 z-score 归一化。这里有一个我们在第一次做 masked modeling 时踩过的坑如果归一化的均值和标准差是在包含全部样本的完整表格上计算的然后你再去做掩码重建模型的输入里多多少少已经携带了被掩码值的统计信息预训练效果会虚高。举个例子你人为掩码了某一行的高收入值但如果归一化时用的全局均值已经包含了这个高收入值那模型的输入向量并不是完全看不到这个值而是隐约含了它的统计贡献。正确做法是归一化参数只在训练集的未掩码部分或者说用全体训练数据统计但预训练时不使用同一批统计去评估重建 loss上计算并保持固定。更稳妥的方案是把归一化和掩码顺序颠倒一下先对原始值做掩码再对未掩码部分计算归一化统计量然后应用到整个 batch。这样模型所看到的输入统计信息确实只来自可见部分。5.2 坑二缺失模式missing pattern本身也是信息很多真实表格里特征缺失不是完全随机的而是和业务逻辑有关。比如欺诈检测数据集中某些字段只在特定渠道的样本里存在缺失本身就是一个强信号。我们在预训练时一开始把缺失值也当成普通值处理用 mask embedding 去填充结果下游任务效果反而变差。后来我们加了一个处理对每个特征额外生成一个 is_missing 标志位和归一化后的数值一起拼接输入模型。这样模型能区分这个值本来就没有和这个值被有意掩码两种情况。这个改动在 IEEE-CIS 数据集上带来了 1 个点以上的提升。有朋友可能会问那 masked modeling 本身不就是在制造缺失值吗我们在设计时确实特意区分了这两种缺失原始数据里真实缺失的字段不参与重建 loss而人为掩码的字段才参与重建。这样才能保证模型的重建目标干净不被预测原本缺失值的假任务干扰。5.3 坑三特征维度太少时masked modeling 可能失效LimiX-2 在 Adult 这种 15 维特征的数据上有收益但如果特征维度只有 5 到 6 个比如一些极简的业务表masked modeling 的预训练收益就非常有限甚至为负。原因也简单特征太少时掩码一部分特征会让剩余信息严重不足模型学到的主要是猜边缘分布而不是捕捉特征间关系。这种情况下预训练带来的收益无法覆盖额外训练成本。解决方法有两个一个是降低掩码率比如从 30% 降到 15%另一个是引入外部相关性比如把一些统计衍生特征加进来再预训练。如果实在没有额外特征我们建议干脆放弃 masked modeling直接用监督学习别把简单问题复杂化。5.4 工程层面的性能优化masked modeling 预训练的一个隐性成本是掩码重建损失要计算所有掩码位置的梯度这比普通监督训练多了一部分计算量。在 LimiX-2 里我们做了一个很小的优化只在被掩码的位置计算 loss反传时通过掩码矩阵把梯度置零。PyTorch 里直接乘掩码即可避免对全量输出做 loss。另外大规模数据集上做整列掩码时会出现某些列在同一个 batch 里完全没有可见样本参与建模的情况这时梯度会不稳定。我们的做法是保证每个 batch 里至少保留 30% 的列完全可见不参与掩码。这个策略没有体现在论文里但实测对训练稳定性帮助很大。5.5 关于 LimiX-2 之后的一些体会LimiX-2 给我最大的启发不是masked modeling 有效而是表格模型的自监督设计必须从数据的结构出发而不是从方法出发。BERT 的掩码建模适合文本ViT 的掩码建模适合图像LimiX-2 的掩码建模只适合表格。直接把其他领域的方案搬过来往往只能得到看似公平、实则无效的结果。如果你打算在自己的表格数据上复现这套方案我的建议是先用一个小数据集跑通整套流程画一下特征嵌入相似度看看模型有没有学到合理的特征关系。这一步只需要半天时间但能帮你省下大量在无用方向上的试错成本。之后再逐步加大数据规模把掩码率、掩码策略做几次网格搜索挑出最适合你数据的配置。至于 LimiX-2 本身我们后续还在尝试把 masked modeling 的目标从特征重建升级成特征关系重建比如让模型预测两个特征之间的相关性是否被破坏。这个方向还在实验阶段但初步结果让我觉得表格数据的自监督学习还远没有到头。