Skip to content

Recall-stage rerank-stage search pipeline - #463

Open
luciaquirke wants to merge 1 commit into
mainfrom
feat/multistage-attribution
Open

Recall-stage rerank-stage search pipeline#463
luciaquirke wants to merge 1 commit into
mainfrom
feat/multistage-attribution

Conversation

@luciaquirke

@luciaquirke luciaquirke commented Sep 13, 2026

Copy link
Copy Markdown
Collaborator

Grosse et. al recommends a search pipeline-style setup where the more expensive method only reranks the top-k queries according to a cheaper and/or recall-optimized method that overfetches.

  • Add score_cfg.candidates on score for the overfetched candidates. ekfac and trackstar pull from this. Two steps in one YAML can compose the stages.
  • Candidates are the union over the earlier run's query columns of each column's top_k or fraction rows at the direction end (proponents / detractors).

TODO claude slop

  • The store keeps every row. Candidates carry the new scores; the rest keep the earlier run's, negated if the runs disagree on higher_is_better, scaled to the candidates' spread and shifted past the weakest candidate, so validate, recall and load_scores see one complete ranking. candidates.npy lists the rescored rows.
  • MemmapSequenceScoreWriter gains rows, mapping subset indices onto store rows.
  • Per-token stores are rejected. The earlier run must have the same number of score columns, or one.

Examples: examples/pipelines/projected_kfac_then_shampoo.yaml, examples/pipelines/trackstar_then_shampoo.yaml. Docs: docs/influence-functions.rst.

@luciaquirke
luciaquirke force-pushed the feat/multistage-attribution branch from 53c3f2d to 44e2a55 Compare September 13, 2026 06:35
@luciaquirke luciaquirke changed the title Rescore the rows a cheaper scoring run ranked highest Score only the rows an earlier run ranked highest Sep 13, 2026
Add score_cfg.candidates to score, ekfac and trackstar. It points at an
earlier run over the same training set and keeps the union over its query
columns of each column's top_k or fraction rows, so e.g. projected KFAC
or TrackStar can run over every row and Shampoo over its top rows. The
Hessian is still fitted on the whole training set.

The store keeps every row: candidates carry the new scores and the rest
keep the earlier run's, oriented to this run, scaled to the candidates'
spread and shifted past the weakest candidate, so validate and recall see
one complete ranking. candidates.npy lists the rescored rows.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JMR4GS4V8QRBXgJ7MusgQa
@luciaquirke
luciaquirke force-pushed the feat/multistage-attribution branch from 44e2a55 to cfb1b6e Compare September 13, 2026 07:23
@luciaquirke luciaquirke changed the title Score only the rows an earlier run ranked highest Recall-stage rerank-stage search pipeline Sep 13, 2026
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