✨ No More K-means: Single-Stage Sparse Coding for Efficient Multi-Vector Retrieval ✨

ICML 2026

Lixuan Guo^*1,2, Yifei Wang^*3, Tiansheng Wen^4,1, Aosong Feng⁵, Stefanie Jegelka^6,7, Chenyu You¹

¹Stony Brook University   ²Xidian University   ³Amazon AGI SF Lab  
⁴Georgia Tech   ⁵Yale University   ⁶TUM   ⁷MIT

🚀 🚀 News

• 2026.06 Code released.
• 2026.05 Accepted by ICML2026.

Setup

Install the project dependencies in the environment you use for training and retrieval. pyproject.toml is the single dependency source; use the data extra when you need preprocessing and the eval extra when you need accelerated retrieval kernels:

cd /home/ubuntu/Project/SSR
pip install -e ".[data,eval]"

Data Preprocessing

Training data is prepared from MS MARCO passage:

python prepare_msmarco.py --subset passage

This writes data/processed/msmarco/passage/ with corpus.jsonl, queries/, qrels/, pairs/, and hard_negatives/.

Evaluation data can be prepared from MS MARCO passage/document or from MTEB/BEIR-style retrieval tasks. For MS MARCO evaluation-only data, generate only the index files:

python prepare_msmarco.py --subset passage --skip-pairs --skip-hard-negatives
python prepare_msmarco.py --subset document --skip-pairs --skip-hard-negatives

For MTEB/BEIR evaluation data:

python prepare_mteb_eval.py --datasets nfcorpus scifact hotpotqa

This writes one directory per dataset under data/processed/mteb/, each with corpus.jsonl, queries/{split}.tsv, and qrels/{split}.tsv. More details are available in docs.

Training

Train the standard token SAE for SSR:

python -m ssr.train \
  --dataset msmarco-passage \
  --data-dir data/processed/msmarco/passage \
  --sae-token-scope non-cls \
  --output-dir output/ssr-token

Train a separate [CLS] SAE for SSR-CLS:

python -m ssr.train \
  --dataset msmarco-passage \
  --data-dir data/processed/msmarco/passage \
  --sae-token-scope cls \
  --output-dir output/ssr-cls

The resulting token SAE checkpoint is used as --model-path; the [CLS] SAE checkpoint is passed to evaluation with --cls-sae-path. More details are available in docs.

Evaluation

For a small direct retrieval run, use the sparse-cache backend:

python -m ssr.retrieval.eval_mteb \
  --task retrieval \
  --variant ssr \
  --model-path output/ssr-token/final \
  --dataset nfcorpus \
  --data-dir data/processed/mteb \
  --mode exact \
  --score-device index

For SSR-CLS, add the CLS SAE:

python -m ssr.retrieval.eval_mteb \
  --task retrieval \
  --variant ssr-cls \
  --model-path output/ssr-token/final \
  --cls-sae-path output/ssr-cls/final \
  --dataset nfcorpus \
  --data-dir data/processed/mteb \
  --score-device index

For larger corpora, build an end-to-end index first:

python -m ssr.retrieval.eval_mteb \
  --task index \
  --variant ssr \
  --model-path output/ssr-token/final \
  --dataset nfcorpus \
  --data-dir data/processed/mteb \
  --index-cache-dir data/cache/mteb_index_e2e/nfcorpus \
  --encode-device cuda:0

Then retrieve from the index:

python -m ssr.retrieval.eval_mteb \
  --task retrieval \
  --corpus-backend e2e-index \
  --variant ssr++ \
  --model-path output/ssr-token/final \
  --dataset nfcorpus \
  --data-dir data/processed/mteb \
  --index-cache-dir data/cache/mteb_index_e2e/nfcorpus \
  --score-device index \
  --index-accum-device hybrid

MS MARCO-style dataset slugs default to MRR@10; other retrieval datasets default to nDCG@10. More details are available in docs.

Citing this paper

If you find this work useful, please cite the accompanying paper:

@inproceedings{guo26ssr,
    title={No More K-means: Single-Stage Sparse Coding for Efficient Multi-Vector Retrieval},
    author={Lixuan Guo and Yifei Wang and Tiansheng Wen and Aosong Feng and Stefanie Jegelka and Chenyu You},
    year={2026},
    booktitle={International Conference on Machine Learning (ICML)},
}

Acknowledgements

This repository was built off of CSR series and Pylate. Thanks for their amazing works!