Source-linked AI summary

Compressing Long Context into Answer-Aligned Memory Embeddings for LLM Inference

Md Mostafizer Rahman, Md Faizul Ibne Amin, Md Shahajada Mia, Yutaka Watanobe, Fang Liu

arXiv:2609.25537v1cs.CLcs.AI

TL;DR

Long-context inference increases self-attention computation and KV-cache memory demands, while existing soft-compression methods lack a joint solution for cross-architecture deployment, answer-targeted training, and query-guided memory selection. CMC compresses contexts into embeddings for frozen decoders and combines selected memory with a local context window. Across benchmarks and model pairings it consistently outperforms the baseline, with reduced inference time, energy use, and reserved GPU memory at 3,000 generation tokens.

  • Problem

    Longer contexts increase self-attention computation and KV-cache memory demands, while existing soft-compression methods do not jointly support cross-architecture deployment, answer-targeted supervision, and query-guided selection.

    Source: §1 · pp. 1–2
  • Method

    CMC compresses contexts into embeddings projected for a frozen decoder, selects CMEs using the question alongside a local context window, and trains with answer-targeted distillation.

    Source: §1 · p. 2
  • Results

    Across nine encoder-decoder combinations and four extractive QA benchmarks, CMC consistently outperforms the baseline in most controlled-data configurations; at T = 3,000 tokens, it reduces inference time and energy by 20% and peak reserved GPU memory by up to 62.5%.

    Source: §6 · pp. 8–9
  • Takeaways & Limitations

    At T = 3,000 tokens, the reported efficiency gains enable deployment at generation lengths that would otherwise exhaust GPU memory.

    Source: §6 · p. 9

    Evaluation is limited to English extractive QA benchmarks; effectiveness on abstractive QA, summarisation, and RAG has not been evaluated.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large language model (LLM) inference is constrained by the quadratic scaling of self-attention and the linear scaling of the KV cache, increasing latency, energy consumption, and GPU memory demand as context length scales. Existing soft-compression methods either lack query-guided memory selection at inference time, train without answer-targeted supervision, or couple compression tightly to a specific decoder architecture. We propose a Context-to-Answer-Aligned Memory Compression (CMC) framework, which compresses long input contexts into compact Context Memory Embeddings (CMEs) aligned to any frozen decoder's embedding space, reducing inference costs without modifying decoder weights. CMC introduces a two-tier KV cache that combines question-guided CME selection with a local context window, and trains the compressor with answer-targeted distillation from a frozen LLM. Experiments across nine encoder-decoder combinations and four QA benchmarks show that CMC consistently outperforms the baseline, achieving up to 7.3 EM and 4.0 F1 point gains on SQuAD, while reducing inference time and energy consumption by up to 20% and peak reserved GPU memory by up to 50% at 3,000 generation tokens. Ablation studies confirm that each architectural component and training objective contributes to the performance.

1 Introduction

Long-context inference faces quadratic self-attention cost and growing KV-cache demands, while existing soft-compression methods do not jointly provide cross-architecture use, answer-targeted training, and query-guided selection. CMC compresses contexts into decoder-aligned memory embeddings and combines selected memory with a local context window, trained with answer-targeted supervision.

1 Introduction

  • Architectural methods require substantial changes, hard-prompt compression discards tokens, and existing soft-compression methods do not jointly provide cross-architecture deployment, answer-targeted training, and query-guided selection.Hard-prompt compression can lose information needed for evidence synthesis across distant passages.Source: §1 · pp. 1–2
  • CMC projects compact Context Memory Embeddings into a frozen decoder’s embedding space and combines question-selected CMEs with a full-resolution local context window.The two-tier KV cache is designed to provide global semantic coverage and local span-level precision without full-context attention.Source: §1 · p. 2
  • CMC trains the ContextEncoder and MemoryBridge in two phases: decoder-alignment objectives followed by answer-targeted distillation and contrastive memory-answer alignment.Phase 2 uses cross-entropy, KL divergence, and contrastive memory-answer alignment losses with a frozen DecoderLLM.Source: §1 · p. 2
  • CMC consistently outperforms the baseline in EM and F1 across nine encoder-decoder combinations and four QA benchmarks, while reducing inference time, energy consumption, MACs, and peak GPU memory.The experiments cover SQuAD, AdversarialQA, HotpotQA, and CovidQA.Source: §1 · pp. 2–3

