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

会议: NeurIPS 2026
arXiv: 2606.11473
OpenReview: https://openreview.net/forum?id=z3XMLStVPC
领域: 优化/理论
关键词: 先验拟合网络, 上下文学习, 最大均值差异, 批量推理, 协变量偏移

一句话总结¶

针对先验拟合网络（PFN）在面对大规模表格训练集时自注意力二次复杂度导致推理不可行、而逐样本 k 近邻检索又无法批量推理的瓶颈，提出 CRUMB 推理包装器，通过对测试查询聚类并在各簇内贪心最小化 MMD 选取小规模分布匹配训练上下文，在保持批量推理吞吐的同时大幅提升精度与抗协变量偏移韧性。

研究背景与动机¶

先验拟合网络（PFN，如 TabPFNv2、TabICLv1/v2）是一类极具前景的表格基础模型。它们在海量合成先验数据上预训练以逼近后验预测分布，因而在下游表格任务上无需微调或逐数据集梯度优化，仅需将全部带标签训练集作为上下文、测试点作为查询，通过单次前向传播即可完成上下文学习。在中小规模表格数据集上，PFN 的表现甚至超越了 CatBoost、XGBoost 等高度竞争的梯度提升树算法。然而，由于 Transformer 中样本间自注意力机制的计算开销和内存占用随训练上下文规模 $N$ 呈二次方增长（$O(N^2)$），直接在超大规模表格数据集（如 $N > 50{,}000$）上运行全局全上下文推理在显存和时间上均不可行。

为了缓解这一扩展瓶颈，现有的上下文筛选方案主要走向了两个极端。简单的均匀随机下采样会丢失大量几何细节与关键判别信息，严重损伤预测性能；基于查询的局部检索（如每个测试点独立执行 $k$ 近邻检索）虽然能为各个测试样本挑选相关性高的训练样本，但因为每个测试点诱导出的上下文各不相同，彻底破坏了 GPU 向量化批量前向传播的前提，导致 $T$ 个测试样本必须退化为 $T$ 次串行的独立 PFN 评估，在大测试集下推理耗时剧增。与此同时，MixturePFN 中提出的 MICP 尝试在训练集侧聚类并路由测试点，但在面对测试分布偏移（协变量漂移）时，其固化的训练聚类中心容易与测试查询分布严重脱节。

解决该矛盾的核心在于：如何在支持批处理（即多个测试查询共享同一个紧凑训练上下文）的前提下，让该上下文在特征空间分布上高保真地对齐这些测试查询。核心 idea：将上下文检索与批处理通过“测试端聚类 + 贪心核分布匹配”相统一，在测试查询侧聚类以实现分批向量化推理，并借助贪心最小化最大均值差异（MMD，Kernel Herding）为每个测试簇挑选兼具相关性与多样性的紧凑训练子集，天然免疫协变量漂移。

方法详解¶

整体框架¶

CRUMB（Clustered Retrieval Using Minimised-MMD Batching）是一个完全免训练、与 PFN 模型架构无关的推理外壳。其整体流程划分为三个串行阶段：首先在测试查询侧进行特征标准化并执行 $k$-means 聚类；随后针对每个测试簇，利用贪心最小化最大均值差异（MMD）从完整训练集中挑选大小为 $n \ll N$ 的代表性训练子集；最后将每个测试簇与所选训练上下文拼接，执行单次向量化 PFN 前向推理，总共仅需 $K$ 次前向传播。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：完整训练集与测试查询"] --> B["测试端划分聚类<br/>按输入特征将测试集聚类为 K 簇"]
    B --> C["贪心 MMD 训练子集检索<br/>Kernel Herding 权衡相关性与多样性"]
    C --> D["轻量化批量推理加速<br/>RFF 投影 + 批次选择 + 早停自适应"]
    D --> E["分簇批量 PFN 前向推理<br/>执行 K 次向量化前向计算输出预测"]

关键设计¶

