SoftLMs: Efficient Adaptive Low-Rank Approximation of Language Models using Soft-Thresholding Mechanism

TL;DR

This paper introduces SoftLMs, a novel method for compressing language models by adaptively determining layer ranks through soft-thresholding, significantly reducing parameters and increasing speed while maintaining performance.

Contribution

It presents a differentiable soft-thresholding technique for automatic low-rank approximation of language models, applicable to various architectures including BERT, GPT2, and TinyLlama.

Findings

Achieves 1.33X to 1.72X speed-up in models.

Reduces total parameters by approximately 50%.

Maintains competitive performance on multiple tasks.

Abstract

Extensive efforts have been made to boost the performance in the domain of language models by introducing various attention-based transformers. However, the inclusion of linear layers with large dimensions contributes to significant computational and memory overheads. The escalating computational demands of these models necessitate the development of various compression techniques to ensure their deployment on devices, particularly in resource-constrained environments. In this paper, we propose a novel compression methodology that dynamically determines the rank of each layer using a soft thresholding mechanism, which clips the singular values with a small magnitude in a differentiable form. This approach automates the decision-making process to identify the optimal degree of compression for each layer. We have successfully applied the proposed technique to attention-based architectures, including BERT for discriminative tasks and GPT2 and TinyLlama for generative tasks. Additionally, we have validated our method on Mamba, a recently proposed state-space model. Our experiments demonstrate that the proposed technique achieves a speed-up of 1.33X to 1.72X in the encoder/ decoder with a 50% reduction in total parameters.