PaperlayerView exact revision on arXiv

Source-linked AI summary

In-Context Multiple Instance Learning

Alexander Möllers, Marvin Sextro, Julius Hense, Gabriel Dernbach, Klaus-Robert Müller

arXiv:2606.06458v1cs.LGcs.AIcs.CV

TL;DR

MIL lacks robust solutions for learning from few labeled bags because flexible models can overfit while rigid models may not fit the task. ICMIL pretrains a Perceiver-style in-context learner on synthetic bag-structured tasks and predicts new labels in one forward pass, achieving the best average performance across twelve MIL benchmarks. The approach’s scope remains centered on binary classification with PCA-reduced features and training bag sizes up to 20 instances.

  • Problem

    Existing MIL methods struggle in the low-label regime, where flexible models may overfit and restrictive models may encode task-mismatched inductive biases.

  • Method

    ICMIL pretrains a Perceiver-style in-context learner on synthetic bag-structured data generated from complementary priors, then classifies new tasks without gradient updates or task-specific training.

  • Results

    ICMIL achieves the best average AUROC and average rank across twelve MIL benchmarks in the low-label regime, outperforming supervised baselines.

  • Takeaways & Limitations

    A mixture of synthetic priors can provide a single in-context MIL model with complementary per-task strengths across diverse low-label benchmarks.

  • Takeaways & Limitations

    The benchmarks focus on binary classification with features reduced to 25 dimensions via PCA, while training uses bag sizes of up to 20 instances.

Abstract

from arXiv · show

Multiple Instance Learning (MIL) addresses problems where supervision is available at the level of bags of instances and has been successfully applied in fields ranging from computational pathology to satellite imagery. Nevertheless, existing algorithms struggle in the low-label regime that characterizes many real-world applications. Flexible models overfit and rigid ones fail to adapt to the task at hand. We show that pretraining an in-context learner with a Perceiver-style architecture on synthetic data yields a model that can solve new tasks from a handful of labeled bags. At inference time, classification happens in a single forward pass and requires no gradient updates. We propose and investigate different synthetic data generators for bag-structured data and find that they capture complementary inductive biases. A model pretrained on a mixture of these generators inherits their per-task strengths and achieves the best average performance across twelve MIL benchmarks, outperforming supervised baselines that require task-specific training.

1 Introduction

ICMIL addresses low-label MIL by learning from synthetic bag-structured tasks and predicting new task labels in one forward pass. Its Perceiver-style architecture and mixture of complementary synthetic priors achieve strong average performance across twelve benchmarks.

  • Motivation: MIL datasets often contain only a few dozen labeled bags, making flexible aggregators prone to overfitting and restrictive models mismatched to the task.This creates a low-label setting in which standard supervised training is fragile.
  • Results: On MNIST-Pos/Neg, ICMIL has the highest median AUROC and lowest variance across 20 resampled training sets, especially with few bags.The result illustrates reduced sensitivity to training-set resampling in the low-bag regime.
  • Approach: ICMIL adapts the PFN paradigm by pretraining on diverse synthetic bag-structured tasks, then inferring the labeling rule from labeled context at inference time.The model predicts labels for unseen bags in a single forward pass without gradient updates or task-specific training.
  • Architecture: A Perceiver-style architecture addresses hierarchical-set challenges involving scalability, task-aware compression, and permutation invariance within bags.The design must preserve bag identity while incorporating bag labels into instance compression.
  • Synthetic priors: Factorized and joint synthetic generators encode different MIL assumptions and produce models with complementary strengths across downstream tasks.Factorized priors retain classical instance-level decomposition, whereas joint priors allow richer dependence structures among instances.
  • Results: ICMIL achieves the best average AUROC and rank across twelve MIL benchmarks in the low-label regime, outperforming supervised baselines without gradient updates or hyperparameter tuning.The reported advantage is measured across the benchmark suite rather than a single downstream task.

2 Background

MIL represents each example as a labeled bag of instances, while PFNs amortize Bayesian inference over synthetic data-generating processes. Prior work applying PFNs to MIL used fixed unsupervised aggregation, whereas ICMIL learns directly from bag-structured inputs and in-context labels.

  • Multiple Instance Learning: MIL assigns one label to a variable-size bag of instances, supporting classification, regression, or time-to-event targets.The original formulation assumes within-bag statistical independence and permutation invariance, with the bag label determined by the maximum instance label.
  • Multiple Instance Learning: Traditional MIL aggregates instance predictions with predefined functions such as max or mean pooling, while attention-based models learn task-relevant instance aggregation.Extensions also model dependencies among instances, broadening the class of supported bag-label relationships.
  • Prior-data Fitted Networks: PFNs amortize Bayesian inference by pretraining on synthetic datasets sampled from priors over data-generating processes and predicting held-out labels from context.This allows inference through the pretrained model rather than fitting a new predictor for each task.
  • Related work: TabPFN-based multi-instance regression collapses each bag through fixed, unsupervised k-means pooling before prediction.ICMIL instead trains directly on bag-structured inputs with priors and an architecture that can use MIL labels in context.

