EcoSpa: Efficient Transformer Training with Coupled Sparsity

TL;DR

Transformer training is computationally intensive; EcoSpa tackles this by treating coupled weight matrices as first-class citizens and sparsifying them in aligned pairs through coupled estimation and sparsification. It leverages empirical Fisher information to assess the joint importance of matrix pairs and employs a tSVD-inspired, aligned removal strategy to preserve interaction patterns in both MHA and FFN (including extensions to GQA and RoPE). The method yields substantial system-level gains across architectures (LLaMA, GPT-2, DeiT), including memory reductions, faster training, and improved perplexity and throughput, all using standard PyTorch operations on commodity hardware. This makes efficient transformer training more accessible across pre-training and fine-tuning phases while preserving performance.

Abstract

Transformers have become the backbone of modern AI, yet their high computational demands pose critical system challenges. While sparse training offers efficiency gains, existing methods fail to preserve critical structural relationships between weight matrices that interact multiplicatively in attention and feed-forward layers. This oversight leads to performance degradation at high sparsity levels. We introduce EcoSpa, an efficient structured sparse training method that jointly evaluates and sparsifies coupled weight matrix pairs, preserving their interaction patterns through aligned row/column removal. EcoSpa introduces a new granularity for calibrating structural component importance and performs coupled estimation and sparsification across both pre-training and fine-tuning scenarios. Evaluations demonstrate substantial improvements: EcoSpa enables efficient training of LLaMA-1B with 50\% memory reduction and 21\% faster training, achieves $2.2\times$ model compression on GPT-2-Medium with $2.4$ lower perplexity, and delivers $1.6\times$ inference speedup. The approach uses standard PyTorch operations, requiring no custom hardware or kernels, making efficient transformer training accessible on commodity hardware.

Figures (4)

• Figure 1: Overview of EcoSpa. Unlike existing methods that independently sparsify weight matrices, EcoSpa jointly estimates the importance of coupled weight matrices and removes aligned row/column pairs, preserving coupling relationships. The approach is applicable to both pre-training and fine-tuning scenarios.
• Figure 2: EcoSpa Method Overview. (Step 1) Coupled Estimation: Jointly evaluate importance of coupled weight matrix pairs using empirical Fisher information. (Step 2) Coupled Sparsification: Remove aligned row/column pairs from less important coupled matrices, preserving dimensional compatibility and interaction patterns. The method ensures that when a column is removed from $\bm{W}_1$, the corresponding row is also removed from $\bm{W}_2$.
• Figure 3: Pre-training DeiT-Small model on CIFAR-10 dataset using different $\bm{W}_{couple}$ compression methods. All the methods remove the same number of parameters within the same $\bm{W}_{couple}$'s in each epoch.
• Figure :

Theorems & Definitions (9)

• Definition 1: Coupled Sparsification
• Proposition 1: Error under Coupled Sparsification
• proof
• Remark 1
• Definition 2: Individual Sparsification
• Proposition 2: Error under Individual Sparsification
• Remark 2: Mismatch Effect
• Proposition 3: Probabilistic Error Comparison
• proof : Argument Sketch