Policy-Gradient Training of Language Models for Ranking

TL;DR

This paper introduces Neural PG-RANK, a policy gradient-based training method for language model retrievers that directly optimizes ranking quality, reducing reliance on heuristics and improving performance in text retrieval tasks.

Contribution

Neural PG-RANK is a novel end-to-end training algorithm that models ranking as a Plackett-Luce policy, aligning training objectives with downstream decision quality.

Findings

Significant in-domain performance improvements.

Enhanced out-of-domain generalization.

Effective unification of training and decision metrics.

Abstract

Text retrieval plays a crucial role in incorporating factual knowledge for decision making into language processing pipelines, ranging from chat-based web search to question answering systems. Current state-of-the-art text retrieval models leverage pre-trained large language models (LLMs) to achieve competitive performance, but training LLM-based retrievers via typical contrastive losses requires intricate heuristics, including selecting hard negatives and using additional supervision as learning signals. This reliance on heuristics stems from the fact that the contrastive loss itself is heuristic and does not directly optimize the downstream metrics of decision quality at the end of the processing pipeline. To address this issue, we introduce Neural PG-RANK, a novel training algorithm that learns to rank by instantiating a LLM as a Plackett-Luce ranking policy. Neural PG-RANK provides a principled method for end-to-end training of retrieval models as part of larger decision systems via policy gradient, with little reliance on complex heuristics, and it effectively unifies the training objective with downstream decision-making quality. We conduct extensive experiments on various text retrieval benchmarks. The results demonstrate that when the training objective aligns with the evaluation setup, Neural PG-RANK yields remarkable in-domain performance improvement, with substantial out-of-domain generalization to some critical datasets employed in downstream question answering tasks.