2 Methods

CMC compresses a long context into decoder-ready memory embeddings for a frozen DecoderLLM, with only the ContextEncoder and MemoryBridge trainable. Its methods combine context denoising, two-phase training, and question-guided memory selection with a local context window.

2 Methods

  • The ContextEncoder and MemoryBridge alone are trainable; the DecoderLLM remains frozen throughout the method.Source: §2 · p. 3

2.1 Context Denoising, Segmentation, and Memory Tokens

  • Tokens lacking sufficient cosine similarity to any of their next g neighbors are discarded; the retained context is segmented and augmented with fixed memory placeholders.The compression rate r = t/m(c) trades a more compact representation against reduced context coverage.Source: §2.1 · p. 3

2.2 CME Extraction with ContextEncoder

  • Causal self-attention lets each segment’s placeholder tokens attend to preceding chunk tokens, yielding hidden-state CMEs that are concatenated into the full context memory.The total memory contains Mb = ns × m(c) vectors.Source: §2.2 · p. 3

2.3 Cross-Architecture Projection with MemoryBridge

  • MemoryBridge projects CMEs between encoder and decoder dimensions, while norm calibration bounds each projected vector to at most twice the decoder’s average token-embedding norm.Source: §2.3 · pp. 3–4

2.4 Two-Phase Training

  • Phase 1 aligns memory through full-context and first-half reconstruction; Phase 2 uses answer-token cross-entropy, educator-distribution KL, and contrastive memory-answer alignment.Removing any Phase-2 loss component degrades performance to the level of Phase 1 alone.Source: §2.4 · p. 4

2.5 Inference with Two-Tier KV Cache Policy

  • Inference selects top-K CMEs by question similarity, reorders them by document position, and combines them with a local context window; cache length stays bounded by K + W + |q| regardless of document length.The local window is centered on the starting token of the answer span, subject to context boundaries.Source: §2.5 · pp. 4–5

3 Experiments

Experiments evaluate CMC across four extractive QA benchmarks, nine encoder–decoder combinations, and two data settings. CMC improves task performance in most configurations while reducing inference costs, with results varying by decoder and compression rate.

3.1 Datasets

  • Evaluation covers single-hop, adversarial, multi-hop, and biomedical QA, with CovidQA tested only in the controlled-data setting because of its limited corpus size.The controlled-data setting uses 10,000 training and 1,000 validation samples for three datasets, and 1,500 training and 200 validation samples for CovidQA.Source: §3.1 · p. 5

3.2 Experimental Settings

  • The evaluation pairs three ContextEncoders with three frozen DecoderLLMs and compares CMC with a controlled local-window baseline and published compression methods.Task metrics are EM and F1; efficiency measures include inference time, energy, MACs, and peak allocated and reserved GPU memory.Source: §3.2 · p. 5, §3.3 · p. 5

3.4.1 Task Performance

  • Controlled-data CMC improves Llama EM by 4.6, 2.5, 5.5, and 7.0 points on SQuAD, AdversarialQA, HotpotQA, and CovidQA, respectively, averaging +4.9 EM.For full-data SQuAD, Llama gains 7.28 EM (48.28→55.56) and Gemma gains 3.22 EM (40.91→44.13); Mistral has mixed results, including a 14.7 EM regression on SQuAD.Source: §3.4.1 · pp. 5–6, Table 1 · p. 6
  • CMC achieves the best F1 against published baselines on all three full-data datasets and leads HotpotQA EM, 47.77 versus 39.97 for PCC-Large.CMC underperforms PCC-Large on EM for SQuAD and AdversarialQA.Source: Table 2 · p. 6, §3.4.1 · p. 6

3.4.2 Generation and Use of CMEs

  • Compression rate r = 4 achieves the best or tied-best EM in 6 of 9 combinations; r = 2 is consistently lowest, while r = 8 can lose context.Llama favors r = 4 by 1.1–1.7 EM points over r = 8, while Mistral and Gemma differ by at most 0.1 EM between those rates.Source: §3.4.2 · p. 6, Figure 2 · p. 7

