Reproduce “Share Your Attention: Transformer Weight Sharing via Matrix-based Dictionary Learning” on OpenResearch

About this paper

Magauiya Zhussip, Dmitriy Shopkhoev, Ammar Ali, Stamatios Lefkimmiatis

A new method called Matrix Atom Sharing in Attention (MASA) reduces the parameter count of transformer attention modules by expressing projection matrices as linear combinations of shared "matrix atoms". This technique achieved up to 66.7% attention parameter compression, often outperforming uncompressed baselines and state-of-the-art compression methods across language models and Vision Transformers, even when adapting pretrained models without fine-tuning.

Key insights

• Inter-block redundancy in transformer attention matrices can be effectively exploited using a dictionary learning framework, leading to substantial parameter reduction without performance degradation.
• Using independent dictionaries for Query, Key, Value, and Output projections (Q, K, V, O) yields optimal performance, suggesting distinct functional roles and varying compressibility among these projections.
• A data-driven grouping strategy, based on KL divergence of layer-wise probability distributions, effectively identifies functionally similar blocks in pretrained models for efficient parameter sharing.

Summary by alphaXiv.