SearchLM NL2BM25 — GRPO v2 Shaped Reward ✅ (Qwen2.5-3B-Instruct)

Part of the SearchLM collection · GitHub

The best-performing SearchLM checkpoint. Trained via GRPO with a shaped reward that eliminated the specification gaming found in GRPO v1 while simultaneously improving retrieval quality. Achieves NDCG@10 = 0.577 on NFCorpus and 0.657 on SciFact.

Pipeline position: base → SFT v1 → GRPO v1 (⚠️) → SFT v2 → GRPO v2 ✅


What it does

The model reasons step-by-step about key concepts, synonym expansion, and boolean structure, then emits a Tantivy-compatible boolean search query:

Input: Do Cholesterol Statin Drugs Cause Breast Cancer?

Output:

<reasoning>
Key concepts:
1. Statin drugs — synonyms: statin, "HMG-CoA reductase inhibitor", simvastatin,
   atorvastatin, lovastatin, pravastatin
2. Causal relationship — cause, risk, association, induce, "increase risk"
3. Breast cancer — "breast cancer", "breast carcinoma", "breast neoplasm"

Strategy: AND the three concept groups; use OR to expand synonyms within each.
Phrase-quote multi-word terms; keep AND chains short to avoid zero-result queries.
</reasoning>
<query>(statin OR "HMG-CoA reductase inhibitor" OR simvastatin OR atorvastatin OR lovastatin)
AND (cause OR risk OR association OR induce)
AND ("breast cancer" OR "breast carcinoma" OR "breast neoplasm")</query>

Compare to GRPO v1's output for the same query:

<reasoning>
</reasoning>
<query>Cholesterol Statin Breast Cancer</query>

GRPO v2 generates 147-token completions with substantive reasoning; GRPO v1 generated 5-token keyword bags with empty reasoning blocks.


How v2 eliminated reward hacking

The v1 reward (0.6 × NDCG@10 + 0.4 × MRR) was gameable with keyword bags on small corpora because BM25 recall on 3–5K doc indexes is high for distinctive nouns. Three mechanisms closed this gap in v2:

# v2 reward function
base   = 0.6 * max(0, ndcg_at_10 - keyword_baseline_ndcg)  # must beat noun-extraction
       + 0.4 * mrr
shaped = complexity_mult * base                              # 1.0 with boolean ops, 0.5 without
       + 0.15 * min(reasoning_tokens / 100, 1.0)            # up to +0.15 reasoning bonus
reward = 0.0 if len(query.split()) < 3 else shaped          # hard gate: ≥3 tokens required
Mechanism Effect
Keyword baseline delta Model earns zero NDCG credit for matching naive noun-extraction
Hard length gate Single/double-word queries unconditionally return 0.0
Reasoning depth bonus Up to +0.15 reward for ≥100-token reasoning blocks
Complexity multiplier Queries without boolean operators earn half credit

All SearchLM checkpoints

Model NFCorpus NDCG@10 SciFact NDCG@10 Mean tokens Boolean ops
base (Qwen2.5-3B-Instruct) 0.455 0.386 120 ~20%
SFT v1 0.441 0.273 95 ~80%
GRPO v1 ⚠️ 0.556 0.608 5–7 0%
SFT v2 0.466 0.358 109 ~65%
GRPO v2 0.577 0.657 147 ~35%

Evaluated on BEIR test splits (NFCorpus: 323 queries, SciFact: 300 queries).


Behavioral comparison (GRPO v1 vs GRPO v2)

Dimension GRPO v1 ⚠️ GRPO v2
NFCorpus NDCG@10 0.556 0.577 (+0.021)
SciFact NDCG@10 0.608 0.657 (+0.049)
Mean completion length 5–7 tokens 147 tokens
Boolean operator usage 0% ~35%
Phrase usage 0% ~25%
frac_reward_zero_std (step 1) 90–96% 0.0%
frac_reward_zero_std (final) 90–96% ~61%
Reasoning block empty substantive

The shaped reward did not sacrifice performance to eliminate gaming — it improved both.


Training Details

Setting Value
Base model searchlm-nl2bm25-sft-v2
Method GRPO (TRL GRPOTrainer + vLLM colocate, single H100)
Reward Shaped: complexity_mult × (0.6 × ΔNDCG + 0.4 × MRR) + 0.15 × reasoning_depth
Training datasets NFCorpus + SciFact + FiQA-2018 (3K queries, 57,638 docs)
num_generations 8 (was 2 in v1)
Epochs 1
Steps 2,879 (~3.3s/step)
Batch size 2 (+ 8 grad accum = effective 16)
Learning rate 1e-6
vLLM GPU utilisation 0.30 (24 GB KV cache)
Max new tokens 1,024
Gradient checkpointing yes
Hardware NVIDIA H100 80 GB
Training time ~3h 3m
Final train loss 0.0012
Final mean reward ~0.29
W&B run supreethrao/searchlm/runs/9x1tg52j

Why these hyperparameters

num_generations=8: v1 used 2, leading to 90-96% of groups having zero within-group reward variance (no gradient signal). With 8 completions, variance emerged from step 1.

vllm_gpu_memory_utilization=0.30: On H100 80GB, Adam fp32 optimizer states for a 3B model require ~24 GB. At 0.45 utilisation, vLLM reserved 36 GB and Adam states OOM'd. 0.30 leaves 56 GB for the training stack.

torch_compile=False: Compiled backward pass materialised fp32 FFN intermediate buffers (~90 MB each) that eager + gradient checkpointing avoids, causing OOM at batch_size=4.


Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "Supreeth/searchlm-nl2bm25-grpo-v2",
    torch_dtype="auto",
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("Supreeth/searchlm-nl2bm25-grpo-v2")

SYSTEM_PROMPT = """You are an expert information retrieval specialist. Convert the \
natural language query into a Tantivy boolean search query.

Output format (strictly follow this):
<reasoning>
Step-by-step concept extraction and synonym expansion.
</reasoning>
<query>your boolean query here</query>"""

nl_query = "effects of climate change on coral reef ecosystems"
messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": f"Convert to a Tantivy boolean search query:\n\n{nl_query}"},
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7, do_sample=True)
print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))

Tantivy Boolean Syntax

Tantivy is a full-text search engine library. The model targets its query language:

Construct Syntax Example
Single term word cancer
Exact phrase "phrase" "bone density"
AND A AND B vitamin AND calcium
OR A OR B cancer OR tumor OR malignancy
NOT NOT A NOT review
Grouping (A OR B) (cat OR feline) AND behavior
Field scope field:term title:"machine learning"
Boost term^N cancer^2 OR tumor

Related resources

Citation

@misc{searchlm2026,
  title  = {SearchLM: Training Small Language Models for Boolean Query Generation via RLVR},
  author = {Rao, Supreeth},
  year   = {2026},
  url    = {https://github.com/SupreethRao99/searchLM},
}
Downloads last month
55
Safetensors
Model size
3B params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for Supreeth/searchlm-nl2bm25-grpo-v2

Base model

Qwen/Qwen2.5-3B
Finetuned
(1432)
this model

Collection including Supreeth/searchlm-nl2bm25-grpo-v2