COMI: Coarse-to-fine Context Compression via Marginal Information Gain

AI summary

Overview

Research area: Natural Language Processing — long-context efficiency for Large Language Models, specifically task-aware context compression for retrieval-augmented and long-document pipelines.

Technical level: Advanced. The paper assumes familiarity with encoder-decoder architectures, attention, cosine-similarity-based retrieval metrics, and submodular greedy selection theory.

Scope in one sentence: COMI is a two-stage context compression framework that scores each unit of text by its "Marginal Information Gain" — query relevance minus semantic redundancy — and uses that score both to redistribute compression budgets across segments and to weight token merging within segments.

What This Paper Is About

LLMs spend most of their capacity on a tiny fraction of a long input, and existing compression methods either ignore the query (losing relevant content) or keep everything relevant (retaining many near-duplicate passages that mislead the model). COMI's goal is to compress context aggressively — up to 32x — while deliberately keeping information that is both relevant to the query and non-redundant with respect to the rest of the retained content.

Key Contributions

1. Marginal Information Gain (MIG). A new scoring metric defined as the cosine similarity between a unit and the query minus the maximum cosine similarity between that unit and any other unit in the context. It converts a relevance-only ranking into a relevance-minus-redundancy ranking, and the authors provide a theoretical appendix arguing that greedy selection under MIG yields higher mutual information with the target label than greedy selection under pure relevance, assuming approximate submodularity.

2. Coarse-Grained Group Reallocation. The context is split into equal-length, non-overlapping groups. An inter-group MIG is computed per group by first selecting a representative vector (the token in the group most similar to the average-pooled query), then assigning each group an output length that is inversely related to its MIG, so low-redundancy, high-relevance groups receive a lower compression rate.

3. Fine-Grained Token Merging. Within each group, every token receives an intra-group MIG score, and tokens are fused into a single compressed token via a softmax over those MIG scores — high-MIG tokens contribute more, so relevant-but-redundant tokens are attenuated rather than accumulated.

4. Extensive evaluation across tasks and backbones. Experiments span four QA datasets and one summarization dataset, two backbones (LLaMA-2-7B-Chat, Qwen2-7B-Instruct) plus Qwen3-4B-Instruct, under 16x and 32x compression constraints, with ablations, a compression pressure test, and a latency study.

Main Findings

• Large gains on QA. Under a 32x compression constraint with Qwen2-7B-Instruct, COMI reaches 49.15 EM / 49.59 F1 on NaturalQuestions, 48.89 / 57.07 on 2WikiMQA, 40.46 / 55.24 on HotpotQA, and 11.18 / 22.30 on NarrativeQA. The paper reports this as approximately a 25-point Exact Match improvement over the suboptimal baseline on NaturalQuestions; the corresponding table row shows GMSA at 24.63 EM on that dataset. On 2WikiMQA and HotpotQA at 32x, COMI's EM is roughly 20 points above the best listed baseline.

• 16x results. With Qwen2-7B-Instruct at 16x compression, COMI scores 56.31 EM / 55.52 F1 on NaturalQuestions, 52.13 / 60.43 on 2WikiMQA, 45.11 / 60.52 on HotpotQA, 13.25 / 24.69 on NarrativeQA, and 35.85 F1 on MultiNews. Notably, the closed-book and Original Prompt baselines for Qwen2-7B on NaturalQuestions report EM 72.35 with F1 38.01 for the full prompt, so COMI's advantage is concentrated in F1 and in high-compression regimes rather than in raw EM against the uncompressed prompt.

• LLaMA-2-7B results. At 16x, COMI scores 22.75 EM on NaturalQuestions versus 15.04 for the Original Prompt; at 32x it scores 16.99. LLaMA-2-7B lags the Qwen2 backbone across the board.

• Pressure test. As the compression constraint rises from 2x to 32x, both COMI and Activation Beacon degrade, but COMI stays consistently higher; the paper reports nearly a 40-point higher EM than Activation Beacon on NaturalQuestions at 32x.

• Every component matters (ablation, NaturalQuestions / 2WikiMQA, 32x, Qwen2-7B). Default: 56.31/55.52 and 52.13/60.43. Removing coarse-grained group reallocation: 54.54/54.26 and 49.82/58.06. Removing fine-grained token merging: 50.81/50.82 and 47.88/56.45. Removing coarse-level redundancy (relevance-only reallocation): 54.59/52.81 and 50.14/58.64. Removing fine-level redundancy: 51.53/52.89 and 49.50/57.87. The reported Default row (56.31 EM on NaturalQuestions) matches the 16x main-table value rather than the 32x value (49.15), which the paper does not explain.

