🤖 AI Summary
Large language models (LLMs) suffer from prohibitive computational and memory overhead, hindering deployment. Existing compression methods primarily focus on intra-layer optimization, overlooking cross-layer redundancy inherent in Transformer’s repetitive architectural structure.
Method: We propose MASA—the first lightweight method to introduce matrix dictionary learning for cross-layer sharing of attention projection matrices in Transformers. MASA learns a shared atomic basis matrix and reconstructs layer-wise weights via linear combinations, enabling structured parameter compression without architectural modification, knowledge distillation, or reliance on key-value caching.
Contribution/Results: On LLMs ranging from 100M to 700M parameters, MASA achieves a 66.7% parameter reduction while outperforming Grouped Query Attention (GQA) and low-rank baselines. In Vision Transformers (ViTs), it reduces attention parameters by two-thirds with no accuracy loss. This work is the first to empirically validate dictionary learning for cross-layer weight sharing, establishing a novel paradigm for efficient, pre-trained model compression.
📝 Abstract
Large language models (LLMs) have revolutionized AI applications, yet their high computational and memory demands hinder their widespread deployment. Existing compression techniques focus on intra-block optimizations (e.g. low-rank approximation, attention head pruning), while the repetitive layered structure of transformers implies significant inter-block redundancy - a dimension largely unexplored beyond key-value (KV) caching. Inspired by dictionary learning in CNNs, we propose a framework for structured weight sharing across transformer layers. Our approach decomposes attention projection matrices into shared dictionary atoms, reducing the attention module's parameters by 66.7% while achieving on-par performance. Unlike complex methods requiring distillation or architectural changes, MASA (Matrix Atom Sharing in Attention) operates as a drop-in replacement - trained with standard optimizers - and represents each layer's weights as linear combinations of shared matrix atoms. Experiments across scales (100M-700M parameters) show that MASA achieves better benchmark accuracy and perplexity than grouped-query attention (GQA), low-rank baselines and recently proposed Repeat-all-over/Sequential sharing at comparable parameter budgets. Ablation studies confirm robustness to the dictionary size and the efficacy of shared representations in capturing cross-layer statistical regularities. Extending to Vision Transformers (ViT), MASA matches performance metrics on image classification and detection tasks with 66.7% fewer attention parameters. By combining dictionary learning strategies with transformer efficiency, MASA offers a scalable blueprint for parameter-efficient models without sacrificing performance. Finally, we investigate the possibility of employing MASA on pretrained LLMs to reduce their number of parameters without experiencing any significant drop in their performance.