← 返回 2026-06-10

上下文多实例学习 In-Context Multiple Instance Learning

Alexander Möllers, Marvin Sextro, Julius Hense, Gabriel Dernbach, Klaus-Robert Müller 📅 2026-06-04 👍 0 2026-07-13 08:37
Prior-data Fitted Networks 上下文学习 元学习 多实例学习 少样本学习

通过合成数据预训练实现少量标签下的零梯度更新多实例分类

前置知识

多实例学习 (MIL)

多实例学习是一种监督学习框架,其中标签仅在bag(实例集合)级别可用,而不是单个实例。在传统MIL中,每个bag B = {x₁, ..., x_I} 包含I个实例,只有一个bag级标签y。经典假设是bag标签通过实例标签的最大值决定:y = max_i{y_i},即bag为正当且仅当其中至少有一个正实例(称为witness)。应用场景包括计算病理学(组织切片作为bag,图像patch作为实例)、药物活性预测、卫星图像等。

本文核心贡献是为MIL引入上下文学习范式,理解传统MIL的bag级监督和实例聚合机制是读懂本文的基础。

Prior-data Fitted Networks (PFNs)

PFN是一种通过在合成数据上预训练来摊销贝叶斯推理的范式。其思想是先定义一个先验p(D)覆盖可能的数据生成过程,从中采样大量合成数据集来训练Transformer。训练时模型学习从上下文样本预测查询点标签。推理时,给定真实数据集作为上下文,模型通过单次前向传播输出预测,无需梯度更新。性能取决于先验与目标任务的匹配程度。

ICMIL是PFN在bag结构数据上的首次直接应用,理解PFN的合成先验设计和推理机制是理解本文方法的关键。

Perceiver架构

Perceiver是一种通用的Transformer变体,设计用于处理任意模态的输入。核心思想使用可学习的latent token通过迭代交叉注意力从输入数据中提取信息,然后通过自注意力在latent之间交换信息。这种设计将计算复杂度从输入大小的平方降低到latent数量的平方。在本文中,Perceiver被改造用于处理分层的bag结构:每个bag有可学习的bag token,通过交叉注意力聚集实例信息,通过自注意力在bag之间交换信息。

本文提出的Perceiver风格架构是解决MIL上下文学习中计算可扩展性、任务相关实例压缩、置换不变性三大挑战的关键创新。

结构因果模型 (SCM)

结构因果模型用有向无环图描述变量之间的因果关系。在本文中,MLP-SCM是一个随机参数化的非线性SCM,用于同时生成实例特征x和实例级输出f(x)。因果图中的特征节点到指定输出节点的映射定义了实例变换f。通过采样不同的SCM参数,可以生成多样化的合成数据集,覆盖不同的假设空间。

本文的合成先验基于MLP-SCM,理解SCM如何同时生成特征和标签有助于理解factorized和joint两类先验的本质区别。

研究动机

多实例学习在实际应用中面临严重的标签稀缺问题。现实世界的MIL数据集通常只有几十个标记bag,例如论文中提到的基准测试大多只有约100个bag。在这种低样本情况下,灵活的模型如基于注意力的MIL和TransMIL容易过拟合,而限制性强的模型编码的归纳偏差可能与任务不匹配。此外,监督MIL预训练需要大量相似的标记源语料,但许多应用缺乏此类资源。自监督方法虽然不需要预训练标签,但必须在不知道下游任务的情况下将实例信息压缩到固定表示,可能丢弃特定问题所需的关键特征。

本文的目标是本文的目标是引入In-Context Multiple Instance Learning (ICMIL),一种无需梯度更新或超参数调优,能够在推理时从少量标记bag中学习新任务的框架。通过在多样化的合成bag结构数据上预训练模型,使其学习合理的bag分类器形式,在测试时能够从标注的上下文推断标签规则并用于预测未见bag。

