ICMIL — In-Context Multiple Instance Learning

📄 Paper (arXiv:2606.06458) · 💻 Code · 🤗 Datasets

ICMIL is an in-context learner for bag-structured data: pretrained on synthetic MIL tasks, it labels the bags of a new task from a handful of labelled bags in a single forward pass, with no gradient updates, fine-tuning or hyper-parameter search.

This repository holds the three trained model seeds behind the paper's ICMIL (Ours) results. The same checkpoints are committed in the code repository under checkpoints/.

Files

File	Seed name
icmil-c5trd795.pt	c5trd795
icmil-ggwsqibd.pt	ggwsqibd
icmil-k337zhz1.pt	k337zhz1

Seed names are opaque run labels. Each .pt holds {model_state_dict, epoch, ...provenance}. The reported ICMIL row is the mean ± cross-seed SEM over the three seeds.

We also provide three 50-feature seeds — icmil-50dim-ldsebntm.pt, icmil-50dim-sjik4hzb.pt, icmil-50dim-zeqy97nm.pt. No paper results are based on them.

Usage

from icmil import load_icmil  # pip install -e . from https://github.com/injurise/ICMIL

model = load_icmil(seed="c5trd795", device="cuda")
# X_train: (1, n_ctx_bags, bag_size, n_features), y_train: (1, n_ctx_bags)
# X_test:  (1, n_query_bags, bag_size, n_features)
logits = model(X_train, y_train, X_test)   # (1, n_query_bags, n_classes)

load_icmil reads checkpoints/ in the repo by default; point it elsewhere with load_icmil(source="/path/to/ckpts", seed=...) or ICMIL_CKPT_DIR. Reproduce the full benchmark table with python -m icmil.reproduce.

Citation

@article{mollers2026incontext,
  title   = {In-Context Multiple Instance Learning},
  author  = {M\"ollers, Alexander and Sextro, Marvin and Hense, Julius and Dernbach, Gabriel and M\"uller, Klaus-Robert},
  journal = {arXiv preprint arXiv:2606.06458},
  year    = {2026}
}