Passage reranking

Overview

A retriever returns several candidate passages for a query. A reranker then scores those candidates again, and only the best one goes on to the reader. In the one card so far, the candidates are the top 16 BM25 passages for the text written so far, and the chosen passage is placed in front of a frozen language model 2302.00083 §6 2302.00083 §3.1. The language model's weights do not change; the reranker only decides which passage is used 2302.00083 §6.2.

The paper's numbers are perplexity on language-modelling sets, reported by the authors for their own setups 2302.00083 tab. 1. Nothing here shows reranking on question answering.

Variants

Zero-shot LM scoring

No training. Each candidate passage is placed before the prefix, and the language model's log-probability of the last s' prefix tokens is computed under that passage. The candidate with the highest score is kept 2302.00083 §6.1. The scoring model does not have to be the generating model 2302.00083 §6.1. The scoring length s' was set to 16 after a search over {4, 8, 16, 32} on WikiText-103 dev 2302.00083 §6.1. Card: In-Context Retrieval-Augmented Language Models.

With BM25 and GPT-2 Small, word-level perplexity on WikiText-103 falls from 29.6 to 28.6 when zero-shot reranking is added. The four other corpora move in the same direction 2302.00083 tab. 1.

Trained predictive reranker

A RoBERTa-base classifier gives each of the k = 16 candidates a relevance score against the prefix. The scores are softmax-normalised over the candidates, and the best one is chosen 2302.00083 §6.2. Training uses a loss built from the frozen language model's likelihood of the next stride under each candidate, so only the reranker is updated 2302.00083 §6.2. The training set has 300,000 examples from WikiText-103 train, scored with GPT-2 110M. The paper reports 10,000 steps, a peak learning rate of 1e-5 and a batch size of 32; the schedule, optimizer and dropout are not stated 2302.00083 §6.2.

With BM25 and GPT-2 Small, WikiText-103 word-level perplexity is 26.8, against 28.6 for zero-shot scoring and 29.6 without reranking. With GPT-2 XL, it is 15.4, against 16.1 and 16.6 2302.00083 tab. 1. The paper reports this variant on WikiText-103 only 2302.00083 tab. 1.

When to use which

Implementation sketch

The two variants share the same loop: score each candidate against the prefix, then pick one. They differ in where the score comes from.

# Zero-shot reranking of k candidate passages  [arxiv:2302.00083 §6.1]
# Input:  prefix tokens, k passages, scoring LM, scoring length s'
# Output: index of the chosen passage

def rerank_zero_shot(prefix, passages, lm, s_prime):
    # the target is the last s' prefix tokens, known at test time  [§6.1]
    target = prefix[-s_prime:]
    context = prefix[:-s_prime]
    # score each passage placed before the context  [§6.1]
    scores = [log_prob(lm, target, p + context) for p in passages]
    return argmax(scores)

# Loss for one training example of the predictive reranker  [arxiv:2302.00083 §6.2]
# Input:  prefix, k passages, next-stride tokens y, reranker f, frozen LM
# Output: scalar loss

def predictive_loss(prefix, passages, y, reranker, lm):
    # relevance scores, softmax-normalised over the k passages  [§6.2]
    weights = softmax([reranker(prefix, p) for p in passages])
    # LM likelihood of the next stride given each passage  [§6.2]
    likes = [log_prob(lm, y, p + prefix) for p in passages]
    # -log of a weighted sum, computed in log space  [§6.2]
    return -logsumexp([log(w) + ll for w, ll in zip(weights, likes)])

Contradictions

None known. Only one card covers this concept, so the page has no cross-paper disagreement to report, and the evidence rests on one paper's setup.

Papers

• In-Context Retrieval-Augmented Language Models: applies zero-shot LM scoring and a trained RoBERTa reranker to the top 16 BM25 passages for a frozen GPT-2 family model.