CAST: Continuous and Differentiable Semi-Structured Sparsity-Aware Training for Large Language Models

TL;DR

This work introduces CAST (Continuous Adaptive Sparse Trainer), a fully continuous sparsity-aware training framework for semi-structured $N:M$ sparsity in large language models. It jointly optimizes masks and weights during training via three innovations: AdamS, a sparsity-aware optimizer with proportional $L_1$ decay; a learnable weight-scaling module to counteract magnitude loss; and knowledge distillation to bolster training efficiency. CAST achieves state-of-the-art performance under $2:4$ sparsity across model families (125M–13B), including near-lossless perplexity and improved zero-shot accuracy on LLaMA2-7B with only a fraction of pretraining data, and it supports practical deployment through quantization and speedups. The work also derives empirical scaling laws to predict sparse model performance and demonstrates robust downstream transfer with sparsity-aware fine-tuning, highlighting CAST’s practical impact for efficient, scalable LLM deployment.

Abstract

Sparsity-aware training is an effective approach for transforming large language models (LLMs) into hardware-friendly sparse patterns, thereby reducing latency and memory consumption during inference. In this paper, we propose Continuous Adaptive Sparse Trainer (CAST), a fully continuous and differentiable sparsity-aware training framework for semi-structured (or "N:M") sparse models. Unlike previous approaches that optimize sparsity patterns and weights separately, CAST enables seamless joint optimization during training, while progressively transforming the model into the desired sparsity format. Specifically, CAST introduces three key components: 1) AdamS, a sparsity-aware optimizer that leverages adaptive L1 decay to promote uniform sparsification across all parameters; 2) Weight Scaling, a module designed to mitigate the magnitude reduction caused by decay while preserving desired sparsity patterns; 3) Knowledge Distillation, which employs the dense model as a self-teacher to enhance training efficiency. We evaluate CAST under 2:4 sparsity patterns across multiple model families, ranging from 125M to 13B parameters. Our results demonstrate significant improvements over previous state-of-the-art methods in both perplexity and zero-shot accuracy with minimal training resources. Notably, on LLaMA2-7B, our 2:4 sparse model achieves a negligible perplexity increase of 0.09 and a 0.36% gain in zero-shot accuracy compared to the dense model using only 2% of the original pretraining tokens. Additionally, we establish an accurate and robust empirical scaling law to predict sparse model performance given adequate training resources. Finally, we demonstrate the practical applicability of our sparse models by evaluating them under quantization and fine-tuning scenarios.

Figures (8)

• Figure 1: Illustration of different sparsity patterns under 50% sparsity ratio. Unstructured sparsity constrains only the total number of non-zero elements; structured sparsity typically removes entire rows or columns; whereas 2:4 sparsity enforces that each 4-element group $g_r^{(i)}$ retains exactly 2 non-zero weights.
• Figure 2: (Left) In the SR-STE training baseline, the forward process is sparse, which can lead to suboptimal performance due to gradient approximation and discontinuities. (Right) In contrast, Continuous Adaptive Sparse Trainer maintains a dense forward process while using AdamS to gradually induce sparsity. CAST further incorporates knowledge distillation and a learnable weight scaling module to enhance training efficiency and improve model performance.
• Figure 3: (Left): Parameters with greater impact on model performance are preserved despite decay. (Right): Parameters with minimal contribution to the loss are gradually driven to zero, allowing AdamS to adaptively identify and retain important weights during training.
• Figure 4: (Left): Distribution of magnitude ratios of the two masked weights within 4-element groups under 2:4 sparsity in the first MLP layer of block 11 in the GPT-2 model with additive decay. (Right): Same distribution under proportional decay. Proportional decay uniformly drives masked weights to zero, whereas additive decay often leaves many masked weights insufficiently reduced, risking degradation during final pruning. Max pooling is applied to enhance readability.
• Figure 5: Validation perplexity on WikiText for GPT-2 under dense forward pass of SR-STE and CAST respectively. In SR-STE based training, the dense model's performance deteriorates rapidly, indicating that masked weights become unreliable upon reactivation. In contrast, CAST maintains the dense model in a lossless state throughout retraining, enabling effective and meaningful mask learning.

Theorems & Definitions (2)

• Lemma 1
• proof