💬 NLP

Document Optimization for Black-Box Retrieval via Reinforcement Learning

This paper introduces a reinforcement learning-based framework that optimizes documents via fine-tuning language models to better align with target retrievers, achieving significant performance gains in code and visual document retrieval tasks while enabling smaller, more efficient models to outperform larger ones.

Original authors: Omri Uzan, Ron Polonsky, Douwe Kiela, Christopher Potts

Published 2026-04-08

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Omri Uzan, Ron Polonsky, Douwe Kiela, Christopher Potts

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a librarian in a massive, chaotic library. You have millions of books (documents), and people come in asking for specific information (queries).

The Old Problem:
In the past, librarians tried to help by adding more tags to the books. If a book was about "cats," they might add tags like "kittens," "felines," or "purring." This is called Document Expansion.

The Catch: In the modern era, the "librarians" (AI search engines) are so smart they don't just look for exact word matches. If you just randomly add tags, you often end up cluttering the book with noise. It's like writing "cat, dog, pizza, galaxy" on a book about cats. The AI gets confused, and the book ends up further down the list when someone searches for "cats."

The New Solution: "Document Optimization"
This paper proposes a smarter way. Instead of just adding random tags, imagine you have a super-smart editor (an AI) whose only job is to rewrite the book's summary specifically to please the librarian's brain.

Here is how it works, broken down into simple steps:

1. The "Black-Box" Librarian

The researchers treat the search engine (the librarian) as a "black box." They don't know exactly how the librarian's brain works inside. They can't ask, "How do you calculate relevance?"

The Trick: They only ask, "If I change this book's summary, does it move up or down the list when someone searches for 'X'?"
The Reward: If the book moves up the list for the right person, the editor gets a "gold star" (reward). If it moves down or confuses the search, the editor gets a "thumbs down."

2. The Editor's Job (Reinforcement Learning)

The AI editor uses a technique called Reinforcement Learning (think of it like training a dog with treats).

Trial and Error: The editor tries rewriting a document in many different ways.
- Attempt A: "This code is a loop." (Too vague)
- Attempt B: "This Python script calculates the factorial of a number using a while loop." (Better!)
- Attempt C: "A function that multiplies numbers 1 through N." (Perfect for the query!)
The Feedback Loop: The system tests these rewritten versions against the search engine. The version that gets the highest ranking gets the "gold star." Over time, the editor learns exactly what style of writing makes the search engine happy.

3. What Does the Rewrite Look Like?

The paper found that the AI doesn't just add words; it restructures them.

For Code: If a programmer searches for "how to sort a list," the original code might just be a block of symbols. The AI rewrites it to add natural language comments like: "This function sorts a list of numbers from smallest to largest." It makes the code "speak the language" of the person searching for it.
For Visual Documents (PDFs): If you have a chart in a PDF, the AI doesn't just say "Image 1." It writes a detailed description: "A bar chart showing sales growth in Q3, with a peak in September." This turns a picture into text that the search engine can understand perfectly.

4. The Results: Small Tools, Big Wins

The most exciting part of the paper is that this method is so effective that small, cheap AI models can beat giant, expensive ones.

The Analogy: Imagine a tiny, efficient car (a small AI model) that has been tuned by a master mechanic (the optimization process). It ends up driving faster and more efficiently than a massive, expensive limousine (a huge AI model) that hasn't been tuned.
Real-world impact: They took a small, cheap search model and optimized the documents for it. The result? It performed better than a model that costs 6.5 times more.

Why This Matters

It's Offline: You do all the rewriting before anyone searches. When a user asks a question, the search is instant. No extra waiting time.
It Works Everywhere: It works for text, code, images, and even different types of search engines.
It's Flexible: You don't need to know how the search engine works inside; you just need to see the results.

In a Nutshell:
Instead of trying to make the search engine smarter (which is hard and expensive), this paper teaches us how to make the documents smarter so they speak the search engine's language perfectly. It's like translating a book into the exact dialect the reader loves, ensuring they find it immediately.

1. Problem Statement

Document Expansion is a classical technique used to improve retrieval quality by augmenting documents with terms likely to appear in relevant queries. Traditionally, this shifts computation offline, avoiding query-time latency. However, when applied to modern neural retrievers (which rely on semantic understanding rather than exact term overlap), naive document expansion often degrades performance by introducing noise that obfuscates the discriminative signal.

The core problem addressed is how to effectively transform documents to better align with a specific target retriever's representation space without requiring access to the retriever's internal gradients (i.e., in a black-box setting). The goal is to learn document transformations that maximize retrieval metrics (like nDCG) rather than simply expanding text.

2. Methodology

The authors recast document expansion as a Document Optimization problem, framing it as a Reinforcement Learning (RL) task where a generative model acts as a policy to rewrite documents.

Core Framework

