GMSA: Enhancing Context Compression via Group Merging and Layer Semantic Alignment

TL;DR

GMSA introduces a novel context compression framework for long-context LLMs, integrating Group Merging and Layer Semantic Alignment to produce compact, semantically rich soft tokens. The training pipeline combines autoencoder training with Knowledge Extraction Fine-tuning to ensure complete semantics and task-specific knowledge extraction, while randomizing compression rates during training. Empirical results demonstrate that GMSA achieves substantial end-to-end speedups (approx. 2x) and outperforms strong baselines on context restoration and multi-document QA tasks, including notable improvements in Exact Match on NaturalQuestions under high compression. The approach offers a scalable, parallelizable solution that preserves semantic integrity with few encoder layers, addressing both efficiency and accuracy in long-context applications.

Abstract

Large language models (LLMs) have achieved impressive performance in a variety of natural language processing (NLP) tasks. However, when applied to long-context scenarios, they face two challenges, i.e., low computational efficiency and much redundant information. This paper introduces GMSA, a context compression framework based on the encoder-decoder architecture, which addresses these challenges by reducing input sequence length and redundant information. Structurally, GMSA has two key components: Group Merging and Layer Semantic Alignment (LSA). Group merging is used to effectively and efficiently extract summary vectors from the original context. Layer semantic alignment, on the other hand, aligns the high-level summary vectors with the low-level primary input semantics, thus bridging the semantic gap between different layers. In the training process, GMSA first learns soft tokens that contain complete semantics through autoencoder training. To furtherly adapt GMSA to downstream tasks, we propose Knowledge Extraction Fine-tuning (KEFT) to extract knowledge from the soft tokens for downstream tasks. We train GMSA by randomly sampling the compression rate for each sample in the dataset. Under this condition, GMSA not only significantly outperforms the traditional compression paradigm in context restoration but also achieves stable and significantly faster convergence with only a few encoder layers. In downstream question-answering (QA) tasks, GMSA can achieve approximately a 2x speedup in end-to-end inference while outperforming both the original input prompts and various state-of-the-art (SOTA) methods by a large margin.

Figures (7)

• Figure 1: Traditional Compression Paradigm v.s. Compression by GMSA. (a) visualizes the attention matrix when processing Life and love. <CT1> <CT2>, where <CT1> and <CT2> are randomly initialized tokens. Original shows the attention changes during processing of life and love. (b) represents the traditional compression paradigm. It first learns summary vectors in an autoregressive manner layer by layer, and then completes coarse-grained semantic alignment through a multi-layer perceptron (MLP), where $N_{\text{Enc}}$ denotes the number of encoder layers. (c) denotes the compression paradigm of GMSA, which first learns summary vectors through group merging and completes semantic alignment between different layers through the Layer Semantic Alignment (LSA) module.
• Figure 2: The Autoencoder Training Process of GMSA. GMSA consists of an encoder and a decoder, trained in an autoencoder manner using cross-entropy loss. GMSA first generates a set of summary vectors that meet the compression rate by performing group merging on the last hidden state of the encoder, and then achieves cross-layer semantic alignment through the Layer Semantic Alignment (LSA) module, which is composed of several Transformer blocks initialized with the weights of lower-layer decoder blocks. Remarkably, we find that using just a single layer of LSA can achieve excellent semantic preservation (see Appendix \ref{['apx:lsa_lays']}), hence $k_{\text{LSA}} << N_{\text{Dec}}$.
• Figure 3: The process of Knowledge Extraction Fine-tuning (KEFT). By fine-tuning only the $\boldsymbol{W}_Q$, $\boldsymbol{W}_K$, and $\boldsymbol{W}_V$ in the attention module of the decoder while keeping other modules frozen, the decoder performs teacher-forcing training using soft tokens $\tilde{m}$, question tokens, and the ground truth answer.
• Figure 4: GMSA-AE v.s. TCP-AE on the context restoration task. Sequence Length represents different context restoration lengths (i.e., 128, 256, 512), and the models are trained with a maximum length of 512.
• Figure 5: Analysis of the Effectiveness of Different Encoder Layers. (a) represents the comparison of convergence speed between GMSA-AE with different encoder layers and TCP-AE. (b) denotes the impact of different encoder layers on the semantic retention of GMSA-AE. The average BERT Score F1 refers to the average F1 score across different context restoration lengths (i.e., 128, 256, and 512).