3 In-Context Learning for Multiple Instance Problems

ICMIL formulates MIL prediction as in-context inference over labeled bags and uses a Perceiver-style architecture to process hierarchical set inputs efficiently. It also compares factorized and joint synthetic priors, which encode different assumptions about instance independence and within-bag correlations.

  • In-Context Learning: The learner receives labeled context bags and predicts a query bag’s label under a prior over MIL datasets.
  • Architectural Challenges: Naive attention over N bags with I instances each costs O(N^2I^2), while MIL also requires task-aware compression and within-bag permutation invariance.These requirements arise because bags may contain hundreds or thousands of instances, and labels can determine which instance features matter.
  • ICMIL Architecture: The Perceiver-style architecture embeds instances into feature groups, updates learnable bag tokens through within-bag cross-attention, and then applies inter-bag column-row attention with label embeddings.Bag-wise processing reduces peak memory for instance aggregation from O(N·I·G) to O(I·G).
  • ICMIL Architecture: After T iterations, query-bag label tokens are decoded into class distributions; alternating instance and inter-bag attention enables task-dependent compression while preserving permutation invariance and bag identity.Inter-bag attention uses O(N^2 + G^2) compute and linear O(N·G) memory with flash attention.
  • Synthetic Priors: Factorized priors use instance-level transforms, permutation-invariant aggregation, and bag-level transforms, whereas joint priors define one function over flattened bags.Factorized priors sample latent causes independently per instance; joint priors sample them once per bag and therefore express inter-instance correlations.
  • Synthetic Priors: The two prior families trade structural alignment with classical MIL assumptions against greater expressiveness through within-bag feature correlations.The paper empirically investigates whether alignment or expressiveness better supports downstream performance.

4 Experiments

Across twelve low-label MIL benchmarks, the experiments test complementary synthetic priors, prior mixtures, scaling, and comparisons with supervised baselines. ICMIL achieves the strongest average performance while avoiding task-specific gradient updates, though it remains weaker on some individual benchmarks.

  • Experimental Setup: The study evaluates in-context learning on twelve MIL benchmarks spanning witness and interaction label rules, using synthetic datasets and low-sample tasks of approximately 100 bags.The experiments compare prior configurations and heterogeneous supervised baselines under the low-label regime.
  • Different Priors Capture Different MIL Regimes: Joint priors achieve the best average rank and favor correlated or interaction-driven settings, whereas factorized priors excel on uncorrelated witness benchmarks.The Factorized(disc, lookup) prior reaches 95.2 ± 0.5 on Letters, 74.9 ± 0.3 on SMIL, and 86.1 ± 1.0 on HEPMASS, while Factorized(continuous, MLP) leads on Musk1 and Musk2.
  • Mixing Priors Yields a Robust Generalist: A mixed prior combining Joint, Factorized(cont, MLP), and Factorized(disc, lookup) retains strength on correlated and interaction benchmarks while recovering performance on uncorrelated witness tasks.Mixed scores 87.3 ± 0.2 on Pos/Neg, 88.3 ± 0.2 on TCGA, 94.4 ± 1.6 on Letters, and 87.2 ± 0.5 on HEPMASS.
  • Scaling Yields Further Improvement on Selected Benchmarks: ICMIL achieves the best average AUROC (84.17) and mean rank (3.62) across the twelve benchmarks.Scaling improves Fox by 6.8 points, Musk2 by 3.2, Letters by 1.3, and Musk1 by 0.8, but causes a 2.7-point regression on HEPMASS.
  • Scaling Yields Further Improvement on Selected Benchmarks: ICMIL outperforms the strongest aggregate supervised baseline, MeanLogReg at 82.37 average AUROC, while no single baseline is competitive across all benchmarks.ABMIL is strongest on some datasets but underperforms on others, and flat TabPFN baselines lose information by flattening MIL inputs.
  • Scaling Yields Further Improvement on Selected Benchmarks: ICMIL avoids cross-validation, refitting, and gradient updates, making its classification runtime competitive with simpler linear and kernel baselines.ABMIL’s per-split cross-validation and refitting dominate total wall-clock time, whereas ICMIL classifies query bags in one forward pass.

5 Limitations

The setup has several scope boundaries and open questions. Training uses bags of up to 20 instances, while benchmarks focus on binary classification with 25-dimensional PCA features.

  • Training uses bag sizes of up to 20 instances, although the model generalizes well to larger bags at inference time.Including larger bags during training is identified as a natural next step.
  • The benchmarks focus on binary classification with feature dimensionality reduced to 25 via PCA.Multi-class targets and higher-dimensional foundation-model embeddings remain promising directions.
  • Scaling gains are not fully uniform across benchmarks, suggesting interactions between prior mixture, model size, and training duration.These interactions warrant closer study.

