Measure Less, Know More: Self-Supervised Test-Time Feature Acquisition

[Paper] | [OpenReview] | [Cite]

Authors: Eeshaan Jain, Linus Bleistein, Bart Deplancke, Charlotte Bunne

Research code for ECHO-k (NeurIPS 2026).

Abstract

Recent progress in multimodal, high-dimensional learning has enabled foundation models to process heterogeneous, large-scale data. However, at test time, acquiring all features or modalities can be prohibitively costly and often redundant. Sequentially selecting informative modalities is therefore critical, yet challenging when the downstream task or prediction target is unknown. To this end, we introduce ECHO-k, a task-agnostic and self-supervised learning principle for modality acquisition: we use a deep model's internal pretrained representations (e.g., from a foundation model) as proxy targets that summarize cross-modal information. We provide theoretical guarantees in a stylized linear setting that motivate a reinforcement learning (RL) policy for sequential modality selection. Across task-agnostic and label-free acquisition baselines, ECHO-k consistently improves budgeted downstream performance across diverse foundation-model backends. Our method provides a principled route to cost-aware test-time deployment, with implications for any multimodal system where measurements are expensive or time-constrained, and downstream tasks unknown a priori.

Installation

Clone the repository with its VirTues and spora-io submodules:

git clone --recurse-submodules https://github.com/bunnelab/echo-k.git
cd echo-k
python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt

For multiplex imaging, install the VirTues dependencies and download the pretrained VirTues backbone:

bash virtues_acquisition/setup.sh
bash virtues_acquisition/run.sh download-model

Usage

1. Prepare your data

Follow the spora documentation to download a dataset or prepare your own.

2. Configure ECHO-k

Refer to virtues_acquisition/configs/spora.yaml for dataset paths, model settings and acquisition parameters.

3. Train and run acquisition

# Verify the setup by embedding one tissue from each split
bash virtues_acquisition/run.sh check config=virtues_acquisition/configs/spora.yaml

# Train the acquisition policy (the VirTues backbone stays frozen)
bash virtues_acquisition/run.sh train config=virtues_acquisition/configs/spora.yaml

# Acquire markers for each test tissue with the best validation checkpoint
bash virtues_acquisition/run.sh rollout config=virtues_acquisition/configs/spora.yaml

4. Use the acquired representations

import torch

result = torch.load("<run_dir>/rollouts/ppo_test.pt", map_location="cpu", weights_only=True)

k = 5  # acquisition budget
for tissue_id, tissue in result["tissues"].items():
    # Markers in acquisition order, starting with the initial marker
    markers = [result["markers"][i] for i in tissue["order"]]
    # Tissue embedding after acquiring the first k markers
    embedding = tissue["embeddings"][k - 1]

For the MNIST, Fashion-MNIST and MiniBooNE examples, see the benchmark guide.

Citation

If you find our work useful in your research or if you use parts of this code, please consider citing our paper:

@inproceedings{jain2026measure,
  title = {{Measure Less, Know More: Self-Supervised Test-Time Feature Acquisition}},
  author = {Jain, Eeshaan and Bleistein, Linus and Deplancke, Bart and Bunne, Charlotte},
  booktitle = {The Fortieth Annual Conference on Neural Information Processing Systems (NeurIPS)},
  year = {2026},
  url = {https://openreview.net/forum?id=bXoyojlawt}
}

Acknowledgements

This work builds on VirTues and spora-io. Policies are trained with Stable-Baselines3 and sb3-contrib.

License and Terms of Use

Copyright (c) ECOLE POLYTECHNIQUE FEDERALE DE LAUSANNE, Switzerland, Laboratory of Artificial Intelligence in Molecular Medicine, 2026

This repository is released under the MIT License; see LICENSE. VirTues, spora-io and the VirTues model weights have their own licenses.

Contact

For questions, please open an issue.