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

TL;DR

COMI introduces Marginal Information Gain (MIG) to jointly optimize query relevance and semantic diversity in long-context compression. It employs a two-stage, coarse-to-fine strategy: inter-group MIG-guided group reallocation and intra-group MIG-guided token merging within an encoder-decoder architecture enhanced by Layer Semantic Alignment (LSA). Empirical results across QA and summarization benchmarks show large gains under high compression (e.g., ~25 EM points at 32x), and improvements on native long-context LLMs, demonstrating the practical value of dynamic, query-aware compression for efficient long-context reasoning. The work highlights MIG as a robust criterion for preserving informative, non-redundant content, enabling more effective use of limited context length in real-world applications.

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.

Figures (4)

• Figure 1: Analysis of Attention Distribution and Similarity of Top Query-Related Tokens. (a) Only a small number of tokens related to the query occupy a large proportion of the attention weight allocation; for example, the 0.75% most relevant tokens occupy 99% of the attention weights. (b) These query-related tokens are highly similar to each other, with the lowest similarity exceeding 0.6.
• Figure 2: The Training Paradigm of COMI. COMI is based on an encoder-decoder architecture. The original context $X$ and query $Q$ are first encoded into hidden states, which are then compressed via a compression process (see Figure \ref{['fig:comi']}). The compressed representation is decoded and trained using cross-entropy loss. During training, the encoder and LSA are fully fine-tuned, while the decoder is partially fine-tuned, updating only the $W_Q$, $W_K$, $W_V$, and $W_O$ matrices in each layer.
• Figure 3: The Compression Process of COMI. Specifically, it sequentially performs three steps: I. Average Pooling of Query Tokens. Obtain a single query vector via average pooling; II. Coarse-Grained Group Reallocation. Reallocate the sizes of compression groups based on inter-group Marginal Information Gain (MIG) (i.e., groups with higher MIG are assigned lower compression rates); III. Fine-Grained Group Token Merging. Compute the intra-group MIG for each token and merge all tokens within a group into a single compressed token according to weights $w_1, ..., w_{L_{i} - 1}$.
• Figure 4: Compression Pressure Test on NaturalQuestions and 2WikiMQA. As the compression constraint increases, although both COMI and Activation Beacon generally show a downward trend in EM and F1, COMI consistently remains higher than Activation Beacon.

Theorems & Definitions (4)

• Lemma A.1: Information Preserved by Pure Relevance vs. MIG
• Proof 1
• Theorem 1: Superiority of MIG under Redundancy
• Proof 2