6 Conclusion

ICMIL brings in-context learning to bag-structured MIL through a specialized architecture and complementary synthetic priors. The mixed-prior model achieves the strongest average performance across twelve low-label MIL benchmarks, while domain-aligned priors and real-world post-training remain future directions.

  • 6 Conclusion: ICMIL addresses scalability, task-dependent compression, and permutation invariance for hierarchical set inputs.Its architecture is paired with synthetic priors encoding complementary inductive biases over MIL tasks.
  • 6 Conclusion: A model trained on a mixture of synthetic priors achieves the best average AUROC and average rank across twelve MIL benchmarks in the low-label regime.
  • 6 Conclusion: Synthetic data-generator choice matters considerably, and combining priors with complementary inductive biases can improve robustness across MIL regimes.
  • 6 Conclusion: Domain-aligned priors, such as spatial priors for computational pathology or sequence-aware priors for time-resolved tasks, are proposed as future directions.
  • 6 Conclusion: Post-training on real-world MIL corpora could further close the synthetic-to-real gap, while richer architectural variants remain to be explored.

A Learning Curves

Learning curves on Fox and Musk2 show performance continuing to improve with longer training. This motivates the extended training schedule used for ICMIL.

  • A Learning Curves: Performance keeps improving with longer training on Fox and Musk2.The curves report mean ± SE across three runs.

B Training Details

Training exposes ICMIL to millions of unique synthetic datasets using a staged curriculum over bag size and instance classes. The scaled model increases embedding and MLP capacity, uses longer training, and applies a lower learning rate with warmup.

  • Synthetic datasets: Across 20,000 steps, the model sees approximately 2.56M independently sampled synthetic datasets, doubling to 5.12M for scaled ICMIL.No synthetic dataset is repeated during training.
  • Curriculum: A three-stage curriculum expands bag sizes from [8] to and instance classes from at most 6 to at most 20.The final stage runs from steps 7,500–20,000.
  • Architecture: The scaled model increases embedding dimension E from 128 to 256 and MLP hidden size from 512 to 1054, while training for 40,000 steps.The number of iterations and attention heads remain unchanged.
  • Optimization: Scaled ICMIL uses learning rate 5×10^-4 with 2,500 warmup steps, whereas the reduced setup uses 1×10^-3 without warmup.Optimization uses schedule-free AdamW.

B.1 Implementation

The experiments use PyTorch-based implementations, synthetic-prior components adapted from TabICL, and a Perceiver-style ICMIL training stack. Training and evaluation run on a high-performance cluster, with code released publicly.

  • PyTorch implements the experiments, with configuration handled through hydra-zen.
  • ICMIL synthetic priors build on TabICL’s MLP-SCM and Tree-SCM generators, hyperparameter samplers, and regression-to-classification transform.
  • The ICMIL Perceiver-style architecture and training loop build on the nanoTabPFN reference implementation and use schedule-free AdamW optimization.
  • Training uses one NVIDIA A100 with 80 GB of VRAM; reduced prior-ablation models take approximately 12 h per run, while the scaled model takes roughly 24 h.
  • Supervised baselines are evaluated on less powerful GPUs in the same cluster, and the authors make their code publicly available.

C Benchmark Details

The benchmark suite contains twelve MIL tasks spanning synthetic and real-world settings, with standardized bag interfaces, split protocols, and preprocessing. Tasks vary in bag semantics, instance representations, class construction, and evaluation design.

  • Benchmark composition: The suite includes twelve benchmarks: Musk1, Musk2, Elephant, Fox, Tiger, SMIL, Pos/Neg, Adjacent Pairs, Letters, HEPMASS, RSNA-ICH, and TCGA.
  • Synthetic benchmarks: MNIST SMIL samples 100 variable-size bags from digit classes and labels them using a uniformly sampled witness digit.
  • Synthetic benchmarks: MNIST PosNeg labels a bag positive when positive-digit evidence outnumbers negative-digit evidence, so a single witness is insufficient.
  • Synthetic benchmarks: MNIST AdjPairs labels bags positive when any sampled adjacent digit pair is jointly present, with bag sizes sampled from 10 to 20.
  • Classical natural-MIL benchmarks: Musk1 and Musk2 classify molecules from conformation instances, using zero-padding, PCA to d=25, and 5-fold stratified cross-validation on bags.
  • Witness-rate benchmarks: Letters and HEPMASS construct 100 bags of size 10 with two witnesses per positive bag, balanced classes, and stratified train/test splits.
  • Natural-image benchmarks: Elephant, Fox, and Tiger use image bags with segmented-region instances and 230 color-texture features, evaluated with 5-fold stratified cross-validation.
  • Medical-imaging benchmarks: RSNA-ICH treats each CT scan as a bag of axial slices, samples 50 train/test draws, and varies the retained slice count uniformly between 15 and 30.
Loading 2606.06458v1…