STORM Qwen3-4B Keyword Generator

EMNLP 2026

Abstract

Modern retrieval increasingly relies on dense and learned-sparse neural models that are effective but require encoding the entire corpus into a specialized index, which must be rebuilt whenever the model changes. Lexical retrievers like BM25 stay efficient, transparent, and run 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, a beam of candidate expansions is scored against the BM25 index and low-reward continuations are 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 and far larger proprietary models while retrieving as fast as plain BM25, and transfers zero-shot to 18 languages (MIRACL), beating dedicated multilingual dense retrievers. STORM is thus a competitive, infrastructure-light alternative to dense neural retrieval.

Install

pip install -U torch transformers accelerate safetensors

Load and run (recommended: group beam search)

STORM must be decoded with group beam search (num_beams=6, num_beam_groups=3, diversity_penalty=1.0, num_return_sequences=3): this is the configuration used in the paper. Since transformers 4.56, group beam search is loaded from the Hub, so custom_generate and trust_remote_code=True must be passed to generate().

The STORM prompt is built into the chat template, so you only pass the raw query:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Arthur-75/storm-qwen3-4B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto")

messages = [{"role": "user", "content": "What are the symptoms of vitamin D deficiency?"}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, enable_thinking=False,
    return_tensors="pt", return_dict=True,
).to(model.device)

with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=64,
        do_sample=False,
        num_beams=6,              # [3, 6]
        num_beam_groups=3,
        diversity_penalty=1.0,
        num_return_sequences=3,   # [1, 3]
        custom_generate="transformers-community/group-beam-search",
        trust_remote_code=True,
    )

keywords = tokenizer.batch_decode(
    output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True
)
print(keywords)  # 3 comma-separated keyword lines

Important: calling model.generate(**inputs) without these arguments falls back to the defaults in generation_config.json: standard beam search (num_beams=6, num_return_sequences=3), not group beam search. Group beam search cannot be set in generation_config.json, because custom_generate and trust_remote_code can only be passed to generate(). Use the default only where you cannot pass these arguments (e.g. pipelines); results are slightly below the paper's.

Writing the prompt yourself

Code that builds the full STORM prompt explicitly still works and gives exactly the same input as the example above:

messages = [
    {"role": "system", "content": "From the query generate new semantic related keywords.\nOutput the result strictly as a single comma-separated line."},
    {"role": "user", "content": "[QUERY]: What are the symptoms of vitamin D deficiency?\n[KEYWORDS]: "},
]

Using the keywords with BM25

Serving with vLLM

vLLM uses the chat template, so sending only the raw query as the user message gives the correct STORM prompt. However, vLLM does not support group beam search and does not read beam settings from generation_config.json: use temperature=0 (greedy) or vLLM's beam search API. Results will be slightly below the paper's; use transformers with group beam search (above) to reproduce them.