3.4.3 Computational Efficiency

  • At 1,000 samples, CMC reduces inference time by 4.6% at 512 tokens and 20.0% at 3,000 tokens, with energy reductions of 4.6% and 20.3%, respectively.Compression overhead is approximately 25 seconds regardless of token budget; decode savings grow with generation length.Source: Figure 3 · p. 7, §3.4.3 · p. 7
  • At 3,000 tokens, CMC’s peak reserved memory is 9.8 GB versus the baseline’s 19.7 GB, a 50% reduction; Mistral shows a 62.5% reduction.Peak allocated memory for CMC remains approximately 8.85 GB as the token budget increases.Source: Figure 4 · p. 7, §3.4.3 · p. 7

4 Ablation Study

Ablations on controlled-data SQuAD show substantial losses when key architectural components or Phase-2 objectives are removed. The results distinguish inference-policy effects from degraded CME quality.

Architectural ablations

  • Removing Tier-2 causes the largest architectural drop, −50.0 EM, while near-identical AE/AR losses to CMC indicate the deficit arises at inference.The ablation setting is SQuAD, GPT2-Large, r = 4, and Llama.Source: §4 · pp. 7–8, Table 3 · p. 8
  • Removing graph-based denoising reduces EM by 26.6 points relative to CMC and puts it 22.0 points below the baseline; random CME selection reduces EM by 8.1 points.Random selection keeps the number of CME slots fixed but removes their question alignment.Source: §4 · p. 7, Table 3 · p. 8

Phase-2 training objectives

  • Removing either LKL or LCL individually is as damaging as removing Phase-2, with all three objectives needed jointly.LCE provides token-level supervision, LKL aligns the learner with the educator distribution, and LCL grounds the CME prefix in answer-embedding space.Source: §4 · p. 8

Additional analyses

  • Appendix analyses examine local-window size, denoising threshold, and Tier-2 localization without gold annotations.Source: §4 · p. 8

5 Related Work

Related work includes architectural approaches, hard-prompt compression, and soft-compression methods. The paper identifies cross-architecture deployment, answer-targeted supervision, and query-guided inference-time selection as gaps addressed together by CMC.

Architectural methods

Hard-prompt compression

  • Hard-prompt methods prune tokens and are model-agnostic, but discarded information cannot be recovered, creating vulnerability on cross-passage evidence-synthesis tasks.Examples include Selective Context (Li et al., 2023), LLMLingua-2 (Pan et al., 2024), and query-guided QGC (Cao et al., 2024).Source: §5 · p. 8

Soft-compression

  • Soft-compression methods encode contexts into dense embeddings, but the paper identifies no existing method combining cross-architecture deployment, answer-targeted training, and query-guided memory selection at inference.Dai et al. (2025) provide cross-architecture projection without the latter two features, while ATACompressor (Li et al., 2025a) adds query-guided selective encoding without cross-architecture deployment or answer-targeted alignment.Source: §5 · p. 8

6 Conclusion

CMC compresses long contexts into compact embeddings for inference with frozen LLMs. Across four QA benchmarks and nine encoder–decoder combinations, it usually outperforms the baseline while reducing compute, energy, and memory use.

Conclusion

  • Across nine encoder-decoder combinations and four extractive QA benchmarks, CMC outperforms the baseline in most controlled-data configurations and achieves the best F1 against published soft-compression baselines under full-data conditions.Source: §6 · pp. 8–9
  • At 3,000 tokens, CMC reduces inference time and energy by 20% and peak reserved GPU memory by up to 62.5%.Source: §6 · p. 9
  • Ablations show substantial performance degradation when graph-based denoising, Tier-2, or any of the three Phase-2 objectives is removed.Source: §6 · p. 9

Limitations

The evaluation is limited to English extractive QA benchmarks, with no cross-dataset evaluation. The fixed compression rates may not suit all context lengths.

Limitations

  • CMC is evaluated only on extractive QA; abstractive QA, summarisation, and RAG remain unevaluated.The authors note those tasks may require adjustments to the answer-targeted distillation objective.Source: Limitations · p. 9
  • Cross-dataset generalization is untested because CMC is trained and evaluated independently within each benchmark.The authors give training on SQuAD and evaluation on HotpotQA as an example.Source: Limitations · p. 9
  • The fixed compression rates r ∈{2, 4, 8} may over-compress short passages or under-compress long ones.Adaptive rate scheduling conditioned on context length is left for future work.Source: Limitations · p. 9
  • Evaluation covers English-language QA only; effectiveness on non-English or multilingual contexts has not been investigated.The authors note that tokenisation and semantic similarity properties differ across these settings.Source: Limitations · p. 9