1. 测试端划分聚类：将离散测试查询聚集成组以赋能批量推理 针对逐样本 $k$NN 检索必须触发 $T$ 次单独前向传播的严重低效性，CRUMB 将聚类操作倒置在测试端。给定测试集 $\mathcal{D}_{\text{test}} = \{\bm{x}_j^*\}_{j=1}^T$，首先进行零均值单位方差标准化，然后利用 $k$-means++ 初始化并运行 Lloyd 算法将测试样本划分为 $K$ 个互不相交的测试簇 $\{C_1, \dots, C_K\}$。这一设计将超参数 $K$ 构造成计算吞吐与上下文特化度之间的可控调节杆：当 $K=1$ 时退化为全测试集共享单一全局上下文；当 $K=T$ 时等价于逐样本独立特化。在实践中取适度的 $K$（如 $K=20$）即可将前向传播次数压缩至原来的 $1/10 \sim 1/250$，使同一簇内的全部查询能拼成张量在 GPU 上进行单次密集矩阵乘法。

2. 贪心 MMD 训练子集检索：显式平衡局部相关性与样本多样性 为特定测试簇 $C_k$ 挑选规模为 $n$ 的训练子集 $\mathcal{S}_k \subset \mathcal{D}_{\text{train}}$ 时，简单按簇中心近邻检索会导致样本在中心过度扎堆并遗漏边缘尾部分布，而 Voronoi 均匀下采样则在密集区产生冗余且在稀疏区遗漏边界。CRUMB 采用最大均值差异（MMD）度量经验分布匹配度，所选子集的目标为： $$ \mathcal{S}{k}^{*}=\argmin{k}\subset\mathcal{D}}},\,|\mathcal{S{k}|=n}\left[\frac{1}{n^{2}}\sum{i},\bm{x}}}\in\mathcal{S{k}}\kappa(\bm{x}},\bm{x{i^{\prime}})-\frac{2}{n|C}|}\sum_{\bm{x{i}\in\mathcal{S}^{}}\sum_{\bm{x}{j}\in C}}\kappa(\bm{x{i},\bm{x}^{})\right] $$ 其中第二项奖励候选训练样本向当前测试簇靠拢（保障语义相关性），第一项惩罚已选子集内部的冗余核相似度（强制保证样本多样性与覆盖广度）。CRUMB 采用核聚集（Kernel Herding）策略进行贪心迭代构造：在第 $t$ 步，从候选未选样本中挑选使下式最小的训练点： $$ \bm{x}{i^{*}}=\argmin{i}\in\mathcal{D}}}\setminus\mathcal{S{k}}\left[-\frac{2}{|C^{}|}\sum_{\bm{x}{j}\in C}}\kappa(\bm{x{i},\bm{x}^{})+\frac{2}{t+1}\sum_{\bm{x}{i^{\prime}}\in\mathcal{S}}}\kappa(\bm{x{i},\bm{x})\right] $$ 核函数采用高斯径向基函数（RBF），带宽 }$\sigma$ 由中位数启发式自适应确定，无需额外超参数调优。值得注意的是，不同测试簇的检索子集 $\mathcal{S}_k$ 允许重叠，高信息量的关键训练样本可在多个相关簇间自然复用。

3. 轻量化批量推理加速：随机傅里叶特征、批次选择与早停自适应机制 在大规模数据集上，朴素高斯核评估开销高达 $O(Nnd)$。CRUMB 引入了三项工程与算法加速。首先是基于随机傅里叶特征（RFF）的显式低维投影映射 $\bm{z}(\bm{x}) = \sqrt{2/D}\cos(\Omega^\top \bm{x} + \bm{b})$，其中 $D=64 \ll N$。通过将核评估转化为线性内积，测试簇与候选样本的亲和项直接简化为单次矩阵-向量乘法 $Z \bar{\bm{z}}_{\mathcal{C}}$，运行冗余和更新仅需 $D$ 维向量累加，使每轮贪心检索代价降至 $O(ND)$。其次，采用批量贪心选择，每轮同时贪心选入得分最高的 $B=50$ 个样本后再更新冗余项，使迭代轮数骤减至 $\lceil n/B \rceil$ 且无精度损失。最后，针对超大数据集引入自适应早停：每累积挑选 $\Delta=50$ 个点检测经验 $\widehat{\mathrm{MMD}}^2$ 的相对降幅，连续 5 次低于 $\varepsilon=10^{-4}$ 时提前终止，使简单测试簇自动采用更小的上下文完成推理。

