STORM: Stepwise Token Optimization with Reward-Guided Beam Search
About
Modern retrieval increasingly relies on dense and learned-sparse neural models that are effective but require encoding the entire corpus into a specialized index, rebuilt whenever the model changes. Lexical retrievers like BM25 stay efficient and transparent on a standard inverted index that need not change as models evolve, but suffer from vocabulary mismatch. LLM query rewriting can help, yet prompted rewriters emit well-formed but retrieval-ineffective or harmful-terms, and training against a retrieval reward gives only delayed, sequence-level supervision that obscures which terms helped. We introduce STORM (Stepwise Token Optimization with Reward-guided beaM search), a self-supervised framework for lexical query expansion. STORM trains the rewriter through generation guided by retrieval metrics: at each step, candidate expansions are scored against the BM25 index and low-reward continuations pruned, turning the retrieval reward into a token-level signal that concentrates exploration on retrieval-effective vocabulary. Across TREC DL and BEIR, STORM lets 0.6B-8B backbones match or surpass competitive LLM rewriters while retrieving as fast as plain BM25; at 8B it rivals far larger proprietary rewriters. It further transfers zero-shot to 18 languages (MIRACL), beating dedicated multilingual dense retrievers on average, making STORM a competitive, infrastructure-light alternative to dense neural retrieval.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Information Retrieval | BEIR | SciFact0.728 | 174 | |
| Information Retrieval | TREC DL 2020 | nDCG@1067.3 | 43 | |
| Information Retrieval | TREC DL 2019 | nDCG@1066.1 | 43 | |
| Information Retrieval | SciFact BEIR | NDCG@1072.8 | 36 | |
| Document Retrieval | MS MARCO document retrieval (DL'19) | NDCG@1066.1 | 32 | |
| Information Retrieval | TREC-COVID BEIR | NDCG@1079.6 | 27 | |
| Information Retrieval | DBPedia-Entity BEIR | NDCG@1040.1 | 27 | |
| Information Retrieval | TREC DL Hard | nDCG@100.342 | 26 | |
| Document Retrieval | MS MARCO document retrieval (DL'20) | NDCG@1067.3 | 25 | |
| Document Retrieval | MS MARCO (dev) | MRR@1022.9 | 20 |