与已有工作不同的是,现有方法要么需要针对每个任务从头训练(如MeanLogReg、SVM、ABMIL),要么依赖于领域特定的预训练(如自监督预训练、迁移学习)。本文的独特切入角度是将上下文学习范式直接应用于bag结构数据,通过精心设计的合成先验和Perceiver风格架构,实现真正的零推理时学习。这与Kopp等人的工作最接近,但他们通过k-means pooling将bag压缩为表格输入,聚合是无监督且推理前固定的,而本文直接在bag结构输入上训练PFN,可以使用MIL标签进行上下文学习。

核心方法

ICMIL基于Prior-data Fitted Network (PFN)范式,在合成bag结构数据上预训练一个Perceiver风格模型来近似bag标签的后验预测分布。训练目标是学习一个查询函数qθ,使其对从先验p(D)采样的数据集D = {(B_i, y_i)}_{i=1}^{N+1}最小化期望负对数似然:L(θ) = E_{D∼p(D)}[-log qθ(y_{N+1} | B_{N+1}, {(B_i, y_i)}_{i=1}^N)]。在推理时,模型接收真实MIL数据集作为上下文,在单次前向传播中分类测试bag,无需梯度更新或任务特定训练。模型通过跨bag、跨实例和跨特征的注意力机制实现任务特定的实例聚合和有效的少样本推理。

ICMIL的核心创新在于两个方面:一是Perceiver风格架构解决bag结构数据的三大挑战(计算可扩展性、任务相关实例压缩、置换不变性),二是两类互补的合成先验设计。factorized priors遵循传统MIL假设,bag标签通过置换不变函数分解:y = g(ϕ{f(x_i)}_{i=1}^I),其中f是实例级变换,ϕ是聚合,g是bag到标签映射。joint priors放弃分解,直接定义整个bag的函数:y = f_S(x_1, ..., x_I)。factorized priors使用MLP-SCM为每个实例独立采样潜在原因,产生不相关的bag;joint priors为每个bag采样一次潜在原因,产生bag内特征相关性。实验表明factorized priors在不相关的witness任务上表现更好,joint priors在相关性和交互任务上更强。

方法步骤详情

ICMIL的训练和推理分为三个主要步骤。步骤1是嵌入:将F个输入特征按组大小s分组,产生G = ⌈F/s⌉个特征组嵌入,维度为E,形成实例嵌入张量X̃ ∈ R^{N×I×G×E}。bag标签被嵌入为ỹ ∈ R^{N×E},测试bag接收训练标签嵌入的均值。步骤2是迭代更新:为每个bag初始化可学习token Q_b ∈ R^{G×E},进行T步迭代。每步首先通过实例聚合更新:Z_{b,g}^{(t)} = Q_{b,g}^{(t)} + MHA(Q_{b,g}^{(t)}, X̃_{b,g}, X̃_{b,g}),独立对每个bag应用。然后将标签嵌入ỹ_b连接到Z_b^{(t)},通过inter-bag attention更新:Q^{(t+1)} = ColRowAttn(Ẑ^{(t)}),其中列注意力在特征组和标签上操作,行注意力允许bag之间信息流动。步骤3是解码:T步迭代后,查询bag的标签token通过解码器产生预测的类别分布。在推理时,整个过程在单次前向传播中完成,无需梯度更新。

技术新颖性

ICMIL的技术新颖性体现在:一是首次将PFN范式直接应用于bag结构数据,而非先压缩成表格形式;二是提出的Perceiver风格架构通过迭代交叉注意力和自注意力解决MIL上下文学习的三大挑战——计算可扩展性(通过分块实例聚合将峰值内存从O(N·I·G)降到O(I·G))、任务相关实例压缩(模型不需要在单次标签无关步骤压缩实例,而是交替处理实例和跨bag注意力)、置换不变性(MHA的置换不变性确保bag内实例顺序不变,同时保持bag间区分);三是两类互补的合成先验设计揭示了结构对齐与表达力的权衡,factorized priors与经典MIL假设对齐,joint priors更表达并能编码bag内相关性;四是先验混合策略使单一模型继承不同先验的优势,在12个基准测试上达到最佳平均性能。

Schematic of the ICMIL architecture, illustrated for T iterations.
Figure 2: Schematic of the ICMIL architecture, illustrated for T iterations.
Taxonomy of bag-structured priors for ICMIL.
Figure 3: Taxonomy of bag-structured priors for ICMIL.

