CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

技术摘要：CRUMB —— 通过分布匹配上下文批处理实现高效的先验拟合网络推理

问题陈述
先验拟合网络（PFNs），如 TabPFN 和 TabICL，代表了一类能够进行上下文学习（in-context learning）的表格基础模型。它们接受一个完整的已标记训练集作为上下文，并通过单次前向传播生成测试查询的预测，在处理中小规模数据集时，其表现通常优于传统的梯度提升方法。然而，一个关键的瓶颈限制了其可扩展性：基于 Transformer 的 PFNs 中的自注意力机制随训练样本数量 ($N$) 呈二次方级增长 ($O(N^2)$)。这使得在处理大规模数据集（$N > 50,000$）时，计算开销变得难以承受。

现有的解决方案在准确性和效率之间面临权衡。均匀子采样通过随机丢弃训练数据来提高速度，但会牺牲信息量。依赖查询的方法（如 $k$-最近邻，$k$NN）通过为每个测试点定制上下文来提高准确性，但却无法进行批处理，需要针对 $T$ 个测试点进行 $T$ 次独立的前向传播。其他方法，例如混合上下文提示器（Mixture of In-Context Prompters, MICP），试图通过对训练数据进行聚类来批处理查询，但它们可能无法使训练上下文与测试数据的特定分布保持一致，特别是在存在协变量偏移（covariate drift）的情况下。

方法论
作者提出了 CRUMB（基于最小化 MMD 聚类检索的上下文批处理），这是一个三阶段、架构无关的推理封装器，无需重新训练模型。该方法通过以下阶段解决了上下文质量与批处理效率之间的矛盾：

1. 测试端聚类： 使用 $k$-means 对标准化后的输入特征对测试查询 ($D_{test}$) 进行划分，将其分为 $K$ 个簇 ($C_1, \dots, C_K$)。这使得系统能够将相似的测试点分组，从而共享一个共同的训练上下文。
2. 分布匹配的上下文选择： 对于每个测试簇 $C_k$，从完整的训练池 ($D_{train}$) 中选择一个小的训练子集 $S_k$（规模 $n \ll N$）。与几何邻近方法（如质心最近邻）不同，CRUMB 通过贪婪地最小化测试簇经验分布与候选训练子集之间的**最大均值差异（Maximum Mean Discrepancy, MMD）**来选择 $S_k$。
   • 这是通过**核硬核算法（kernel herding）**实现的，这是一种贪婪算法，通过迭代选择训练点来最小化 MMD 目标函数。该目标函数平衡了两个项：奖励与测试簇的接近度（相关性）以及惩罚所选训练点之间的冗余度（多样性）。
   • 为了加速这一过程，作者采用了随机傅里叶特征（Random Fourier Features, RFF）来近似高斯 RBF 核，并使用批处理贪婪选择。
3. 批处理推理： PFN 模型运行 $K$ 次独立的前向传播。在每次传播中，模型接收选定的训练子集 $S_k$ 作为上下文，并将整个簇 $C_k$ 作为查询。这使得总前向传播次数从 $T$（针对每个查询的 $k$NN）减少到 $K$，在保持分布对齐的同时，实现了高效的 GPU 利用。

核心贡献

• 基于 MMD 的上下文检索： 本文引入了贪婪 MMD 最小化（核硬核算法）作为选择训练子集的机制。作者证明，这种分布匹配方法在相同的聚类框架内，优于质心最近邻和 Voronoi 均匀检索策略。
• 通过测试端聚类的批处理上下文选择： 通过对测试查询进行聚类，并为每个簇分配一个分布匹配的训练上下文，CRUMB 在不牺牲查询相关性的情况下实现了高效的批处理。
• 对协变量偏移的鲁棒性： 该方法能自然地处理协变量偏移。由于训练上下文是根据当前的测试批次分布（而非固定的训练分布）进行选择的，CRUMB 能够适应数据的偏移。相比之下，依赖固定训练簇的 MICP 在测试分布偏离训练分布时性能会显著下降。

实验结果
作者在 TabArena 基准测试（包含 51 个多样化的表格数据集）上对三种 PFN 架构进行了评估：TabPFNv2、TabICLv1 和 TabICLv2。

• 性能对比基准： 在 TabArena 基准测试中，在相同的上下文预算下，CRUMB 的表现显著优于均匀子采样和 MICP 技术。虽然 CRUMB 在性能上在统计上接近逐查询 $k$NN，但它通过固定的 $K$ 次前向传播（例如 $K=20$）而非 $T$ 次完成了任务，从而实现了推理时间的实质性缩减（例如，在报告的实验中为 35s 对比 361s）。
• 大规模数据集： 在高达 120,000 个样本的数据集（如 GiveMeSomeCredit, Higgs）上的实验表明，采用早停机制（当 MMD 改善趋于平缓时终止上下文选择）的 CRUMB 比 MICP 更频繁地获胜，尤其是在使用 TabPFNv2 和 TabICLv1 时。
• 协变量偏移： 在受控的漂移实验中（测试分布沿第一主成分发生偏移），CRUMB 相对于 MICP 的优势随漂移强度的增加而单调增长。在零漂移时，CRUMB 领先 +4.9%；在最大漂移（完全超出样本范围）时，差距扩大至 +17.1%。这归功于 CRUMB 能够重新调整训练上下文以匹配偏移后的测试分布，而 MICP 的固定训练簇则会发生失配。

意义与主张
本文声称 CRUMB 使大规模表格数据集的先验拟合网络推理变得切实可行，且无需改变架构或进行模型微调。通过结合测试端聚类与分布匹配的上下文检索，CRUMB 在逐查询检索的准确性与均匀批处理的效率之间找到了“甜点区”。

作者强调，核心创新在于其测试自适应设计：即通过聚类测试点并基于分布对齐（MMD）而非仅仅基于几何邻近度来选择上下文。这种设计提供了针对协变量偏移的天然鲁棒性机制，而在这种场景下，现有的上下文选择方法往往会失效。文章总结道，虽然 CRUMB 继承了底层 PFN 预训练的性能上限，但它有效地消除了阻碍 PFN 应用于大规模数据集的计算障碍。

局限性
作者承认，CRUMB 需要预先获取完整的测试查询集以进行聚类，因此在不进行修改（例如缓存质心，这实际上退化为 MICP 式的路由）的情况下，并不适用于真正的在线流式场景。此外，对欧几里得 $k$-means 和各向同性高斯核的依赖可能会限制其在高度维度或混合类型特征空间中的表现。最后，作为一个纯粹的推理封装器，CRUMB 无法补偿 PFN 预训练先验与特定数据集之间较差的匹配度。