MOOSComp: Improving Lightweight Long-Context Compressor via Mitigating Over-Smoothing and Incorporating Outlier Scores

TL;DR

MOOSComp tackles the inefficiency of processing long contexts by introducing a task-agnostic, hard-prompt compressor built on a BERT-based encoder. It mitigates token over-smoothing with an inter-class cosine similarity loss and preserves rare but important tokens through an outlier-score mechanism, forming a robust compression strategy. Across multiple benchmarks and target models, MOOSComp outperforms prior task-agnostic methods and delivers significant speedups on mobile and edge devices, enabling practical long-context understanding and reasoning in resource-constrained environments.

Abstract

Recent advances in large language models have significantly improved their ability to process long-context input, but practical applications are challenged by increased inference time and resource consumption, particularly in resource-constrained environments. To address these challenges, we propose MOOSComp, a token-classification-based long-context compression method that enhances the performance of a BERT-based compressor by mitigating the over-smoothing problem and incorporating outlier scores. In the training phase, we add an inter-class cosine similarity loss term to penalize excessively similar token representations, thereby improving the token classification accuracy. During the compression phase, we introduce outlier scores to preserve rare but critical tokens that are prone to be discarded in task-agnostic compression. These scores are integrated with the classifier's output, making the compressor more generalizable to various tasks. Superior performance is achieved at various compression ratios on long-context understanding and reasoning benchmarks. Moreover, our method obtains a speedup of 3.3x at a 4x compression ratio on a resource-constrained mobile device.

Figures (4)

• Figure 1: An overview of the proposed MOOSComp. Training Phase (Top): The compressor, containing a BERT-based encoder and a classifier in series, is trained for binary token classification with the cross-entropy loss and the inter-class cosine similarity loss to mitigate over-smoothing. Compression Phase (Bottom): Each token's classification score is combined with an outlier score to enhance the preservation of rare but important tokens.
• Figure 2: Inter-class cosine similarity of different layers in xlm-roberta-large and in llmlingua-2-xlm-roberta-large-meetingbank.
• Figure 3: Inter-class cosine similarity of different layers in xlm-roberta-large and in MOOSComp with varying values of $\beta$.
• Figure 4: Examples of the original and compressed prompts from the LongBench dataset, with the preserved words highlighted in red. Left: Results of LLMLingua-2. The number above each word indicates the probability of that word being preserved. Right: Results of MOOSComp. The two numbers above each word represent the probability of that word being preserved (up) and the outlier score of that word (down), respectively.