Ethics Statement

The study did not require ethical approval.

Ethics Statement

A.1 Hyperparameters

The ContextEncoder and MemoryBridge are jointly optimised, while the DecoderLLM remains frozen; two-tier cache settings are tuned for each dataset and model configuration.

A.1 Hyperparameters

  • The ContextEncoder and MemoryBridge are optimised jointly, while the frozen DecoderLLM receives no gradient updates.Source: §A.1 · p. 11
  • The two-tier KV cache hyperparameters K and W are tuned per dataset and encoder-decoder configuration.The supplied passages do not state the specific tuned values.Source: §A.1 · p. 11
  • MemoryBridge LayerNorm uses PyTorch defaults, with learnable per-dimension scale and shift parameters trained jointly with the bridge.The stated epsilon is 10^-5.Source: §A.1 · p. 12

A.2 Detailed Results per Dataset

Across the detailed results and analyses, CMC's outcomes vary little with encoder scale or decoder in several settings, while cache efficiency and localization show distinct trade-offs. The appendices also examine compression rate, window size, denoising, and example-level behavior.

A.6 Local Window Size Sensitivity

  • EM rises from 0.140 at W = 50 to 0.354 at W = 200, then plateaus at 0.355 at W = 256.The default W = 130 gives EM = 0.304; the figure passage also reports EM and F1 rising monotonically before plateauing at W ≥200.Source: §A.6 · p. 13, Figure 5 · p. 13

A.7 Denoising Threshold (τ) Sensitivity

  • 19.9-point EM loss without denoising contrasts with identical EM = 0.507 and F1 = 0.739 for τ ≥0.02.At τ ≥0.02 the keep_min_every = 16 floor binds; the default τ = 0.05 lies in the stable region.Source: §A.7 · p. 14, Figure 6 · p. 15

A.2 Detailed Results per Dataset

  • EM = 0.4157 and F1 = 0.6678 at r = 4 lead the overall averages; CovidQA improves across all nine configurations with 1,500 training samples.OPT encoders produce near-identical results, GPT2-Large matches or outperforms them, and Gemma varies by ≤0.001 EM across encoder and compression-rate combinations.Source: §A.2 · p. 12, Table 9 · p. 15

A.3 Inference Efficiency

  • ≈20% inference-time savings at T = 3,000 for Mistral match Llama; compression overhead scales with sample count, not generation length.Prefill time is marginally lower than baseline across the reported conditions.Source: §A.3 · p. 12, Table 12 · p. 17

A.4 Peak GPU Memory

  • At T = 3,000, peak reserved memory reaches 19.7 GB for Llama and 24.3 GB for Mistral with the baseline, versus CMC's near-constant ≈9.8 GB and ≈9.1 GB.CMC's reserved-memory footprint stays near constant across sample sizes and token budgets.Source: §A.4 · p. 12, Table 14 · p. 18, Table 15 · p. 19

A.5 Attention MACs Analysis

  • 27.0% average attention-MAC reduction (range 6.6%–51.3%) follows from CMC's bounded K + W cache, with larger average savings on SQuAD.The SQuAD average reduction is 34.0%; shorter baseline prompts explain smaller reductions on some HotpotQA configurations.Source: §A.5 · p. 12, Table 16 · p. 19

A.8 Tier-2 Window Localization Analysis

  • Oracle-free localization retains 78–99% of oracle performance across datasets; on SQuAD, BM25 costs 0.4 EM points and CME-guided localization costs 2.0.HotpotQA alternatives retain 97–98%; AdversarialQA has the largest oracle dependence, at −0.052 to −0.069 EM.Source: §A.8 · p. 15, Table 17 · p. 19

A.9 Case Study

  • CMC distinguishes £76 million from an earlier £43 million figure, but compression loses the month from “June 1979”; it also improves partial-answer F1 from 0.40 to 0.55.The cases illustrate both global-context benefits and possible loss of fine-grained temporal detail.Source: §A.9 · pp. 15–16, Table 18 · p. 20
Loading 2609.25537v1…