GRC: Unifying Reasoning-Driven Generation, Retrieval and Compression

Abstract: Text embedding and generative tasks are usually trained separately based on LLMs nowadays. This causes a large amount of training cost and deployment effort. Context compression is also a challenging and pressing task, which is vital to reasoning-driven generation, and agentic tasks requiring long context and continual learning. In this paper, we explore how to unify reasoning-driven generation, reasoning-enhanced text representation and context compression tasks in one forward pass for LLMs. Through meta latent tokens and a unified generative, representative and compressive tuning approach, we propose a training framework named GRC that bridges the three tasks. The trained models can accomplish three objectives in a single forward pass while maintaining modular, LEGO-style flexibility during inference. This design greatly reduces the deployment effort for retrieval-augmented generation (RAG) and achieves efficient inference and three times data utilization during training. Furthermore, this framework design enables a new paradigm for text embedding: self-reason-latent embeds, and a new generation paradigm, latent memory-augmented generation, where compressed and internalized KV cache with O(1) length is used as the updatable memory. We also propose hybrid paged attention to speed up the inference of our models. Extensive experiments on reasoning-intensive retrieval benchmarks, generative tasks, document compression, latency evaluation, and RAG settings demonstrate the effectiveness of our method and may shed light on the truly unified model that can handle reasoning-driven generation, embedding and compression tasks seamlessly.

• The paper’s main contribution is a unified framework that integrates reasoning-driven generation, retrieval, and compression through meta latent tokens and custom masking.
• It demonstrates superior performance on benchmarks like BRIGHT and GSM8K, achieving enhanced retrieval accuracy, generative quality, and inference speed.
• Implications include reduced memory usage and computational overhead, supporting scalable retrieval-augmented generation and continual learning applications.

Unifying Reasoning-Driven Generation, Retrieval, and Compression with GRC

Motivation and Problem Statement

Current LLM deployment practices treat generative modeling, text embedding for retrieval, and context compression as distinct tasks, typically handled by separate models with disjoint representations and inference paths. This separation results in redundant computation, increased memory requirements, and suboptimal information reuse, particularly prominent in retrieval-augmented generation (RAG) and long-context or continual learning frameworks reliant on storage and manipulation of latent “memories.” The proliferation of reasoning-centric benchmarks and workflows further amplifies these inefficiencies, as reasoning traces require not only high-fidelity generation but also effective information distillation for retrieval and memory management. Previous attempts at unification, such as GritLM [26], failed to fully harmonize attention masking, latent memory management, and embedding extraction, limiting their practical applicability and efficiency.

GRC Framework: Design and Architecture

The proposed GRC framework introduces a decoder-only LLM paradigm wherein text generation, reasoning-augmented embeddings, and context compression are tightly integrated, operating under a causal attention mask and a rigorous, joint training protocol. The core architectural innovation is the introduction of meta latent tokens—trainable, non-intrusive registers that mediate the compression of contextual information and subsequent derivation of semantic embeddings. Unlike previous approaches, these meta tokens are not part of the external vocabulary and do not contaminate the generative process, but rather serve as specialized, reusable storage slots analogous to processor registers for neural computation.

Key technical elements include:

• Unified training with custom causal attention masks: The context is partitioned into segments for user instructions, queries, model generations (reasoning traces), meta latent tokens, and reconstruction instructions. Attention patterns are carefully designed to mask previous segments when encoding the reconstruction, forcing the model to rely solely on latent information for regeneration of context.
• Self-reason-latent-embed paradigm: Embeddings are constructed by first generating explicit reasoning traces, followed by the production and pooling over meta latent representations, enabling embeddings that capture both sequential reasoning and latent semantics.
• Latent memory-augmented generation: Context compression leverages the meta token KV cache, which has $O(1)$ size regardless of input sequence length, as an updatable latent memory, supplanting token-level document caches and enabling rapid context switching or extension in RAG and agentic settings.
• Hybrid paged attention (HPA): Extending paged attention [29], HPA manages both standard and compressed KV caches in block-oriented device memory allocation, substantially improving inference throughput and enabling seamless multi-task inference within a single forward pass.