实验结果

实验在12个MIL基准测试上评估ICMIL,包括6个witness标签规则任务(SMIL、Musk1、Musk2、Letters、HEPMASS、RSNA-ICH)和6个跨实例交互任务(Elephant、Fox、Tiger、TCGA、Adjacent Pairs、Pos/Neg)。消融实验表明不同先验编码了互补的归纳偏差。Joint先验达到最佳平均排名(2.60),在bag内相关性和/或交互存在的设置中表现最强。但factorized priors在不相关的witness基准测试上大幅超过它:Factorized(disc, lookup)在Letters上达到95.2±0.5(比Joint高约4个百分点),在SMIL上74.9±0.3,在HEPMASS上86.1±1.0;Factorized(continuous, MLP)在Musk1上93.9±0.6,在Musk2上85.4±3.2。先验混合模型(Mixed)保持Joint先验在相关性和交互驱动基准测试上的领先优势(如Pos/Neg 87.3±0.2,TCGA 88.3±0.2),同时恢复了factorized priors在不相关witness任务上的差距(Letters 94.4±1.6 vs 95.2±0.5;HEPMASS 87.2±0.5,整体最佳)。进一步扩展的ICMIL模型(加倍训练步数到40000,增加嵌入维度到256,MLP隐藏大小到1054)在12个基准测试上达到最佳平均AUROC 84.17和最佳平均排名4.08。提升主要由Fox (+6.8)、Musk2 (+3.2)以及较小程度的Letters (+1.3)和Musk1 (+0.8)驱动。ICMIL避免了ABMIL因少量训练bag导致的超参数选择和早停不可靠,且总运行时与简单线性和核基线竞争。但ICMIL并非在所有基准测试上占优,在SMIL、TCGA和Letters上基线仍有明显优势(如ABMIL在SMIL上+11.0,TabPFN-Cluster在Letters上+3.5)。

Per-benchmark mean AUROC (%) ± standard error for each model.
Table 1: Per-benchmark mean AUROC (%) ± standard error for each model.
Left: Mean rank across MIL benchmarks for models trained on joint and factorized priors, broken down by group and pooled over all datasets (Overall). Right: Within-bag feature correlation (ICC) per benchmark and colored by task groups.
Figure 4: Left: Mean rank across MIL benchmarks for models trained on joint and factorized priors, broken down by group and pooled over all datasets (Overall). Right: Within-bag feature correlation (ICC) per benchmark and colored by task groups.
Total wall-clock time on the Pos/Neg benchmark for ICMIL and the three supervised MIL baselines.
Figure 5: Total wall-clock time on the Pos/Neg benchmark for ICMIL and the three supervised MIL baselines.
Learning curves on Fox (left) and Musk2 (right), mean ± SE across three runs.
Figure 6: Learning curves on Fox (left) and Musk2 (right), mean ± SE across three runs.
查看结构化数据
任务指标本文基线提升
Musk1 AUROC (%) 93.3±0.6 MeanLogReg 91.6±0.5 +1.7
Musk2 AUROC (%) 90.8±1.1 MeanLogReg 82.1±1.6 +8.7
TCGA AUROC (%) 88.9±0.6 ABMIL 90.7±1.6 -1.8 (基线更好)
Letters AUROC (%) 95.7±0.7 ABMIL 97.6±0.8 -1.9 (基线更好)
Pos/Neg AUROC (%) 87.0±0.2 MeanLogReg 85.4±0.5 +1.6
Elephant AUROC (%) 74.4±0.3 MeanLogReg 73.9±0.1 +0.5
Fox AUROC (%) 77.6±0.5 MeanLogReg 74.6±0.0 +3.0
SMIL AUROC (%) 74.5±0.5 ABMIL 85.4±0.5 -10.9 (基线更好)
Overall Mean AUROC (%) 84.17 MeanLogReg 82.37 +1.80
Overall Mean Rank 4.08 MeanLogReg 6.33 +2.25 (越低越好)

