DEPT: Document Embedding Preservation Tuning for Unified Query Expansion and Retrieval

Abstract: 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.

• The paper introduces DEPT, an end-to-end method that uses document-embedding preservation, fixed whitening, and straight-through expansion training to adapt query generation without destabilizing document representations.
• DEPT achieves the best reported average BEIR nDCG@10 on Qwen3-4B (42.59) and LLaMA-3.2 (39.60), while DEPT-K offers a competitive short-expansion setting with about nine generated tokens.
• The paper shows that preserving document embeddings supports reuse of a pre-tuning FAISS index and protects generation quality, retaining near-baseline language-task performance while standard contrastive learning collapses dramatically.

Motivation and problem statement

The paper addresses the gap between query expansion and dense retrieval in LLM-based systems. A single decoder-only LLM can in principle generate an expansion for an underspecified query and encode both the expanded query and candidate documents, making expansion a trainable retrieval action rather than a prompted preprocessing step. Prior systems stop short of this: training-free methods such as HyDE and Query2Doc leave the generator unchanged (Wang et al., 2023, 2305.13614); independent-training methods such as InPars and Promptagator optimize generator and retriever through separate interfaces; ExpandR uses indirect retrieval feedback rather than the deployed contrastive loss; and UniRAG, the closest staged unified framework, trains augmentation and representation in separate phases. The common deficiency is that the text-producing component is never continuously shaped by the retrieval objective and embedding geometry that determine final ranking.

Directly applying contrastive learning to a unified generator–encoder creates a distinct failure mode. Because the same parameters define document embeddings and generation behavior, the InfoNCE gradient on the raw score $r(q_i,y_i,d)=\mathbf{q}_i^\top\mathbf{d}$ decomposes into a desired query-side term plus a document-side term that relocates retrieval targets. The paper shows empirically that the result is a bottleneck rather than collapse: ranking improves, but expansion quality degrades and the gain is limited. The authors therefore argue the training problem is asymmetric: the query path must remain plastic, while document embeddings must serve as approximately stable contrastive targets.

Method

DEPT trains one decoder-only LLM end to end under the objective $$\mathcal{L}_{\mathrm{DEPT}}=\mathcal{L}_{\mathrm{RET}}+\lambda\mathcal{L}_{\mathrm{DEP}}$$, with three coupled mechanisms.

Document Embedding Preservation (DEP) loss. Cached reference embeddings $\mathbf{d}^0=f_{\theta_0}(d)$ are precomputed from the initial model and receive no gradient. The loss penalizes angular drift $\delta(d)=1-\cos(\mathbf{d},\mathbf{d}^0)$ over the positives and sampled negatives in each minibatch, scaled by $s$ and exponentiated with $p>1$ to emphasize large departures. Notably, the paper claims a secondary benefit: since document encoding shares token embeddings and transformer blocks with document-conditioned language modeling, preservation acts as lightweight functional rehearsal that protects general generation ability.

Fixed whitening. Because decoder-only LLM embeddings are anisotropic, DEPT estimates a whitening transform $W=U(\Lambda+\varepsilon I)^{-\alpha}$ from the cached document embeddings and freezes it throughout training. This design is explicitly tied to preservation: whitening remains valid only if the DEP loss keeps current embeddings close to the distribution used to estimate $(\mu, W)$. For baselines whose document embeddings drift, a pre-fitted whitening transform would degrade into an arbitrary fixed linear map.

Straight-through expansion training. Hard token selection would sever the gradient path from the retrieval loss to the expansion logits. DEPT uses straight-through decoding: the forward pass consumes the greedy token embedding, while the backward pass uses the gradient of a top-$k$ ($k=64$) soft mixture over the vocabulary. Retrieval gradients thus update both the expanded-query representation and the logits that produced the expansion.

The stable document side enables two practical capabilities. Index reuse: the pre-tuning FAISS index over cached whitened embeddings remains compatible with the tuned query encoder, avoiding full corpus re-encoding after training. Online hard-negative mining: at each step, the current expanded query searches the fixed cached index, and negatives are sampled from a high-ranking interval (rank 5) after removing the positive, giving query-conditioned negatives that track current expansion behavior without refreshing the corpus index.

Experimental results

