Source-linked AI summary
COMI: Coarse-to-fine Context Compression via Marginal Information Gain
Jiwei Tang, Shilei Liu, Zhicheng Zhang, Yujin Yuan, Libin Zheng, Wenbo Su, Bo Zheng
TL;DR
Long-context LLM deployment faces computational inefficiency and information redundancy. COMI addresses these problems with MIG-guided coarse-to-fine compression, and experiments report substantial gains, including approximately 25-point EM improvement at 32x compression on NaturalQuestions.
Problem
Long-context LLM deployment is hindered by computational inefficiency from long sequences and redundant information that can degrade performance.
Method
COMI defines MIG as query relevance minus semantic redundancy, then reallocates group compression and merges within-group tokens using MIG-based weighting.
Results
COMI outperforms existing methods across QA and summarization tasks, including approximately 25-point NaturalQuestions EM improvement over a suboptimal baseline at 32x compression with Qwen2-7B.
Takeaways & Limitations
MIG provides a criterion for prioritizing information that is both query-relevant and low redundant during high-rate context compression.
Takeaways & Limitations
Latency phases for SnapKV and Activation Beacon cannot be measured separately because compression and generation are coupled.
Abstract
from arXiv · showhide
Large Language Models (LLMs) have demonstrated exceptional capabilities across diverse tasks. However, their deployment in long context scenarios remains hindered by computational inefficiency and information redundancy. Context compression methods address these challenges by significantly reducing input length and eliminating redundancy. We propose COMI, a coarse-to-fine adaptive context compression framework that jointly optimizes for semantic relevance and diversity under high compression rates. We introduce Marginal Information Gain (MIG), a metric defined as the relevance of a unit to the input query minus its semantic redundancy with other units, guiding the compression process to prioritize information that is both relevant and low redundant. The framework operates in two stages: (1) Coarse-Grained Group Reallocation, where the context is partitioned into groups and dynamically assigned compression rates based on inter-group MIG, ensuring compression budgets align with information value distribution; and (2) Fine-Grained Token Merging, where tokens within each group are fused via an intra-group MIG-based weighting mechanism, thereby preserving key semantics while avoiding the accumulation of redundancy. Extensive experiments across question-answering (e.g., NaturalQuestions, 2WikiMQA, HotpotQA and NarrativeQA), summarization (e.g., MultiNews) with various backbones (e.g., LLaMA-2-7B, Qwen2-7B) show that COMI outperforms existing baselines by a large margin, e.g., approximately 25-point Exact Match (EM) improvement under 32x compression constraint with Qwen2-7B on NaturalQuestions.
1 INTRODUCTION
Long-context LLM use is limited by attention cost and redundant information. COMI addresses both by allocating compression according to marginal information value and merging tokens to preserve relevant, diverse content.
- Long contexts create quadratic attention costs and information redundancy that can degrade LLM performance.
- Task-agnostic compression can lose query-relevant information, while relevance-only task-aware methods may retain highly similar redundant content.
- MIG measures a unit’s query relevance minus its semantic redundancy, prioritizing information that is both relevant and low redundant.
- COMI reallocates compression across groups using inter-group MIG, then merges tokens within groups using intra-group MIG-guided weights.
- 32x compression with Qwen2-7B improved NaturalQuestions Exact Match by approximately 25 points over a suboptimal baseline.
2 PRELIMINARY: GMSA
COMI adopts Layer Semantic Alignment because its encoder-decoder architecture can create a semantic gap between high-level summaries and low-level input representations.
- COMI uses Layer Semantic Alignment to align high-level summary vectors with low-level original-input semantics across encoder-decoder layers.
- The method follows GMSA by setting LSA to a single layer.
3 RELATED WORK
Prior work includes task-agnostic and task-aware context compression, while COMI is presented within an encoder-decoder training framework and a multi-stage compression process.
- Task-Agnostic Context Compression Methods: Task-agnostic methods compress context without incorporating the input query, aiming to preserve broad semantic information.
- Task-Aware Context Compression Methods: Task-aware methods use the context and query to merge, summarize, or delete content relevant to downstream tasks.
- KV-cache Compression: KV-cache methods compress layer-wise caches through sharing, head reduction, or discarding less important key-value states, but require matching compression and response models.
- COMI’s training paradigm encodes context and query, compresses hidden states, decodes the compressed representation, and optimizes cross-entropy loss.
- Its compression process averages query tokens, reallocates group sizes using inter-group MIG, and merges within-group tokens using intra-group MIG weights.
4 COMI
COMI encodes the context, measures relevance and redundancy with MIG, reallocates compression across segments, and performs weighted token merging. Training uses compressed representations with an encoder, LSA, and selected decoder parameters fine-tuned.
- 4.1 MARGINAL INFORMATION GAIN: MIG for a token is defined as its cosine similarity to the query minus its maximum cosine similarity to another context token.
- 4.2 COARSE-GRAINED GROUP REALLOCATION: The context hidden states are divided into equal-length segments because fixed compression rates are suboptimal when segment relevance and redundancy differ.
- 4.2 COARSE-GRAINED GROUP REALLOCATION: COMI pools the query into one vector, selects each segment’s most query-relevant representative, and computes segment-level MIG.
- 4.2 COARSE-GRAINED GROUP REALLOCATION: Segments with higher MIG receive smaller compression rates, allocating more output length to information that is more informative and less redundant.
- 4.3 FINE-GRAINED TOKEN MERGING: Within each segment, tokens are merged into one output token using softmax MIG weights, giving greater contribution to tokens with higher information gain.
- 4.4 TRAINING OBJECTIVE: COMI is fine-tuned on compressed representations while training the encoder and LSA fully and updating selected decoder attention matrices.
5 EXPERIMENTS
Experiments evaluate COMI across QA, summarization, compression-pressure, ablation, efficiency, and native long-context settings. COMI remains competitive under severe compression while reducing generation cost and benefiting from all major components.
- Experimental Setup: COMI is evaluated on four QA benchmarks and MultiNews summarization using EM and F1 for QA and F1 for summarization.The experiments use LLaMA-2-7B and Qwen2-7B backbones, with baselines spanning hard prompt, soft prompt, and KV-cache compression methods.
- Compression Pressure: As compression increases from 2x to 32x, COMI and Activation Beacon generally decline, but COMI remains higher and is nearly 40 EM points ahead on NaturalQuestions at 32x.The comparison covers both EM and F1 on NaturalQuestions and 2WikiMQA.
- Ablation Study: Removing any COMI component causes clear drops in all metrics, while omitting reallocation, token merging, or redundancy handling respectively misallocates budgets, dilutes details, or retains redundant information.The ablation study examines coarse-grained reallocation, fine-grained token merging, and redundancy handling at both levels.
- Efficiency Analysis: At a 32x compression constraint, COMI achieves more than a 2× end-to-end speedup over the Original Prompt on both QA and summarization tasks.The analysis separates compression and generation latency; group reallocation and pooling add lightweight overhead because their inputs are small.
- Native Long-Context LLMs: With Qwen3-4B-Instruct, COMI outperforms the full Original Prompt across datasets at 16x and 32x compression; on NaturalQuestions at 16x, F1 is 34.79 versus 16.90.The backbone natively supports a 256K context length, so the comparison tests COMI beyond models lacking native long-context capability.
6 CONCLUSION
COMI combines Marginal Information Gain with coarse-to-fine compression to optimize relevance and semantic diversity under high compression rates, outperforming existing methods in extensive experiments.
- COMI jointly optimizes task relevance and semantic diversity under high compression rates.
- Marginal Information Gain penalizes redundancy while rewarding query relevance, guiding adaptive compression-budget reallocation and within-group token merging.
- COMI significantly outperforms existing methods, achieving up to a 25-point EM gain under 32x compression.
ETHICS STATEMENT
The work uses open-source data and models from open platforms and reports no additional ethical concerns beyond existing text-compression methods. It also notes that some phase latencies cannot be measured separately for SnapKV and Activation Beacon.
- The data and models are released under open-source licenses and sourced from open platforms.
- The authors report no additional ethical concerns compared with existing text-compression methods.
- SnapKV and Activation Beacon phase latencies cannot be measured individually because compression and generation are coupled.
A THEORETICAL ANALYSIS: MIG VS. PURE RELEVANCE
This analysis defines relevance, redundancy, and Marginal Information Gain (MIG) for token selection. MIG combines query relevance with a penalty for similarity to already selected tokens, favoring information that is both relevant and less redundant.
- Token relevance is measured by cosine similarity between a token embedding and the query vector.The token embeddings and query are assumed zero-mean and unit-norm.
- Token redundancy is measured by its maximum cosine similarity with any token already selected.This captures similarity to the most similar member of the selected set.
- MIG is introduced as a relevance criterion that explicitly accounts for redundancy with previously selected tokens.The framework is inspired by max-relevance min-redundancy selection.
- When no tokens have been selected, the redundancy term is set to 0, so MIG reduces to relevance.
- Sequential MIG selection subtracts the highest similarity to the selected context, penalizing tokens that duplicate existing information.
A.3 COMPARISON: PURE RELEVANCE VS. MIG
The comparison argues that pure relevance can select highly similar tokens, whereas MIG favors relevant but collectively diverse tokens. Under redundancy and approximate submodularity assumptions, this improves the preservation of mutual information with the target.
- Pure Relevance: Pure relevance selects tokens with the highest query correlation, even when those tokens are mutually redundant.
- MIG: MIG selects tokens that are highly relevant to the query while dissimilar to tokens already selected.
- Theoretical comparison: Under approximate submodularity and frequent relevance–redundancy co-occurrence, greedy MIG selection is theoretically superior under redundancy.
- Information overlap: Highly correlated tokens contribute diminishing additional information, with simplified Gaussian overlap approximately proportional to τ^2.
- MIG: MIG reduces the score of highly relevant tokens that duplicate an already selected token, discouraging redundant selections.
- Conclusion: MIG favors diverse, low-redundancy information and aims to preserve more mutual information with the target than relevance-only compression.
B EXPERIMENTAL EVIDENCE: MIG VS. PURE RELEVANCE
COMI’s diagnostic and baseline studies show that MIG identifies answer-critical segments more effectively and with less redundancy than pure relevance. Additional experiments show robust EM performance across compression rates, adaptive budget allocation, and scalability to 64K-token contexts.
- MIG versus pure relevance: 0.5809 AUC for MIG exceeds 0.5423 for relevance when predicting answer-containing segments on NaturalQuestions.The diagnostic evaluates whether representative-token scores identify segments containing the ground-truth answer.
- MIG versus pure relevance: MIG consistently reduces redundancy across different retention ratios compared with relevance-based selection.Redundancy is quantified using average pairwise cosine similarity among retained compressed embeddings.
- Baseline comparison: COMI achieves the highest Exact Match scores across 2× to 32× compression rates on NaturalQuestions and 2WikiMQA with Qwen2-7B-Instruct.The comparison includes SnapKV, LongLLMLingua, and LLMLingua-2.
- Adaptive group reallocation: At 32× compression, MIG-guided reallocation gives an informative segment a final size of 18 instead of its initial size of 32, while less informative segments expand.The example uses a 256-token context and reallocates group sizes according to inter-group MIG.
F DATASETS
The evaluation covers open-domain, multi-hop, long-form narrative, and multi-document summarization settings. These datasets test extraction, cross-document reasoning, long-context comprehension, and synthesis of concise low-redundancy summaries.
- Question answering datasets: NaturalQuestions evaluates open-domain question answering using authentic Google search queries paired with Wikipedia documents, including one ground-truth document and distractors.Answers may be short spans or longer passages, testing both precise extraction and long-document comprehension.
- Question answering datasets: HotpotQA requires multi-hop answers that synthesize information from multiple Wikipedia articles.Questions involve several entities and facts, requiring cross-document reasoning and linking.
- Question answering datasets: 2WikiMQA tests complex multi-hop reasoning across multiple Wikipedia documents, including comparison, contrast, and causal inference.Its questions require models to identify relevant facts and understand relationships among them.
- Summarization dataset: MultiNews is a multi-document summarization dataset whose news clusters require extracting, integrating, and synthesizing information into concise, coherent, low-redundancy summaries.Each cluster contains multiple articles reporting on the same event.
- Long-form comprehension: NarrativeQA evaluates comprehension and summarization over full-length novels and movie scripts, with questions requiring inference about plots, characters, and events.The dataset is designed for long-form understanding rather than isolated fact extraction.