One-sentence Summary

Researchers from Apple and the University of Edinburgh propose CLaRa (Continuous Latent Reasoning), a unified framework that compresses retrieved documents into embedding vectors via the $\text{SCP}$SCP key-preserving synthesis framework and trains the reranker and generator end-to-end using a differentiable $\text{top-}k$top-k estimator, achieving state-of-the-art compression and reranking at a text compression rate of $16$16 across multiple QA benchmarks.

Key Contributions

• Introduces CLaRa, a unified framework that compresses RAG documents into continuous latent vectors and jointly optimizes the reranker and generator through a differentiable top-k estimator, enabling label-free retriever learning via the generator's language modeling loss.

• Proposes SCP, a key-preserving data synthesis framework that builds pretraining supervision from QA pairs, including compositional reasoning cases, and paraphrased documents to produce semantically rich, retrievable compressed vectors that shorten generator context.

• Evaluates on four single-hop and multi-hop QA benchmarks with Mistral-7B and Phi-4B, achieving state-of-the-art compression and reranking performance and outperforming supervised, unsupervised, and text-only DRO baselines at a text compression ratio of 16.

Introduction

Retrieval-Augmented Generation (RAG) is a powerful way to ground large language models in external evidence, yet most RAG systems suffer from a structural flaw: retrieval and generation are optimized separately. Retrievers rank documents by surface-level similarity while generators provide no feedback on what is actually needed, leading to two intertwined problems. First, efficiency suffers from an architectural mismatch, since dense retrievers operate in embedding space while generators consume raw text, causing redundant processing, higher inference costs, and context overflow. Second, optimization is blocked because discrete document selection prevents gradients from flowing back to the retriever, so it cannot align with the generator's task objective. Prior solutions, such as embedding-based compression or reinforcement learning for joint training, remain dependent on raw text, require query-specific recompression, or face unstable and computationally heavy training.

The authors contribute CLaRa (Continuous Latent Reasoning), a unified framework that performs retrieval and generation over shared continuous document representations. Each document is encoded once into a compact set of memory tokens that serve both retrieval and generation, eliminating redundant computation and enabling true end-to-end optimization. They introduce a salient-information-aware pretraining objective to preserve semantic fidelity, and use Straight-Through estimation to make the top-k document selection differentiable, allowing generator gradients to update the retriever directly without explicit retrieval labels. The authors show this unified objective yields valid gradients for retriever learning and evaluate CLaRa on single-hop and multi-hop QA benchmarks, outperforming both supervised and unsupervised baselines while achieving strong results with a high text compression ratio.

Dataset

For their salient compressor pretraining, the authors build a synthetic dataset designed to explicitly expose the semantic core of documents through question answering and paraphrasing. The construction pipeline is as follows:

• Source data: 2 million randomly sampled documents from the Wikipedia-2021 corpus.
• Supervision generation: A locally deployed LLM (Qwen-32B) is used to produce three complementary supervision signals for each document:
  • Simple QA pairs: Each question targets a single atomic fact, promoting fine-grained factual retention. Redundancy is avoided by extracting only facts not covered by earlier questions.
  • Complex QA pairs: Each pair integrates multiple facts to encourage relational reasoning, explicitly connecting previously unrelated information to increase coverage.
  • Paraphrases: Surface structure is altered while semantics are preserved, demonstrating more compact expression of the same content.
• Verification and regeneration: For every document, the generated QA pairs and paraphrases are checked by the local LLM for factual consistency and information coverage. If missing content is detected, the LLM reviews the original text and existing pairs to generate additional questions capturing uncovered facts. This iterative process runs for up to 10 rounds, and samples that still fail coverage criteria are discarded. This filtering ensures only fully covered, factually faithful examples are retained in the training set.

The authors then use this dataset during both pretraining and instruction tuning:

• Compressor pretraining: The model receives a document appended with learnable memory tokens, and only the compressor LoRA adapter is active. The hidden states of the memory tokens form the compressed representation. Training uses a cross-entropy loss for answer generation (with QA or paraphrasing as the instruction) plus a mean squared error term that aligns the averaged hidden states of the document tokens with those of the memory tokens, keeping the compressed latent space faithful to the original text.
• Instruction tuning: To adapt the pretrained compressor for downstream QA and to obtain an answer generator that can work with continuous document representations, the authors optionally finetune both the compressor and generator LoRA adapters jointly. Downstream training datasets are used, where retrieved documents are paired with task instructions as inputs, and the outputs are reference responses produced by a teacher model conditioned on the same documents and instructions.

Method

The authors propose a two-stage framework that first distills documents into compact semantic representations and then jointly optimizes retrieval and generation. The initial stage, Salient Compressor Pretraining (SCP), focuses on learning a highly informative document representation. To prevent the model from wasting capacity on trivial token-level reconstruction, the authors construct a synthetic dataset that explicitly exposes salient information through question-answering pairs and paraphrased documents. This pipeline involves iterative verification and regeneration to ensure complete factual coverage. Based on this curated data, the compression framework is trained. Refer to the framework diagram

The underlying architecture employs a shared base model equipped with distinct Low-Rank Adaptation (LoRA) adapters for modular control. During compression, a compressor adapter $\theta_c$θc​ processes the original document tokens $d_i = \{t_1, \ldots, t_m\}$di​={t1​,…,tm​} appended with $l$l learnable memory tokens $(m_1, \ldots, m_l)$(m1​,…,ml​). The final-layer hidden states of these memory tokens yield the compressed representation:

To ensure this compact vector faithfully captures the original semantics, the training objective combines a cross-entropy loss for text generation with a mean squared error alignment loss. This alignment loss minimizes the distance between the averaged hidden states of the original document tokens and the memory tokens, forcing the compressor to summarize the same semantic space without drifting:

Once the compressor is pretrained, the framework transitions to CLaRa, which unifies retrieval and generation within a single language model through a differentiable retrieval module. Refer to the end-to-end training architecture

The pretrained compressor remains frozen to allow for efficient offline document encoding into dense embeddings. A query reasoner, implemented as a LoRA adapter initialized from the compressor, learns to encode queries into the identical representation space using the same number of memory tokens. This design enables the query reasoner to anticipate relevant document content through next-token prediction training. Retrieval is performed by computing the cosine similarity between the query embedding and the frozen document embeddings:

The top-$k$k most relevant document embeddings are then concatenated with the query and passed to a generator adapter $\theta_g$θg​, which produces the final answer. Both the query reasoner and the generator are updated simultaneously via a unified language modeling loss, allowing the retriever to receive implicit supervision directly from the generation objective without requiring explicit relevance labels.

To bridge the discrete nature of top-$k$k selection with continuous gradient flow, the authors employ Straight-Through (ST) estimation. This technique acts as a soft lens during training, preserving discrete retrieval behavior in the forward pass while enabling smooth gradient backpropagation through a softmax-weighted soft selection. The aggregated top-$k$k document representation is computed by applying the ST estimator to the candidate embeddings. This gradient coupling ensures that the retriever receives complementary signals: it is encouraged to rank documents that maximize generation likelihood while simultaneously receiving representation-level feedback from the generator. To validate the internal reasoning capabilities of the query reasoner, the authors apply logit lens analysis to project memory embeddings through the output head. Refer to the token decoding analysis

This analysis reveals that the query reasoner implicitly decodes reasoning-relevant knowledge and evidence that may not appear explicitly in the question, demonstrating effective semantic alignment and enhanced retrieval accuracy compared to baseline systems.

Experiment

The study validates a jointly optimized compression, reranking, and generation framework across four QA benchmarks, showing that soft compression preserves essential reasoning while filtering noise, often surpassing raw-text baselines. Joint training remains robust across compression ratios and benefits from instruction-tuned initialization in normal settings, while retrieval quality improves notably under oracle conditions. Notably, the weakly supervised retriever outperforms fully supervised baselines, and ablations confirm that diverse pretraining objectives and an MSE alignment loss enhance semantic consistency and overall performance.

The proposed document compressor consistently outperforms both soft and hard compression baselines across QA datasets, with larger gains when retrieval noise is removed. It also surpasses uncompressed text baselines, indicating that well-trained soft compression can preserve reasoning-critical information while filtering out irrelevant content. Performance degrades at extreme compression ratios, though the decline is moderate under normal retrieval conditions. Achieves average gains of 1.13% in Normal and 5.35% in Oracle settings over the best soft baseline. Exceeds uncompressed text baselines by 2.36% with Mistral-7B and 6.36% with Phi-4-min. Extreme compression beyond 32x hurts Oracle performance more noticeably but has a moderate impact in Normal settings.

The proposed method delivers end-to-end QA performance comparable to strong baselines while compressing the context by 16x. Under noisy retrieval, performance stays stable across compression ratios and even surpasses a text-based baseline on some datasets, while accurate retrieval notably boosts scores. The choice of initialization matters more when retrieval is noisy, but its impact diminishes as retrieval quality improves. At 16x compression, the method outperforms a text-based baseline on NQ and 2Wiki F1 scores under normal retrieval. Providing annotated positives in the oracle setting lifts F1 above 75% on both NQ and HotpotQA. Instruction-tuned initialization yields stronger results than pretraining-based initialization under noisy retrieval, but the gap narrows when retrieval is reliable.

Pretraining on QA-style and paraphrase objectives consistently improves instruction-tuning performance under oracle retrieval at a 32x compression ratio, with combined and diverse objectives yielding the largest gains. Paraphrase alone already provides substantial improvements, and adding complex QA types further boosts results, especially for the larger model. Any pretraining data composition outperforms the no-pretraining baseline across both models, confirming that both factual reasoning and paraphrastic rewriting enrich compressed representations. Combining all objective types (simple QA, complex QA, and paraphrases) produces the highest average gains, with particularly large improvements on multi-hop and open-domain benchmarks. Paraphrase-only pretraining offers a strong boost over baseline, sometimes rivaling or exceeding the addition of simple QA, indicating the value of diverse semantic coverage.

The table shows instruction-tuning performance with and without MSE loss at two compression ratios across four datasets. Adding MSE loss yields mixed effects: it improves scores on Musique and 2Wikiqa at both compression ratios, but lowers performance on NQ and sometimes on HotpotQA. The overall magnitude of change is modest, typically within about 1.5 points in either direction. MSE loss consistently improves Musique and 2Wikiqa at both 32x and 128x compression, with gains up to 1.56 points. NQ performance drops when MSE loss is added at both compression ratios, with the largest decline of 1.28 points at 128x. HotpotQA shows a slight gain at 32x but a slight loss at 128x, indicating the effect is not uniform across configurations.

The proposed document compressor consistently outperforms both soft and hard compression baselines and even uncompressed text on QA tasks, with larger gains when retrieval noise is removed, although extreme compression ratios lead to moderate degradation. Initialization choice matters more under noisy retrieval, while pretraining with combined QA and paraphrase objectives improves instruction-tuning, particularly for larger models. Adding an MSE loss yields mixed effects, boosting performance on some datasets but lowering it on others, with the impact varying across compression ratios and configurations.