损失函数 / 训练策略¶

CRUMB 是完全训练无关的测试时推理包装方法，不涉及任何参数更新、梯度下降或损失函数优化。预训练 PFN 权重完全保持冻结。预处理阶段在 CPU 上并行完成特征标准化、测试集聚类与 RFF 贪心 herding，各测试簇构成的张量输入送入单卡 GPU 执行标准的自注意力/交叉注意力前向计算，推理全过程不需要微调适配器。

实验关键数据¶

主实验¶

在涵盖 51 个分类与回归数据集的 TabArena 基准上，以 TabICL-v2 为骨干网络，对比 Full Context、per-query $k$NN、CRUMB、MICP 与 Uniform Subsampling（除全量上下文外，预算均为 $n=0.1N$），10 个随机种子共 510 组配对观测，评估平均排名（数值越低越优）及显著性检验结果。

方法	平均排名 (Avg. Rank)	$p_{\text{adj}}$ (vs. CRUMB)	统计显著	每轮注意力开销	前向计算次数
Full Context	2.202 ± 0.098	$\sim 10^{-19}$	是	$O(N^2)$	1
per-query $k$-NN	2.782 ± 0.107	0.106	否	$O(n^2)$	$T$ (通常 200+)
CRUMB (本文)	2.975 ± 0.096	—	—	$O(n^2)$	$K$ (固定 20)
MICP	3.304 ± 0.105	$\sim 10^{-4}$	是	$O(n^2)$	$K^*$
Uniform	3.737 ± 0.105	$\sim 10^{-18}$	是	$O(n^2)$	1

在 38 个 TabArena 分类数据集、跨 3 种 PFN 架构（TabPFNv2、TabICLv1、TabICLv2）及 5 种采样预算比率 $p \in \{0.05, 0.1, 0.2, 0.5, 0.8\}$ 的全量对比中： - CRUMB 平均排名达到 1.809 ± 0.027； - MICP 平均排名为 2.022 ± 0.023（$p_{\text{adj}} \sim 10^{-8}$，显著弱于 CRUMB）； - Uniform 平均排名为 2.169 ± 0.026（$p_{\text{adj}} \sim 10^{-9}$）。

消融实验¶

在受控协变量漂移实验中，利用真实数据集的第一主成分（PC1）施加漂移强度参数 $\tau \in [0, 1]$（$\tau=0$ 为无漂移，$\tau=1$ 为完全脱离分布的域外测试），跨 3 个数据集、3 种 PFN 模型与 10 个随机种子评估分类准确率：

漂移强度 $\tau$	CRUMB 准确率	MICP 准确率	差值 $\Delta_{\text{CRUMB}-\text{MICP}}$
$\tau = 0.0$ (同分布)	0.794 ± 0.029	0.746 ± 0.033	+0.049 ± 0.006 (+4.9%)
$\tau = 0.5$ (中度偏移)	0.788 ± 0.019	0.663 ± 0.031	+0.125 ± 0.015 (+12.5%)
$\tau = 1.0$ (极度脱离分布)	0.699 ± 0.014	0.528 ± 0.031	+0.171 ± 0.021 (+17.1%)

同时，在 $2 \times 2$ 阶乘设计（聚类 vs 不聚类，均匀抽样 vs MMD 抽样）消融中： - Case A（不聚类 + 均匀）与 Case B（测试聚类 + 全局均匀）表现垫底且性能相仿； - Case C（不聚类 + 全测试集直接 MMD）提升微弱； - 仅有 Case D（CRUMB：测试聚类 + 簇级 MMD）排名显著最佳，证实测试局部切分与分布对齐的协同作用是核心优势来源。

关键发现¶

