Summary
The paper introduces MaskLLM, a method for introducing semi-structured sparsity (N:M mask patterns) in LLMs. The authors show that existing model pruning methods such as SparseGPT result in significant loss in model quality at smaller scales (800M ~ 15B parameters) when using semi-structured methods. They formulate the problem of find a good N:M pattern as a mask selection problem from a candidate set of mask $S$. They formulate this further as a sampling problem, whereby they sample masks for all $\mathcal{W}^{1\times4}$ parameters in a layer to measure model quality - since exact mask computation is an intractable problem for large models. They propose using Gumbel-softmax sampling to figure out soft / differentiable masks that can learned through a training process, to optimize the sampling problem. The authors further find that learning of the masks by their proposed method can result in vanishing gradients throughout the network, and enable a sparse weight regularization to promote higher gradient norms through the network. They further find that using methods such as SparseGPT as mask priors enables more effective sampling (ie, learning) of the final masks. They follow these up with different experiments and ablations to highlight the strengths of their proposed method.
Strengths
1. The paper formulates the selection of N:M masks as an optimization problem, which enables scaling the method to large datasets and models.
2. The paper walks through the math required to understand their optimization formulation step-by-step, making it simple to understand.
3. The method focuses on common problems introduced by pruning (such as vanishing gradients of pruned weights etc.) and proposes ways to resolve them through regularization.
4. For each of the parts of the proposed method, the authors present ablations and results that validate their design choices.
5. The authors also showcase how their method can scale for transfer learning to downstream tasks (both via fine-tuning or via mask-transfer, as proposed via their prior initialized method)
6. The paper shows performance results from using their method (~1.3~1.4x faster) on A6000 GPUs.
Weaknesses
1. The best MaskLLM results presented rely on SparseGPT as a prior for mask intialization and then compare against SparseGPT for efficacy. This seems to be an unfair comparison, since you're using the method's best result and then improving on top of that. For e.g, without prior masks, the method seems to be much closer in performance to the SparseGPT method for the LLaMA-2 7B model (9.12 vs 10.42). It will be good to see how the method performs for say the 13B model without using any priors.
2. The authors show that SparseGPT has fundamental limits for improvement based on the number of samples (which is well documented and understood) - but the comparisons shown, for e.g., in Figure 4 are with different datasets? Did the authors test the collated datasets for the MaskLLM training with SparseGPT and what plateaus were observed there?
4. One inherent limitation to the method seems that there are many hyper-parameters to tune to find good masks ($\alpha$ for SparseGPT prior, $\mathcal{k}$ for logit scaling and $\tau$ for softmax temperature, $\lambda$ for the regularization coefficient). Do these hyper-parameters scale optimally for all model scales? Or is more fine-grained optimization needed as models scale up?
- Also for the $\lambda$ parameter, the authors show that using $1e^{-4}$ results in highest gradient norm [Table 12], but the hyper-parameter used in Table 7 is $1e^{-5}$, can the authors clarify this aspect?
Questions
1. For the LLaMA-2 models, the authors mention using the official dataset from the paper. However, the paper has no mention of any datasets used for training the models. Can the authors clarify this discrepancy?
2. In Table 11, when the authors mention the score from RIA, does that include the channel permutation method enabled in the paper?
3. Can the authors verify their magnitude pruning results for downstream tasks for Llama2-7B (Table 1). For such a large perplexity difference between magnitude pruning and SparseGPT / Wanda - the downstream results seem too be high?
Limitations
1. There are some inconsistencies in the links in the paper (for e.g., Appendix D mentioned in page 8 - for weight regularization maps to Section 4.1 in the paper - please fix this).
2. Throughout the paper, it is unclear what results map to what experiments. For example, for Figure 5, which models were used for the ablation? Same for the results in Tables 3, 5, and 13. Scanning the Appendix also did not help clarify this. It makes the results presented somewhat non-trivial to parse.