🤖 AI

STORM: Stepwise Token Optimization with Reward-Guided Beam Search

STORM is a self-supervised framework that enhances lexical query expansion by using reward-guided beam search to optimize token-level generation against retrieval metrics, enabling efficient, transparent, and high-performing retrieval that rivals dense neural models without requiring specialized indexing.

Original authors: Arthur Satouf, Giulio D'Erasmo, Yuxuan Zong, Habiboulaye Amadou Boubacar, Pablo Piantanida, Benjamin Piwowarski

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Arthur Satouf, Giulio D'Erasmo, Yuxuan Zong, Habiboulaye Amadou Boubacar, Pablo Piantanida, Benjamin Piwowarski

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find a specific recipe in a massive, chaotic library. You ask the librarian (the search engine), "How do I make a cake?"

The Problem: The Vocabulary Gap
The librarian is very fast and uses a standard filing system (called BM25). However, the librarian only looks for the exact words you typed. If the best recipe in the library is titled "Delicious Vanilla Sponge with Buttercream," but you only asked for "cake," the librarian might miss it because the word "cake" isn't in the title.

This is the "vocabulary mismatch." You need to help the librarian understand that "cake" might also be called "sponge," "dessert," or "baked good."

The Old Solutions

The Human Helper: You could hire a human to rewrite your query for you. But humans are expensive and slow.
The AI Guess: You could ask a smart AI (a Large Language Model) to guess better words. But the AI often guesses words that sound right but don't actually help find the recipe. It might suggest "baking" or "flour," which are too common and clutter the search.
The Blind Trial-and-Error: Some methods let the AI guess a whole sentence, then check if it worked. If it failed, the AI has to guess again from scratch. This is slow and inefficient because the AI doesn't know which specific word caused the failure.

The New Solution: STORM
The paper introduces STORM (Stepwise Token Optimization with Reward-guided beaM search). Think of STORM as a smart, real-time editor that helps the AI write the perfect search query, one word at a time, while constantly checking the library's filing system.

Here is how it works, using a creative analogy:

The "Tree Climber" Analogy

Imagine the AI is trying to climb a tree to find the best fruit (the best search terms).

Standard AI: The AI picks a branch and climbs all the way up. Only when it reaches the top does it check, "Oh no, this branch has no fruit." It has to climb all the way back down and try a different branch. This is slow and wasteful.
STORM: STORM uses a reward-guided beam search. Imagine the AI is climbing, but at every single step (every new word it adds), a "scorekeeper" immediately checks the library index.
- If the AI adds a word like "sponge," the scorekeeper says, "Great! That matches a real recipe. Keep going!"
- If the AI adds a word like "delicious," the scorekeeper says, "Stop! That word is too common; it won't help us find the specific recipe. Cut that branch off immediately."

STORM prunes (cuts away) the bad branches as they grow. It only keeps the paths that are actually leading to good results.

Why This is a Big Deal

1. It Learns by Doing, Not by Memorizing
STORM doesn't need a human to teach it what a "good" search query looks like. It learns entirely by itself. It generates a query, checks if it finds good documents, and if it does, it remembers that path. If it doesn't, it forgets that path. It's like a dog learning to fetch a ball by only being praised when it brings the right ball back.

2. It's Lightning Fast
Many AI search tools are slow because they write long, fancy paragraphs to explain what you want. STORM is different. It writes short, punchy lists of keywords (like "sponge," "vanilla," "butter").

The Result: It finds better results than the old methods, but it does it as fast as the basic, dumb search engine (BM25). It doesn't need to rebuild the library's filing system; it just speaks the librarian's language better.

3. It Speaks Many Languages (Even if it only learned English)
The most surprising part of the paper is that STORM was trained only on English data. Yet, when tested on 18 different languages (like French, Chinese, or Swahili), it performed better than specialized, expensive AI models built specifically for those languages.

The Analogy: It's like teaching someone to play chess using only English rules, but then they can play perfectly in Spanish, French, and Japanese without ever being taught those languages. The AI learned the logic of finding good words, not just the specific English words.

The Bottom Line

STORM is a new way to make search engines smarter without making them slower or more expensive. It acts like a real-time coach for the AI, whispering, "Good word, keep going!" or "Bad word, stop right there!" at every single step. This allows even smaller, cheaper AI models to find information as well as (or better than) massive, expensive ones, all while using the simple, fast filing systems that libraries have used for decades.

Technical Summary: STORM (Stepwise Token Optimization with Reward-Guided Beam Search)

Problem Statement

Modern Information Retrieval (IR) faces a trade-off between efficiency and effectiveness. Dense and learned-sparse neural models offer high retrieval performance but require specialized indices that must be rebuilt whenever the model changes, incurring significant infrastructure costs. Conversely, lexical retrievers like BM25 remain highly efficient and transparent on standard inverted indices but suffer from the "vocabulary mismatch" problem, where relevant documents contain terms absent from the user's query.

While Large Language Models (LLMs) can rewrite queries to bridge this gap, existing approaches have limitations:

Prompt-based rewriters often generate well-formed but retrieval-ineffective terms or long pseudo-documents, incurring high inference latency without consulting the retriever during generation.
Reinforcement Learning (RL) approaches that train against retrieval rewards typically observe the reward only after the full sequence is generated. This sequence-level supervision creates an exploration bottleneck, making it difficult to identify which specific lexical terms contributed to the reward, leading to slow convergence or collapse onto generic expansions.

Methodology

The authors introduce STORM (Stepwise Token Optimization with Reward-guided beaM search), a self-supervised framework for lexical query expansion that converts retrieval rewards into token-level signals.

Core Mechanism: Reward-Guided Beam Search

Unlike standard beam search which selects candidates based on likelihood, STORM uses the retrieval system itself as a guide during the generation process.

Stepwise Scoring: At each generation step, candidate keyword sequences are scored against a BM25 inverted index.
Pruning: Branches with low retrieval rewards are pruned immediately. This concentrates exploration on vocabulary that is demonstrably effective for retrieval, rather than relying on the LLM's internal parametric knowledge alone.
Training Loop: The framework adapts the Generative Cooperative Networks (GCN) architecture. Instead of a learned discriminator, it uses a predefined retrieval reward function.
- Proposal Distribution: Training samples are drawn from a mixture of standard nucleus sampling and a reward-guided beam search.
- Objective: The policy is updated to minimize the KL divergence between the policy and a target distribution weighted by the retrieval reward.

Reward Function

The reward function $R(\hat{y})$ combines three components to optimize for retrieval quality while penalizing inefficiency:

Retrieval Quality: Measured by nDCG@100 against cross-encoder pseudo-labels (distilled from OpenSearch Project MS-MARCO Hard Negatives).
Vocabulary Penalty: A penalty term based on the document frequency (DF) of generated terms, discouraging the model from emitting high-frequency, non-discriminative terms.
Length Regularization: A log-likelihood term that acts as a prior for token selection and regularizes sequence length, preventing the generation of overly long or uninformative strings.

Key Contributions

Self-Supervised Training: A method for lexical query rewriting that requires no human-written rewrites; supervision is derived entirely from the retriever's scoring of generated expansions.
Token-Level Exploration: The introduction of reward-guided stochastic beam search, which transforms delayed sequence-level retrieval feedback into an immediate token-level signal, guiding the model to discover retrieval-effective vocabulary during generation.
Infrastructure-Light Efficiency: STORM achieves competitive effectiveness while maintaining inference latencies comparable to plain BM25, as it generates short keyword strings rather than full pseudo-documents.

Experimental Results

The authors evaluated STORM across in-domain, out-of-domain, and multilingual benchmarks using Qwen3 backbones ranging from 0.6B to 8B parameters.

In-Domain (TREC DL, MS-MARCO): STORM consistently outperforms prompt-based rewriters (HyDE, MuGI, W2P) and RL-based baselines (QUESTER) across all model sizes. At 8B, it rivals or surpasses competitive LLM rewriters.
Out-of-Domain (BEIR): STORM achieves the highest mean nDCG@10 across 12 BEIR collections. Notably, even the 0.6B STORM model outperforms all 0.6B baselines on average, demonstrating robust generalization.
Zero-Shot Multilingual (MIRACL): Trained exclusively on English MS-MARCO data, STORM-8B transfers zero-shot to 18 languages, achieving an average nDCG@10 of 48.4. This surpasses dedicated multilingual dense retrievers (mDPR, mContriever, mColBERT) on average, despite relying on a purely lexical backend.
Comparison with Proprietary Models: STORM-8B (Qwen3) outperforms several query rewriters based on the much larger GPT-4.1 model (e.g., MuGI-GPT4, Q2D-GPT4) on average across in-domain and BEIR datasets, suggesting the reward signal is more effective than chain-of-thought prompting for retrieval tasks.
Efficiency: STORM maintains retrieval latency equivalent to plain BM25 (~0.01s) despite issuing queries with significantly more terms (up to 8x more). This is achieved because the reward function penalizes high-frequency terms that would otherwise cause costly index traversals. Generation latency is also lower than competing LLM rewriters (e.g., 4.39s for STORM-8B vs. 34.2s for W2P).

Significance and Claims

The paper positions STORM as a competitive, infrastructure-light alternative to dense neural retrieval. Its primary significance lies in demonstrating that:

Lexical expansion can be optimized effectively without dense indices or massive proprietary models.
Reward-guided decoding is a viable general tool for aligning generation with non-differentiable retrieval objectives, solving the exploration bottleneck inherent in sequence-level RL.
Zero-shot transfer to multilingual settings is achievable with a purely lexical approach, challenging the assumption that dense semantic matching is strictly necessary for cross-lingual retrieval.

The authors conclude that STORM offers a robust balance of effectiveness and efficiency, making it a practical solution for search engines and Retrieval-Augmented Generation (RAG) systems where index rebuilds are costly or latency is critical.

Limitations

The authors acknowledge that while STORM narrows the gap to learned sparse retrievers (like SPLADE-v2) on in-domain tasks, it does not fully match them, likely due to the inability to capture rich contextual relationships encoded by term-weighting mechanisms in dense models. Additionally, the optimal generation length is dataset-dependent, requiring a trade-off between latency and effectiveness.