Token-Level Credit Assignment Optimization for Generative Document Retrieval

2026-08-12

AI Summary

This work addresses the granularity mismatch in generative document retrieval, where document-level relevance supervision is misaligned with the token-level generation process, hindering effective credit assignment in reinforcement learning. To resolve this, the authors formulate retrieval as a sequential decision-making problem and introduce a token-level relevance reward mechanism that estimates the contribution of each token selection to final retrieval quality through trajectory-based reward modeling. By integrating autoregressive DocID generation with policy gradient optimization, the proposed approach significantly outperforms conventional sequence-level reward methods across multiple benchmarks, demonstrating both the efficacy and necessity of fine-grained supervision for enhancing generative retrieval performance.

Abstract

Generative retrieval models perform document retrieval by autoregressively generating document identifiers (DocIDs). This process naturally forms a sequential decision problem, where each decoding step selects a DocID token and the complete token sequence determines the retrieved document. However, retrieval effectiveness is typically evaluated only after the full DocID is generated, creating a mismatch between token-level generation and document-level relevance supervision. As a result, existing reinforcement learning methods for generative retrieval mostly rely on sequence-level rewards, where the same document-level feedback is propagated to all decoding steps. Such coarse-grained feedback makes it difficult to identify which token decisions are responsible for successful or failed retrieval. In this work, we propose a fine-grained reinforcement learning framework for generative retrieval with token-level relevance rewards. Instead of assigning a single reward to the entire generated DocID, we estimate step-wise rewards by measuring how each token decision changes the expected retrieval quality of the corresponding generation trajectory. This enables more precise credit assignment and encourages the policy to favor token decisions that contribute more directly to document-level relevance. We further develop practical reward estimation strategies tailored to the DocID generation process and incorporate them into a policy optimization framework. Experiments on retrieval benchmarks show that our method consistently outperforms sequence-level reward baselines, demonstrating the effectiveness of fine-grained supervision for aligning autoregressive DocID generation with retrieval objectives.

Problem

• generative retrieval
• credit assignment
• token-level reward
• document retrieval
• reinforcement learning

Innovation

• token-level credit assignment
• generative retrieval
• fine-grained reinforcement learning
• DocID generation
• step-wise reward