Policy ( $\pi_\phi$ ): An instruction-tuned Large Language Model (LLM) or Vision-Language Model (VLM) acts as the policy. It takes a document $d$ and a fixed prompt $p_{trans}$ as input and generates a transformed document $\tilde{d}$ .
Reward Signal (Counterfactual Retrieval): The reward is defined based on the change in retrieval performance if the original document were replaced by the transformed one.
- Counterfactual Collection: For a specific document $d_i$ , a baseline transformed collection $\tilde{D}$ is created. To evaluate a new candidate transformation $\tilde{d}_i$ , the system constructs a counterfactual collection $\tilde{D}_i$ where only $d_i$ is replaced.
- Metric: The reward is calculated as the change in nDCG@k (Normalized Discounted Cumulative Gain) for positive queries ( $Q^+$ ) minus the change for negative queries ( $Q^-$ ).
- Formula: $r_i(\tilde{d}_i) = \frac{1}{|Q^+|}\sum \Delta nDCG - \frac{1}{|Q^-|}\sum \Delta nDCG$ .
- This approach isolates the effect of the transformation on ranking, penalizing transformations that improve matches for irrelevant queries.

Optimization Algorithm

Group Relative Policy Optimization (GRPO): The authors use GRPO to optimize the policy. Instead of training a separate value function, GRPO samples a group of candidate transformations for a single document, computes their relative advantages based on the counterfactual rewards, and updates the policy. This reduces variance and is computationally efficient.
Training Loop:
1. Sample a batch of documents.
2. Generate $G$ candidate rewrites per document.
3. Evaluate candidates against the fixed collection to compute rewards.
4. Update the policy using GRPO advantages.
5. Periodically refresh the entire corpus with the new policy to maintain an up-to-date index.

Applicability

The method is architecture-agnostic. Since the learning signal depends solely on retrieval rankings (black-box access), it works with:

Single-vector retrievers (e.g., dense embeddings).
Multi-vector retrievers (e.g., ColBERT).
Lexical retrievers (e.g., BM25).

3. Key Contributions

Reframing Document Expansion: Moving from heuristic expansion to an optimization problem where documents are rewritten (compressed, restructured, or rephrased) specifically to align with a target retriever's latent space.
Black-Box Compatibility: The method requires no access to retriever weights or gradients, making it applicable to proprietary models (e.g., OpenAI embeddings) and diverse open-weight models.
Offline Efficiency: All computation happens offline during the indexing phase. Query-time latency and complexity remain unchanged.
Novel Reward Mechanism: The use of counterfactual $\Delta nDCG$ rewards allows the model to learn not just to increase similarity to relevant queries, but to actively suppress spurious matches with irrelevant queries.

4. Experimental Results

The authors evaluated the approach on Visual Document Retrieval (VDR) and Code Retrieval tasks using various retrievers (OpenAI, Qwen, Jina-ColBERT, BM25).

Key Findings

Performance Gains: Document optimization consistently improves retrieval quality (nDCG@5) across all settings.
- Code Retrieval: Applying optimization to OpenAI text-embedding-3-small improved nDCG@5 from 58.7 to 66.8, outperforming the much larger and more expensive text-embedding-3-large (66.3).
- VDR: OpenAI small improved from 53.3 to 57.6, surpassing the large model (57.0).
Lexical & Multi-Vector Models: The gains were most pronounced for models that traditionally struggle with semantic nuance.
- BM25: Improved dramatically from 15.6 to 46.6 on code retrieval.
- Jina-ColBERT-V2: Improved from 48.6 to 61.8 on code retrieval when combined with joint adaptation.
Comparison with Fine-Tuning:
- In black-box settings, optimization is the only viable path and yields significant gains.
- In white-box settings, optimization is competitive with fine-tuning the retriever.
- Joint Adaptation: Combining document optimization with retriever fine-tuning yields the best results in most settings (e.g., Jina-ColBERT-V2 on VDR jumped from 55.8 to 63.3).
Nature of Transformations: Analysis showed that the policy does not simply expand text (which increases noise). Instead, it produces compact, refined representations (e.g., adding natural language comments to code, wrapping code in blocks, or summarizing visual content) that better match the query distribution.

5. Significance and Implications

Democratizing High-Performance Retrieval: The ability to make smaller, cheaper models (like text-embedding-3-small) outperform larger, expensive models through document optimization offers a cost-effective alternative to scaling model size.
Offline Optimization: By shifting the computational burden to the indexing phase, this approach preserves the low-latency requirements of production search systems.
Flexibility: The framework is versatile enough to handle diverse data types (code, images, text) and retrieval architectures, suggesting that "optimizing the document space" is a fundamental and underexplored direction in information retrieval.
Practicality: The method works with weak supervision (using hard negatives mined from the retriever itself), reducing the dependency on large, manually annotated query-document pairs.

In conclusion, the paper demonstrates that treating documents as learnable, optimizable objects via reinforcement learning is a powerful strategy to bridge the gap between document content and retriever representation spaces, significantly boosting retrieval performance without increasing inference costs.