Skip to content

Improve reranker ordering, low-confidence threshold and memory use - #1

Merged
space-sphere merged 1 commit into
mainfrom
improve-reranker
Oct 19, 2026
Merged

space-sphere merged 1 commit into
mainfrom
improve-reranker

Conversation

@space-sphere

Copy link
Copy Markdown
Contributor

Summary

  • Order reranked candidates by 1/(1 + reranker rank) + 0.5/(1 + retrieval rank) instead of by the reranker score alone. The reranker gives many candidates a probability close to 1, so ordering by it alone could push the best retrieval match out of the top results.
  • Write low_confidence_below: 0.9 instead of 0.5 into new indexes. Queries for results in StatLean score close to 1, so the higher threshold also flags queries for results StatLean does not contain that score between 0.5 and 0.9. Search reads the threshold from each index's index.json, so indexes built before this change keep 0.5.
  • Compute reranker logits for the last position only (logits_to_keep=1). Full-sequence logits took about 3 GB per batch of ten 1024-token prompts in half precision and ran out of memory on 11 GB GPUs.

Checks

  • pytest and pyflakes src tests
  • test_real_hybrid with a hybrid index of StatLean e436d8a1 on an RTX A6000
  • On an RTX A6000, computing logits for the last position only gives the same top five results and highest relevance for every evaluation query as full-sequence logits
  • On an RTX A6000, reranking the longest query from the StatLean website (600 characters) peaks at 3.6 GiB of GPU memory

- Order reranked candidates by 1/(1 + reranker rank) +
  0.5/(1 + retrieval rank) instead of by the reranker score alone. The
  reranker gives many candidates a probability close to 1, so ordering
  by it alone could push the best retrieval match out of the top
  results.
- Write low_confidence_below: 0.9 instead of 0.5 into new indexes.
  Queries for results in StatLean score close to 1, so the higher
  threshold also flags queries for results StatLean does not contain
  that score between 0.5 and 0.9.
- Compute reranker logits for the last position only. Full-sequence
  logits took about 3 GB per batch of ten 1024-token prompts in half
  precision and ran out of memory on 11 GB GPUs.
@space-sphere
space-sphere marked this pull request as ready for review October 19, 2026 05:00
@space-sphere
space-sphere merged commit 9adeabb into main Oct 19, 2026
3 checks passed
@space-sphere
space-sphere deleted the improve-reranker branch October 19, 2026 05:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant