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

2025-04-23

AI Summary

To address the high inference overhead and weak discriminative capability of lightweight compressors in long-context reasoning under resource-constrained settings, this paper proposes a lightweight token-classification-based context compression method. We innovatively introduce an inter-class cosine similarity loss to mitigate representation over-smoothing and design a dynamic outlier scoring mechanism to preserve critical rare tokens, enabling joint optimization across both training and inference stages. Built upon the BERT architecture, our method supports end-to-end differentiable training. Extensive evaluations on multiple long-context understanding and reasoning benchmarks demonstrate significant improvements over state-of-the-art approaches: at a 4× compression ratio, it achieves a 3.3× inference speedup on mobile devices and a 12.7% gain in classification accuracy, while maintaining robust compression performance and strong task generalization ability.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Data Mining & Knowledge Management: Data Compression
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Mitigating over-smoothing in BERT-based compressors
• Incorporating outlier scores to preserve critical tokens
• Improving long-context compression efficiency and speed

Innovation

• Token-classification-based long-context compression
• Inter-class cosine similarity loss term
• Outlier scores for critical token preservation