Training Methodology

Training leverages a composite objective:

• Generative cross-entropy loss for reasoning-trace generation and QA,
• InfoNCE-style contrastive loss over pooled meta latent representations for semantic retrieval,
• Reconstruction loss over compressed latent states, enforcing high-fidelity context regeneration solely from meta token-derived information.

The unified dataset format ensures that both generation and retrieval data are mapped into scenarios where queries, positive/negative documents, and reasoning traces can be simultaneously exploited for all three training signals. Losses are balanced through scalar coefficients, and careful attention is paid to gradient management due to the substantial sequence lengths and memory demands.

Experimental Results

Retrieval Performance

On the BRIGHT benchmark [24], GRC-1.7B outperforms ReasonIR-8B, demonstrating the capability of small GRC models to act as advanced reasoning retrievers. Ablations show that the retrieval accuracy benefits from longer, self-generated reasoning traces—though plateaus beyond a certain context length—confirming the utility of explicit reasoning for embedding quality.

Generative Performance

GRC models show competitive generative results on BBH [33] and GSM8K [32], maintaining strong reasoning synthesis with both compressed and plain context. The architecture’s reasoning-driven generation matches or exceeds established decoder-only LLMs, even when grounding responses in compressed latent memory.

Context Compression

Evaluated on the PwC [10] and new Wikipedia Markdown datasets, GRC demonstrates robust context compression and reconstruction, with particularly pronounced gains for the 4B model at short sequence lengths. SacreBLEU and ChrF metrics show sustained high fidelity for compressed-context regeneration, with graceful degradation on longer contexts.

Retrieval-Augmented Generation

On Natural Questions (NQ) with BEIR corpus [34,35], GRC-4B exceeds GritLM-7B for both plain-text and compressed latent memory settings. The results indicate that latent KV cache memory is both information-rich and faithfully utilizable in downstream generation.

Inference Efficiency

HPA delivers 10x or greater speedups in average latency compared to naive implementations across all patterns (generation, embedding, compression, RAG), with compressed latent memory allowing context reuse and task switching at minimal inference cost, even on large models.

Theoretical and Practical Implications

GRC demonstrates that with careful meta token engineering and unified causal masking, it is possible to collapse the traditional separation between generation, semantic retrieval, and context compression without sacrificing performance on any axis. The primary implications are:

• Reduced deployment and operational complexity: A single model file, single forward pass, and a unified cache mechanism can cover diverse NLP and agentic workflow needs.
• Resource efficiency: Orders-of-magnitude savings in device memory due to $O(1)$ meta token-based caches compared to $O(N)$ raw token caches.
• Stronger cross-task generalization: Internal representations trained under joint objectives are better aligned, supporting tasks requiring rapid alternation between synthesis, lookup, and summarization.
• Foundational groundwork for continual learning and agentic research: The possibility of compositional “LEGO-style” modular inference, enabling flexible context updates and recycling.

Limitations and Future Directions

The current pooling strategy for producing embeddings may introduce conflicts between next-token prediction and embedding quality due to shared representation at the last meta latent token. More targeted architectural modifications, or selective decoupling, may further stabilize embedding quality. Investigating larger model scales, extension to multimodal memory, or adaptive latent register sizing are natural next steps.

Further research can explore tighter integration with emerging agent protocols, advanced cache server architectures, and extension toward compressing not just textual context but also images or structured data.

Conclusion

GRC advances the NLP model architecture landscape by fusing generation, retrieval, and compression into a single, efficient LLM framework. The method demonstrates—empirically and architecturally—that such unification yields not only resource and deployment benefits but measurable improvements in reasoning-intense and context-heavy tasks. GRC’s mechanisms are well-positioned to become foundational in scalable, agentic, and continual learning systems seeking robust and efficient operation across the spectrum of language understanding and manipulation tasks (2605.09100).