Source-linked AI summary
DEPT: Document Embedding Preservation Tuning for Unified Query Expansion and Retrieval
Jingyuan Wang, Richong Zhang, Zhijie Nie, Mingxin Li, Yanzhao Zhang
TL;DR
Unified query expansion and dense retrieval suffer when retrieval updates move the document embeddings that serve as targets. DEPT preserves document embeddings while training expansion and retrieval jointly, achieving the best average retrieval score for both backbones across five BEIR tasks while preserving generation and cached-index compatibility.
Problem
Unified query expansion and dense retrieval lack stable document targets, allowing retrieval supervision to improve local rankings by moving documents rather than adapting expansions.
Method
DEPT preserves current document embeddings near cached references with fixed whitening while keeping the query expansion path trainable through retrieval feedback.
Results
DEPT achieves the best average retrieval score with both backbones across five BEIR tasks, while generation declines only from 77.56 to 76.28 on average.
Takeaways & Limitations
Explicitly stabilizing document embeddings enables unified generative and representational retrieval while retaining generation and cached-index compatibility.
Abstract
from arXiv · showhide
Large language models (LLMs) can both expand underspecified queries and encode text as dense representations, suggesting a unified model for query expansion and retrieval. Existing systems usually rely on prompted expansions, independently trained modules, or staged optimization, leaving generated expansions only indirectly aligned with the retrieval loss that judges them. We train a single decoder-only LLM end to end, where the same model generates the expansion and encodes both the expanded query and candidate documents. This unified setting creates a moving-target problem: retrieval supervision should improve query-side expansion, but the same update also shifts the document embeddings that serve as retrieval targets. We introduce Document Embedding Preservation Tuning (DEPT), which keeps tuned document embeddings close to cached initial embeddings while allowing retrieval gradients to pass through straight-through decoding into the generator. DEPT converts joint query--document movement into query-side adaptation against approximately stable, whitened document embeddings that support index reuse and online hard-negative mining. Experiments with Qwen3-4B-Instruct-2507 and LLaMA-3.2-3B-Instruct on five datasets in BEIR benchmark show that DEPT improves average retrieval quality over training-free, independently trained, and staged unified baselines, while ablations isolate the effects of preservation, whitening, end-to-end expansion training, and online negatives. Code is available at https://github.com/ILSparkle/DEPT.
Introduction
The introduction frames unified query expansion and dense retrieval as a promising but unstable decoder-only LLM objective: retrieval updates can improve representations while degrading expansions and shifting document targets. DEPT addresses this asymmetry by preserving cached document embeddings while training expansion and query representations end to end.
- Motivation: Unified decoder-only LLMs can make expansion a trainable retrieval action while keeping generated text inspectable.The same model family can generate fluent expansions and dense representations, enabling retrieval loss to directly shape expansion.
- Limitations of Existing Methods: Existing training-free, independently trained, and staged methods do not continuously optimize the text-producing component with the deployed retriever’s final contrastive objective.These approaches leave generation unchanged, separate generator and retriever supervision, or use distinct augmentation and representation phases.
- Failure Mode: Ordinary contrastive learning can improve retrieval while document-embedding drift and altered generation behavior limit the benefit of unified expansion.The same parameters define query and document embeddings and generation behavior, so updates move both retrieval targets and the expansion process.
- DEPT: DEPT preserves current document embeddings near cached initial embeddings while allowing retrieval gradients to update expansion and expanded-query representations.This implements asymmetric training: the query side remains plastic, while the document side provides stable contrastive targets and supports cached indexing.
- DEPT: Fixed whitening, online hard-negative mining, and index reuse complement document preservation in DEPT.Whitening addresses anisotropic, poorly calibrated LLM embeddings, while stable document embeddings allow cached indices to remain useful during training and inference.
Related Work
Prior work shows that LLM-generated expansions and unified generative-representational models can improve retrieval, but typically lacks direct optimization through the final retrieval loss. Related studies also motivate preserving retrieval-target geometry when whitening is used in trainable systems.
- LLM-Based Query Expansion: HyDE, Query2Doc, and ExpandR use LLM-generated text to bridge query gaps, augment queries, or align expansions with retrieval preferences.UniRAG is the closest staged unified framework, using a decoder-only LLM for query expansion and encoding.
- LLM-Based Query Expansion: These expansion methods optimize generation outside the final contrastive retrieval loss or through indirect feedback, leaving shared-model retrieval targets difficult to stabilize.The unresolved issue arises when one decoder-only model handles both generation and encoding.
- Generative and Representational LLMs: Instruction-tuned embeddings and GritLM demonstrate that generative LLMs can produce strong representations and jointly support generation and embedding objectives.The paper distinguishes its retrieval-specific constraint from generic multitask unification because corpus embeddings serve as retrieval targets during training.
- Representation Geometry and Stability: Whitening can improve semantic similarity and retrieval by centering and decorrelating embeddings, but trainable systems require embeddings to remain near the distribution used to estimate the transform.This motivates treating representation geometry and stability as coupled concerns rather than applying whitening only as post-processing.
Method
DEPT uses one decoder-only LLM to generate query expansions and encode expanded queries and documents, while preserving cached document embeddings to stabilize retrieval targets. It combines straight-through expansion training, document embedding preservation, and fixed whitening for end-to-end retrieval optimization and index reuse.
- Unified model and asymmetric adaptation: The shared decoder-only LLM generates expansions and encodes expanded queries and documents, but preservation lets the query path adapt while keeping document retrieval targets stable.This asymmetry addresses gradient interference between query-side learning and document-side movement, with cached-index reuse as a downstream benefit.
- Document Embedding Preservation: DEPT caches each document’s initial embedding and penalizes angular drift of current document embeddings, with fixed references receiving no gradient.The preservation strength is controlled by λ, while scale s and exponent p > 1 make small deviations visible and emphasize larger departures.
- Whitening: Fixed whitening transforms cached document embeddings into a better-conditioned coordinate system before cosine retrieval, remaining useful because preservation limits subsequent document movement.The transform is estimated from cached references and held fixed throughout tuning; ε prevents unstable inversion of small eigenvalues.
- End-to-end retrieval training: DEPT retains discrete expansions during the forward pass but uses a differentiable backward approximation so retrieval loss reaches the expansion logits and expanded-query representation.Candidates include the positive document, in-batch documents, and sampled or mined negatives, while the shared model encodes every candidate document.
- Index Reuse: Preserving document embeddings avoids incompatibility between tuned queries and corpus vectors stored in a pre-tuning nearest-neighbor index.Without preservation, end-to-end tuning can require re-encoding the corpus before serving retrieval.
Experiments
Experiments show that DEPT improves unified retrieval by training expansion and encoding together while preserving document embeddings. Across BEIR evaluations and targeted analyses, preservation, whitening, end-to-end expansion training, and hard negatives support retrieval quality, index reuse, and retained generation ability.
- Evaluation Setup: DEPT is evaluated zero-shot on SciFact, ArguAna, NFCorpus, FiQA, and SCIDOCS using nDCG@10 and their unweighted average.Both Qwen3-4B-Instruct-2507 and LLaMA-3.2-3B-Instruct instantiate the methods alongside training-free, independently trained, and staged baselines.
- Main Results: Long-expansion DEPT achieves the best average retrieval score with both backbones, while DEPT-K remains competitive with roughly nine generated tokens.On Qwen, DEPT is the highest-quality point and DEPT-K defines the short-expansion frontier; gains are not explained by expansion length alone.
- Ablation Study: Removing the DEP loss, whitening, or expansion training reduces performance, showing that preservation stabilizes targets while end-to-end expansion training adapts the query side.The ablations also identify whitening as the largest individual factor and show that detaching expansion as an offline input loses straight-through training benefits.
- Index Reuse: With DEPT, cached-index retrieval remains close to re-encoding documents with the final model, whereas removing DEP makes the original index invalid for tuned document embeddings.Preservation prevents document-key drift while allowing the query path and LoRA parameters to adapt.
- Generation Ability: DEPT’s average generation score decreases from 77.56 to 76.28, while standard contrastive learning collapses to 9.32.The evaluation covers GSM8K, IFEval, and MMLU-Pro, indicating that preservation improves retrieval without reducing the decoder to an embedding-only model.
- Case Study: Case studies show DEPT-K compresses queries into high-value retrieval cues, while full DEPT preserves those mechanisms in readable explanatory expansions.Standard contrastive learning and UniRAG instead produce degraded or repetitive strings on the illustrated NFCorpus query.
Conclusion
The conclusion frames unified query expansion and dense retrieval as a moving-target problem because the same decoder-only LLM generates expansions and encodes candidate documents. DEPT addresses this dynamic through document embedding preservation and fixed whitening.
- The same decoder-only model generates retrieval-oriented query expansions and encodes the candidate documents used to judge them.
- Ordinary contrastive updates can improve local retrieval scores while drifting document embeddings and degrading the expansion generator.
- DEPT addresses the moving-target problem with document embedding preservation and fixed whitening.