Pruning Mask

The primary function of a pruning mask is to definitively specify which parameters are active (1) and which are pruned (0). It acts as a blueprint for the model's sparse architecture. The mask has the exact same shape as the target weight tensor. During inference, this binary map is used by sparse kernels to skip computations involving zero-valued weights, directly translating to a reduction in FLOPs.

• Example: A mask value of 0 at position [i, j] means the multiplication involving weight W[i, j] is omitted.
• Formats: Often stored efficiently as a bitmask, where each bit corresponds to a weight or a block of weights (e.g., for N:M sparsity).

After initial pruning, a model typically loses accuracy. The pruning mask is fixed and used as an anchor during sparse fine-tuning or retraining. Only the weights indicated as active (1) in the mask are updated via gradient descent; pruned weights (0) remain frozen at zero. This prevents the optimizer from reviving pruned connections and ensures the sparsity pattern—and thus the computational savings—is preserved while accuracy is recovered.

• Process: Also known as masked retraining or sparse training.
• Importance: Decouples the structure of sparsity (defined by the mask) from the values of the remaining weights.

During sparse inference, the pruning mask is parsed by the sparse inference engine to guide efficient computation. It enables zero-skipping at the kernel level. Specialized SpMM (Sparse Matrix-Matrix Multiplication) and sparse convolution kernels use the mask's indices to perform gather-scatter operations, loading only the non-zero weights and their corresponding activations for computation.

• Hardware Support: On GPUs with Sparse Tensor Cores, the mask defines the valid 2:4 sparsity pattern for accelerated computation.
• Performance: The mask's encoding (e.g., CSR Format, COO Format) directly impacts kernel efficiency by influencing memory access patterns and load imbalance.

The pruning mask is a core component of a serialized, deployable sparse model. Alongside the array of non-zero weight values, the mask allows the full weight tensor to be reconstructed. Storing only non-zero values and their mask/indices leads to significant model compression. For example, a weight tensor with 90% weight sparsity requires storing only 10% of the original parameters plus the relatively small overhead of the index data.

• Combined Techniques: In sparse quantization, the mask is used in conjunction with low-precision (e.g., INT8) storage for the non-zero weights, yielding multiplicative size reduction.
• Formats: Integrated into runtime formats like TFLite sparse schema or framework-specific checkpoint formats.

For structured pruning techniques that remove entire filters, channels, or blocks, the pruning mask enforces hardware-friendly, regular sparsity patterns. This is crucial for achieving actual speedups on general-purpose hardware without requiring exotic sparse kernels. The mask ensures that entire rows/columns of weight matrices or blocks are zeroed out, leading to more efficient dense matrix operations on the remaining structure.

• N:M Sparsity: A specific constraint where the mask enforces that in every block of M weights, only N are non-zero. This pattern is natively supported by modern GPU Sparse Tensor Cores.
• vs. Unstructured Pruning: Unstructured pruning masks are irregular and offer higher theoretical compression but require specialized kernels for speedup.

The pruning mask is an inspectable artifact that allows engineers to analyze the characteristics of the compressed model. By examining layer-wise masks, one can measure the sparsity distribution, identify layers that are overly aggressive or conservative in pruning, and correlate this with accuracy metrics. This analysis is part of sparse model profiling and is essential for understanding the compression-accuracy tradeoff.

• Metrics: Used to calculate exact sparse FLOPs and model size for deployment planning.
• Debugging: Helps diagnose issues in sparse training pipelines by verifying the intended sparsity pattern is correctly applied and maintained.

A data structure for efficiently storing and operating on tensors where the majority of elements are zero. Instead of storing all values, it encodes only the non-zero values and their indices. Common formats include:

• CSR (Compressed Sparse Row): Optimized for row-wise operations.
• CSC (Compressed Sparse Column): Optimized for column-wise operations.
• COO (Coordinate Format): Stores explicit (row, column, value) tuples. A pruning mask is often the first step in generating a sparse tensor representation from a dense weight matrix.

The fundamental computational kernel for sparse inference, which multiplies a sparse matrix by a dense matrix. It is the workhorse operation for sparse linear layers and convolutions. Efficient SpMM kernels skip multiplications where the sparse operand is zero. Performance depends critically on:

• The sparse data layout (e.g., CSR vs. blocked formats).
• Memory access patterns to avoid bottlenecks.
• Load balancing across parallel threads, as non-zero elements are often irregularly distributed. The pruning mask directly defines the sparsity pattern that the SpMM kernel must exploit.

A specialized hardware unit within modern GPUs (e.g., NVIDIA's Ampere, Ada, and Hopper architectures) designed to accelerate sparse matrix operations. These cores leverage structured sparsity patterns, most notably the 2:4 sparsity constraint (2 non-zero values in every block of 4). This allows the hardware to effectively double theoretical compute throughput for eligible operations. The pruning mask must often be formatted to comply with such hardware-specific sparsity patterns to unlock this acceleration.

Fundamental parallel computing primitives essential for sparse computation on vector architectures like GPUs and NPUs.

• Gather: Reads data from a set of non-contiguous memory addresses (specified by an index array) into a contiguous vector register.
• Scatter: Writes data from a contiguous vector to a set of non-contiguous addresses. These operations are used to collect non-zero weights and their corresponding activations before computation and to write results back to correct locations. The indices in a sparse tensor representation drive these gather-scatter patterns.

The process of retraining a neural network after a pruning mask has been applied. The primary goal is to recover accuracy lost during pruning by updating the remaining non-zero weights. Crucially, the pruning mask is typically fixed during this phase; gradients are not computed for masked (zero) weights, and their values remain zero. This ensures the sparse architecture is preserved. Techniques like gradual magnitude pruning interleave pruning and fine-tuning steps to minimize final accuracy loss.

A software runtime or framework component specifically designed to load and execute sparse neural network models. It contains optimized kernels (like SpMM) for target hardware (CPU, GPU, NPU). Examples include sparse execution paths in TensorFlow Lite, PyTorch with torch.sparse, and proprietary SDKs for mobile NPUs. The engine is responsible for:

• Parsing the model file, which includes weight values and the associated pruning mask or sparse format metadata.
• Mapping the sparse computational graph to efficient kernels.
• Managing the memory layout of sparse tensors for optimal performance.