局限与改进

作者承认的局限性包括:训练课程当前使用的bag大小最多20个实例,虽然模型在推理时能泛化到更大bag(如在Musk数据集上的强性能所示),但在训练中包含更大bag是自然的下一步。基准测试专注于通过PCA降维到25维的特征的二分类,多类目标和更高维的基础模型嵌入是有前景的方向。扩展实验表明,增加容量和训练时间的收益在基准测试上不完全均匀,先验混合、模型大小和训练持续时间之间存在值得进一步研究的交互。此外,作者观察到在HEPMASS上扩展模型有2.7个百分点的回归,说明不同任务对模型规模的响应不同。我自己观察的局限性是:基线方法在部分任务上仍明显优于ICMIL,说明合成先验与真实任务分布间存在gap;PCA降维到25维可能丢失有用信息,限制了在真实高维嵌入(如UNI2的1536维)上的性能;训练数据大约256万个合成数据集(扩展模型512万),虽然数量庞大,但先验设计仍可能无法覆盖某些真实MIL任务的复杂模式。

独立分析的弱点

ICMIL的主要弱点体现在三个方面。首先是合成到真实的分布差距:factorized priors基于实例独立性假设,而真实MIL数据(如组织病理学图像)常有空间依赖性;joint priors虽然能编码bag内相关性,但MLP-SCM生成的依赖模式可能与真实病理或卫星图像的依赖结构不同。改进方向可以设计领域特定的先验,如计算病理学的空间先验或时间分辨任务的序列感知先验。其次是特征维度限制:当前通过PCA降维到25维,可能丢弃高维嵌入(如UNI2的1536维)的判别信息。改进方向是直接训练更高维的合成数据,同时扩展模型容量和训练时间。第三是部分基准测试上基线优势明显:ABMIL在SMIL上+11.0,TabPFN-Cluster在Letters上+3.5,说明ICMIL在某些任务上仍欠优化。改进方向包括设计更丰富的先验、更长预训练、或在真实MIL语料上进行后训练以缩小合成到真实gap。另一个独立观察的弱点是训练课程限制在最大20个实例的bag,而某些基准测试(如Musk2有最多1044个实例)可能需要更大bag的先验支持。

未来方向

作者提出的未来研究方向包括:设计更大bag的训练课程,在更高维合成数据上训练以支持多类目标和高维基础模型嵌入。基于成果可延伸的方向包括:领域特定先验设计,如计算病理学的空间先验(模拟组织切片的空间连续性)、时间序列任务的序列感知先验(类似时间序列预测PFNs);在真实MIL语料上进行后训练以进一步缩小合成到真实gap;探索更丰富的架构变体,如引入跨bag的显式实例依赖建模、或使用更复杂的解码器;研究先验混合、模型大小和训练持续时间之间的交互,优化扩展策略;将ICMIL扩展到回归和时间到事件目标;探索零样本应用,如利用视觉语言模型提供bag描述作为额外上下文。另一个有趣的方向是研究ICMIL的可解释性,通过注意力权重分析模型学到的实例聚合模式。

复现评估

论文提供开源代码:https://github.com/injurise/ICMIL。训练基于PyTorch实现,使用hydra-zen配置。合成先验构建自TabICL的MLP-SCM和Tree-SCM生成器。ICMIL Perceiver风格架构和训练循环基于nanoTabPFN参考实现。基线包括MIL-Lab的ABMIL实现、TabPFN-v2模型、以及scikit-learn的经典方法。基准测试加载器(RSNA-ICH)来自torchmil。实验在高性能集群上运行,ICMIL预训练在单块NVIDIA A100 80GB显卡上执行。减少设置(用于先验消融)每次运行约12小时,扩展ICMIL模型约24小时。评估在集群较弱的GPU上执行。合成数据从预生成池中采样,20000步批大小128的训练中模型接触约256万个独立采样的合成数据集,扩展ICMIL模型加倍到512万。每个batch在同质先验配置上。复现难度中等,主要依赖是单块A100 GPU,但代码和配置已开源,实验细节在附录B完整描述。