Training uses LoRA ($$\mathcal{L}_{\mathrm{DEPT}}=\mathcal{L}_{\mathrm{RET}}+\lambda\mathcal{L}_{\mathrm{DEP}}$$0) on an eight-source mixture (ELI5, FEVER, HotpotQA, MS MARCO passage and document, NQ, SQuAD, TriviaQA) with 1k steps at batch 256, evaluated zero-shot on five BEIR tasks with nDCG@10, using Qwen3-4B-Instruct-2507 and LLaMA-3.2-3B-Instruct backbones.

Backbone	Method	Avg. nDCG@10	Exp. tokens
Qwen3-4B	Original query	36.14	0
Qwen3-4B	HyDE	39.85	364.50
Qwen3-4B	ExpandR	41.09	174.82
Qwen3-4B	UniRAG	38.62	56.86
Qwen3-4B	DEPT	42.59	161.87
Qwen3-4B	DEPT-K	41.45	9.15
LLaMA-3.2	Original query	36.18	0
LLaMA-3.2	ExpandR	38.17	164.86
LLaMA-3.2	DEPT	39.60	148.82

Long-expansion DEPT achieves the best average on both backbones, and DEPT-K is competitive with roughly nine generated tokens, giving a short-expansion operating point. SCIDOCS is the consistent exception, where concise topical cues beat long expansions. The gains are not a length artifact: training-free methods produce longer text that is not optimized for the embedding scorer.

The ablations on Qwen attribute the gains causally. Removing whitening is the largest single drop (average 50.82 → 42.09 on the three-task subset), removing the DEP loss drops the average to 46.95 even when documents are re-encoded, showing preservation improves optimization beyond index reuse; removing online negatives, detaching the expansion, or removing expansion entirely each cost 2.5–8 nDCG points.

Two results are particularly strong claims. First, on cached-index compatibility, retrieval against the pre-tuning index stays close to full re-encoding with DEPT, whereas without the DEP loss the original index contains invalid keys for the tuned encoder. Second, on generation preservation, DEPT retains near-baseline performance on GSM8K, IFEval, and MMLU-Pro (77.56 → 76.28 average), while standard contrastive learning on the same unified model collapses to 9.32 — a stark demonstration that unconstrained joint training destroys the generator. A qualitative case study on an NFCorpus query corroborates this: standard CL and UniRAG emit degenerate repetitive strings ("Starving tumors Starving Starving tumor Star Star Star tumors"), while DEPT-K compresses the query into high-value cues and full DEPT preserves readable explanatory text.

Limitations and open questions

The paper concedes several boundaries. Evaluation is restricted to five BEIR tasks with a single training run per configuration (seed 42), and UniRAG is a same-backbone reimplementation rather than the original system, which weakens the staged-unified comparison. The whitening transform is estimated from only 1,024 cached documents, and its validity is conditional on the DEP loss holding embeddings near the reference distribution — the paper does not quantify how much drift the transform tolerates before degrading. Straight-through decoding is applied only at the top-$$\mathcal{L}_{\mathrm{DEPT}}=\mathcal{L}_{\mathrm{RET}}+\lambda\mathcal{L}_{\mathrm{DEP}}$$1 level, and the sensitivity of end-to-end expansion training to this approximation is not analyzed. The DEP loss regularizes only documents appearing in minibatches, so embeddings of documents never sampled during training are preserved only implicitly through parameter sharing; whether this suffices for very large corpora is left unexamined. Finally, the fixed cached index means retrieval quality is bounded by the initial model's document representation; whether preservation constraints cap achievable retrieval quality relative to fully retuned encoders is an open question the paper raises implicitly but does not answer.

Conclusion

DEPT formulates unified query expansion and dense retrieval as a single decoder-only LLM training problem and identifies document-embedding drift as the obstacle preventing contrastive learning from fully exploiting generated expansions. Its asymmetric principle — aggressive query-side adaptation under a preservation loss and fixed whitening on the document side — yields state-of-the-art average nDCG@10 among compared expansion paradigms on two backbones, while retaining generation ability and cached-index compatibility. The design principle the paper distills is that unified generative–representational retrieval benefits from explicitly stabilizing document embeddings rather than treating their movement as a harmless side effect of joint training.