• Latency (32x, seconds). On NarrativeQA: End-to-end 7.04 for Original Prompt, 2.10 for SnapKV, 3.93 for Activation Beacon, 70.95 for LongLLMLingua, 2.10 for GMSA, 3.27 for COMI. On MultiNews: 8.58 for Original Prompt, 5.03 for SnapKV, 7.94 for Activation Beacon, 8.25 for LongLLMLingua, 3.86 for GMSA, 4.09 for COMI. COMI's compression stage costs 2.76s on NarrativeQA and 0.66s on MultiNews. The paper states COMI achieves more than a 2x end-to-end speedup over the Original Prompt at 32x compression on both QA and summarization.

• Native long-context backbones also improve. With Qwen3-4B-Instruct (native 256K context), F1 at 16x is 34.79 on NaturalQuestions (versus 16.90 for the full original prompt), 45.01 on 2WikiMQA (22.51), 44.44 on HotpotQA (34.16), and 15.71 on NarrativeQA (11.35). At 32x the corresponding figures are 28.89, 41.19, 39.71, and 13.25.

• Attention is extremely concentrated. The paper's motivating analysis reports that the 0.75% most query-relevant tokens occupy 99% of attention weights, and that the top query-related tokens are highly similar to each other, with the lowest similarity exceeding 0.6.

Methodology in Plain English

COMI starts from an encoder-decoder setup, where the encoder reads the original context and produces hidden states. The context is chopped into equal-length groups, and the query is collapsed into one average-pooled vector. For each group, COMI picks the token most similar to that query vector as the group's representative, then scores the group by its relevance to the query minus how similar it is to other groups' representatives — the inter-group MIG. Groups scoring higher get allocated a larger share of the allowed output length, using a softmax over the negated MIG values so that high-gain groups are compressed less. Within each group, every token gets a similar MIG score (relevance to the query minus similarity to its most similar in-group peer), and the group's tokens are merged into one vector using a softmax over those scores as weights. The merged vectors are concatenated, passed through a single-layer Layer Semantic Alignment module borrowed from GMSA to bridge encoder and decoder semantics, and the decoder generates the answer. Training is a single run of standard negative log-likelihood on 20,000 examples each from NaturalQuestions, HotpotQA, 2WikiMQA, NarrativeQA, and MultiNews, with all sequences at or under 32K tokens, batch size 64, learning rate 1e-5, and linear decay. The encoder and LSA are fully fine-tuned; the decoder is partially fine-tuned (only its W_Q, W_K, W_V, W_O matrices). Training randomly samples a compression rate per example. GMSA was retrained on the same data and settings for fairness, and all experiments ran on 8 NVIDIA H20 (94GB) GPUs.

Why This Matters

Research impact: COMI reframes compression as a relevance-and-diversity problem rather than a relevance-only problem, and offers MIG as a reusable scoring criterion that could be plugged into other selection-based pipelines (retrieval reranking, KV-cache eviction, memory selection). Its ablation shows that removing redundancy awareness at either granularity costs accuracy, which is a concrete argument against the conditional-independence assumption common in task-aware compression.

Real-world applications:

• Retrieval-augmented generation over large document collections, where retrieved chunks are long and overlapping.
• Long-document question answering, including narrative and legal or financial filings tested here via NarrativeQA.
• Multi-document summarization, evaluated on MultiNews.
• Cost and latency reduction for LLM serving, where the paper reports more than 2x end-to-end speedup over the uncompressed prompt at 32x compression.

Industry relevance: Compression that keeps accuracy at 32x reduces the token budget and attention cost of every inference call, and COMI's encoder-decoder design is model-agnostic enough to be applied across different backbones (demonstrated on LLaMA-2-7B, Qwen2-7B, and Qwen3-4B). The work was supported by Alibaba Group through the Alibaba Research Intern Program, and code is slated for release at https://github.com/Twilightaaa/COMI.

Future Directions

• Extending beyond QA and summarization. The evaluation covers four QA datasets and one summarization dataset; whether MIG-based compression preserves performance on reasoning, code, or multi-turn dialogue is not reported.
• Broadening backbone coverage. Results are reported for LLaMA-2-7B, Qwen2-7B, and Qwen3-4B; behavior on larger models and other architectures is not reported.
• Combining with KV-cache compression. The paper frames KV-cache methods as a separate family with different constraints (identical compression and response models, model-specific modifications), but does not test a hybrid of COMI-style selection with cache eviction.
• Refining the redundancy measure. Redundancy is defined as the maximum pairwise cosine similarity to another unit; alternatives such as set-level diversity or learned redundancy estimators are not explored.
• Human evaluation. All reported metrics are automatic (EM, F1); no human judgment of compressed-context faithfulness is reported.

Target Audience

Researchers and engineers working on long-context LLM inference, retrieval-augmented generation, and prompt or context compression. It is most useful to readers who already understand attention mechanics and embedding similarity, since the method and its theoretical appendix lean on submodularity and mutual-information arguments. Practitioners optimizing inference cost for document QA or summarization pipelines will find the latency and compression-rate tables directly actionable; readers seeking an introductory treatment of context compression would need background reading first.

Authors’ abstract

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.