批量加速不失精度：相比 per-query $k$NN，CRUMB 将平均推理耗时由 361 秒缩减至 35 秒（实现 10 倍以上提速），在 Wilcoxon 检验中与 $k$NN 无统计学显著差异（$p_{\text{adj}} = 0.106$）。
抗漂移断层优势：当协变量偏移加剧到极限（$\tau=1.0$）时，MICP 准确率暴跌 21.8%（由 0.746 降至 0.528），而 CRUMB 仅平缓下降 9.5%（保持在 0.699），领先优势扩大至 +17.1%。
RFF 与批量选择无损性能：$D=64, B=50$ 的组合不仅让 CPU 预处理加速约 10 倍，且其在分类与回归任务中的性能与精确高斯核 MMD 差距在 0.2 个百分点以内。

亮点与洞察¶

测试端聚类破除批量瓶颈：颠覆了传统在训练集做划分或逐样本近邻检索的思维，先按测试查询分布聚类，既能保留局部几何特化，又天然兼容 GPU 批处理。
MMD Herding 实现分布自适应覆盖：不仅选择最靠近测试簇的训练样本，还通过非负排斥项抑制冗余，自动在测试簇边界和中心之间形成最优支撑点分布。
天然的分布漂移追踪器：当流式测试数据面临环境变化时，只需在滑动时间窗内周期性刷新测试簇中心与对应 MMD 上下文，即可无缝应对动态协变量演化。

局限与展望¶

离线测试依赖：算法原生依赖一批测试样本进行整体聚类，单样本低延迟在线流式输入需依赖缓存质心退化机制。
几何各向同性假设：$k$-means 与标准径向基 RBF 核依赖各向同性欧氏距离与固定中位数带宽，难以直接自适应极度高维或异质混合类型的表格特征流形。
受制于基础 PFN 先验表达上限：作为外层封装器，其性能上限仍然受限于底层预训练 PFN 自身的后验拟合能力，当基础模型对目标领域的先验假设失效时，上下文优化无法从根本上消除泛化误差。

评分¶

新颖性: ⭐⭐⭐⭐☆ 颠覆训练端聚类范式，将测试端聚类与 MMD Kernel Herding 精巧结合用于 PFN 上下文批处理选择。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 51 个真实数据集、3 种主流 PFN 架构，深入进行了超大样本早停、漂移模拟及严格的阶乘消融分析。
写作质量: ⭐⭐⭐⭐⭐ 动机清晰直白，数学推导完备，实验设计与图表论证逻辑非常扎实。
价值: ⭐⭐⭐⭐⭐ 彻底解决了 PFN 无法拓展至 5 万到 10 万以上样本规模的实际应用痛点，为表格基础模型落地带来实用高效的通用外壳。

方法	平均排名 (Avg. Rank)	\(p_{\text{adj}}\) (vs. CRUMB)	统计显著	每轮注意力开销	前向计算次数
Full Context	2.202 ± 0.098	\(\sim 10^{-19}\)	是	\(O(N^2)\)	1
per-query \(k\)-NN	2.782 ± 0.107	0.106	否	\(O(n^2)\)	\(T\) (通常 200+)
CRUMB (本文)	2.975 ± 0.096	—	—	\(O(n^2)\)	\(K\) (固定 20)
MICP	3.304 ± 0.105	\(\sim 10^{-4}\)	是	\(O(n^2)\)	\(K^*\)
Uniform	3.737 ± 0.105	\(\sim 10^{-18}\)	是	\(O(n^2)\)	1

漂移强度 \(\tau\)	CRUMB 准确率	MICP 准确率	差值 \(\Delta_{\text{CRUMB}-\text{MICP}}\)
\(\tau = 0.0\) (同分布)	0.794 ± 0.029	0.746 ± 0.033	+0.049 ± 0.006 (+4.9%)
\(\tau = 0.5\) (中度偏移)	0.788 ± 0.019	0.663 ± 0.031	+0.125 ± 0.015 (+12.5%)
\(\tau = 1.0\) (极度脱离分布)	0.699 ± 0.014	0.528 ± 0.031	+0.171 ± 0.021 (+17.1%)