TLDR: MASA (Matrix Atom Sharing in Attention) is a novel framework that uses dictionary learning to share weights across Transformer layers, significantly reducing the attention module’s parameters (e.g., 66.7% reduction) while maintaining or improving performance. It’s a simple, plug-and-play solution compatible with standard training, applicable to both language and vision models, and effective for compressing pretrained models with minimal performance degradation.
Large language models (LLMs) have transformed artificial intelligence, but their significant computational and memory requirements often limit their widespread use. While many techniques exist to compress these models, most focus on optimizing individual parts of the model (intra-block optimizations), such as reducing the rank of matrices or pruning attention heads. However, the repetitive, layered structure of Transformers, which are the backbone of LLMs, suggests a lot of untapped redundancy across different layers (inter-block redundancy).
Inspired by dictionary learning, a concept from signal processing, researchers have introduced a new framework called Matrix Atom Sharing in Attention (MASA). This innovative approach systematically exploits the inter-block redundancy by enabling structured weight sharing across Transformer layers. Instead of each layer having entirely independent attention projection matrices (Q, K, V, O), MASA decomposes these matrices into shared ‘dictionary atoms’.
How MASA Works
At its core, MASA allows each layer’s weights to be represented as simple linear combinations of these shared matrix atoms. Imagine having a common vocabulary (the dictionary atoms) that different layers can draw from to construct their specific attention mechanisms. This method significantly reduces the number of parameters in the attention module, for instance, by as much as 66.7% (reducing parameters from 226.5 million to 75 million in a 700-million-parameter model) while maintaining comparable performance.
A key advantage of MASA is its simplicity. Unlike other complex compression methods that might require additional training steps like distillation or fundamental changes to the model’s architecture, MASA operates as a ‘drop-in replacement’. This means it can be integrated into existing training pipelines and trained with standard optimization techniques, eliminating the need for auxiliary components.
Also Read:
- Making Large AI Image Models Accessible: A Hierarchical Approach to Compression
- Keeping LLMs Sharp: General Samples Replay for Continual Learning
Performance and Versatility
Extensive experiments across various model scales (from 100 million to 700 million parameters) demonstrate that MASA achieves better benchmark accuracy and perplexity compared to other leading compression techniques like grouped-query attention (GQA), low-rank baselines, and recently proposed Repeat-all-over/Sequential sharing methods, all while operating within similar or even tighter parameter budgets.
The benefits of MASA extend beyond language models. When applied to Vision Transformers (ViT) for image classification tasks, MASA also matches performance metrics with a substantial 66.7% reduction in attention parameters. Furthermore, the research explores the application of MASA to large pretrained models, showing that it can reduce their parameter count with only a marginal drop in performance, making it highly practical for real-world deployment where fine-tuning might not be feasible.
The study also revealed interesting insights into which parts of the attention module are most compressible. It was found that compressing the Query (Q), Key (K), and Value (V) projections (MASA-QKV) often yields better results than compressing all four projections, including the Output (O) projection (MASA-QKVO). This suggests that Q, K, and V projections might contain more cross-layer redundancy, while the O projection plays a more specialized role that benefits from independent parameters.
By combining the principles of dictionary learning with Transformer efficiency, MASA offers a scalable blueprint for creating parameter-efficient models without sacrificing performance. This work highlights the power of inter-layer matrix decomposition and opens new avenues for future research in inter-layer compression methods. For more details, you can read the full research paper: Share Your Attention: Transformer Weight Sharing via Matrix-based Dictionary Learning.
