Merging Feed-Forward Sublayers for Compressed Transformers

Thanks for your time and thoughtful review of our paper. We address the stated weaknesses and questions below:

Weakness 1: limited practical use

Compressing models without necessarily reducing speed is still a major practical contribution. While we agree that structured pruning methods do frequently result in some speed increase, many methods exist to train or compress models to achieve small sizes without focusing on run-time as well. For example, small language models with built-in weight sharing, some unstructured pruning methods and some quantization methods share this property. For example, LLM.int8 quantization is actually slower than its unquantized counterpart for GPT-Large.

Weakness 2: More baselines:

Thanks for your comment regarding comparisons and discussions. We agree that even if our method takes a very different approach to compression, it should be compared to alternatives. Since we propose a general compression method not specific to ViT, we choose a general pruning method for comparison. Many recent papers on structured pruning of Transformers (which results in compression unlike many unstructured pruning methods) have centered around choosing layers for dropping and then fine-tuning the resulting model [1,2,3]. We use a strong baseline of 1) picking the best layers to drop after evaluation (which generalizes and strengthens many of these layer-dropping papers) and 2) fine-tuning the same as our method, for all parameters. We choose the number of layers dropped to cover a similar range as ⅓ and ½ FFs removed. We present the results in the revised PDF (Section 4.4 details, Section 5.1 results) and link them here: https://anonymous.4open.science/r/temp-C34A/README.md . We outperform or match the baseline consistently judging by the curves across different parameter reduction ratios. Although we cannot achieve exact comparisons across specific parameter reduction ratios due to the block wise reduction nature of both methods, our method trends better in general, as seen in the figure.

Weakness 3: Writing quality:

Thanks for pointing out these 2 typos. We have corrected them in the updated PDF and rechecked thoroughly for others. However, we would like to point out that reviewer zq1y has stated the paper is well-written. We hope that this is just a minor issue at the typo level rather than something more systematic. We are happy to receive further feedback if it happens to be the latter.

Question 1: Other attention conclusions:

The first cited work re-trains only translation models from scratch with shared attention weights. Additionally, their similarity analysis is on weight distributions rather than the similarity between attention activations. Training the model from scratch with this inductive bias is very different than our comparison between attention layers, and seemingly easier to reuse attention states. The second cited work compares attention score matrices between layers whereas we compare attention output activations just after the linear projection following multi-headed attention. The output incorporates also the value vectors, as well as the output projection, leading to a very different output than the attention score matrix. In summary, direct similarity evaluations are applying on different objects in all 3 papers, and cited paper 1 is training models from scratch with this sharing paradigm.

Question 2: Why no P matrix in $b^{\text{out}}$ equation:

The $P_i^T$ matrix applies to the input dimension of the $W_i^{\text{out}}$ matrix, whereas the bias term is added to the output of the $W_i^{text{out}}$ projection. For FF sublayer $x_\text{out} = W_{\text{out}} \sigma(W_\text{in} + b_\text{in}) + b_{\text{out}}$, applying permutations yields $x_\text{out} = W_{\text{out}}P^T \sigma(P(W_\text{in} + b_\text{in})) + b_{\text{out}}$, showing this more clearly.

Question 3: vanilla baseline clarification:

Yes, this is correct. We include this baseline to observe the effect of the permutation alignment.

Question 4: Pre-tuning results:

These results are in section 5.2. Fine-tuning clearly contributes to the performance recovery, but we limit the amount of fine-tuning as described in Section 4. Through our results and analysis, we show that the permutation alignment and weight sharing provides a solid starting point for limited downstream fine-tuning.

Question 5: Sliding window:

We use the sliding window strategy to exhaust all sets of k adjacent feed-forward layers. We require adjacency for 1) combinatorial ease (ex, 36 GPT-2 layers choose 12 is > 1B) as well as some evidence of similarity that aligns with adjacency from prior work [1, 2] as well as our own (i.e. Figure 4).

[1] Pires et al. "One wide feedforward is all you need." arXiv preprint arXiv:2309.01826 (2023).

[2] Kornblith, Simon, et al. "Similarity of neural network representations revisited." International conference on machine learning (ICML), 2019.

Thanks for your time and thoughtful review of our paper, and thanks for noting our contributions of a novel method, diversity in models tested, and detailed ablations. We address the stated weaknesses and questions below:

Weakness: Weight sharing novelty: The cited Albert paper is a great example of effective weight sharing for model size reduction. We emphasize that the novelty of our approach is the post-training integration of weight sharing. Most weight sharing approaches in the literature occur at the initialization of the model. Our work provides a realistic way to introduce weight sharing into models that have already been pre-trained, going beyond models that have to do this from the start, like Albert.

Weakness: Comparing to pruning:

Thanks for your comment regarding comparisons. We agree that even if our method takes a very different approach to compression, it should be compared to alternatives. Since we propose a general compression method not specific to ViT, we choose a general pruning method for comparison. Many recent papers on structured pruning of Transformers (which results in compression unlike many unstructured pruning methods) have centered around choosing layers for dropping and then fine-tuning the resulting model [1,2]. We use a strong baseline of 1) picking the best layers to drop after evaluation (which generalizes and strengthens many of these layer-dropping papers) and 2) fine-tuning the same as our method, for all parameters. We choose the number of layers dropped to cover a similar range as ⅓ and ½ FFs removed. We present the results in the revised PDF (Section 4.4 details, Section 5.1 results) and link them here: https://anonymous.4open.science/r/temp-C34A/README.md . We outperform or match the baseline consistently judging by the curves across different parameter reduction ratios. Although we cannot achieve exact comparisons across specific parameter reduction ratios due to the block wise reduction nature of both methods, our method trends better in general, as seen in the figure.

[1] Men, et al. "Shortgpt: Layers in large language models are more redundant than you expect." arXiv preprint arXiv:2403.03853 (2024).

[2] Gromov, et al. "The unreasonable ineffectiveness of the deeper layers." arXiv preprint arXiv:2403.17887 (2024).

Weakness: Permute merge gains:

The gains attributable to permutation are a smaller addition onto solely weight sharing between feed-forward sublayers. However, they do provide consistent improvements over vanilla averaging. Another major contribution, as highlighted above, is the introduction of weight sharing as a post-training compression technique.

Weakness: MLPs only

We intentionally focus on MLP layers in this work. Our motivation is detailed in section 3.1, and we reiterate that these subcomponents are a majority of parameters for enc-only or dec-only models.

Question 1: GPT-2 permute v vanilla:

Permute FF merge does not perform worse than vanilla merge in GPT-2 (ref Figure 2, image B). However, we had a typo in from arranging results in the corresponding table in Appendix A that may have caused this confusion. We had also caught this typo on our end soon after submission and addressed it in the updated PDF. The figure in 5.1 was originally correct, and the corresponding table in the appendix is addressed to reflect this.

Question 2: all model layers:

The method was extended to all model layers in the main results, in 5.1. N-1 FFs removed is N FFs merged,which is all layers. However, we focus on smaller ranges of model layers as the method degrades at these very heavy compression ratios.

Thanks for your time and thoughtful review of our paper, and thanks for noting our contributions of ablations, visualizations, and contextualizing our method with prior work and other compression methods. We address the stated weaknesses and questions below:

Weakness 1: pruning comparisons:

Thanks for your suggestion regarding comparisons. We have included a new layer-pruning baseline that realizes compression ratios similar to ours, and discussed the choice of this baseline in new subsection 4.4 in our updated PDF. We do not consider methods like Wanda as baselines due to their unstructured sparsity patterns. Although it can achieve 50% sparsity in some LLMs, model weights are not actually compressed with this sparsity pattern without specific storage considerations (like COO, CSR, DOK sparse formats via sparse libraries), and these generally require ratios > 50% sparsity for actual compression on 2D weight matrices. This discussion of merging, unstructured pruning, and structured pruning is expanded in the PDF.

In our revision, we use a strong layer-pruning baseline of 1) picking the best layers to drop after evaluation which generalizes and strengthens many proposed layer-pruning methods [1,2] and 2) fine-tuning the same as our method, for all parameters. We choose the number of layers dropped to cover a similar range as ⅓ and ½ FFs removed. We present the results in the revised PDF (Section 4.4 details, Section 5.1 results) and link them here: https://anonymous.4open.science/r/temp-C34A/README.md. We outperform or match the baseline consistently judging by the curves across different parameter reduction ratios. Although we cannot achieve exact comparisons across specific parameter reduction ratios due to the block wise reduction nature of both methods, our method trends better in general, as seen in the figure.

[1] Men, et al. "Shortgpt: Layers in large language models are more redundant than you expect." arXiv preprint arXiv:2403.03853 (2024).

[2] Gromov, et al. "The unreasonable ineffectiveness of the deeper layers." arXiv preprint arXiv:2403.17887 (2024).

Weakness 2: Benchmarking GPT-2:

Thanks for your suggestion regarding LM type and evals. While we agree that LLaMA models would be ideal due to their popularity, they are still quite large for our experimentation, whereas GPT-2 Large is much smaller and still performant for its size. Regarding zero-shot eval, we wholeheartedly agree that for large language models, evals like MMLU are important to test knowledge retention alongside language modeling performance. However, on a smaller model like GPT-2, evaluating on MMLU is more appropriate after fine-tuning (like in the original MMLU paper). Since we are dealing with smaller scale models, and trying to measure just language modeling capability rather than also world knowledge, which is the goal of MMLU, we believe PPL is a sufficient metric in our reduced parameter case.

Question 2: Comparisons to pruning and quantization We have now added a strong pruning baseline, and respectfully disagree with comparing directly with quantization. Merging and quantization address orthogonal dimensions of compression (precision, redundancy), and we demonstrate their complementary and orthogonal performance, yielding benefits greater than either method alone.

Finally, thank you for your careful reading and catching our typo, it is addressed in our revision.

Thanks for your time and thoughtful review of our paper, and thanks for noting our contribution of our novel compression method and its experimental effectiveness. We address the stated weaknesses and questions below:

Weakness: Thanks for your comment regarding comparisons and discussions. We agree that even if our method takes a very different approach to compression, it should be compared to alternatives. Since we propose a general compression method not specific to ViT, we choose a general pruning method for comparison. Many recent papers on structured pruning of Transformers (which results in compression unlike many unstructured pruning methods) have centered around choosing layers for dropping and then fine-tuning the resulting model [1,2]. We use a strong baseline of 1) picking the best layers to drop after evaluation (which generalizes and strengthens many of these layer-dropping papers) and 2) fine-tuning the same as our method, for all parameters. We choose the number of layers dropped to cover a similar range as ⅓ and ½ FFs removed. We present the results in the revised PDF (Section 4.4 details, Section 5.1 results) and link them here for quick viewing: https://anonymous.4open.science/r/temp-C34A/README.md . We outperform or match the baseline consistently judging by the curves across different parameter reduction ratios. Although we cannot achieve exact comparisons across specific parameter reduction ratios due to the block wise reduction nature of both methods, our method trends better in general, as seen in the figure.

[1] Men, et al. "Shortgpt: Layers in large language models are more redundant than you expect." arXiv preprint arXiv:2403.03853 (2024).

[2] Gromov, et al. "The unreasonable ineffectiveness of the deeper layers." arXiv preprint arXiv:2403.17887 (2024).

Question 1: Training from scratch: Parameter sharing from scratch is a useful framework for pre-encoding parameter efficiency that has seen success in the past, as discussed in Section 2.1. However, our method showcases the effectiveness of weight sharing as a lightweight, post-training compression method that applies to pre-trained models, whereas the methods discussed in Section 2.1 train models from scratch. To retrain the models in this work with this new sharing structure would require extensive pretraining which is out of scope of this work.

Question 2: Adjacency: We use the sliding window strategy to exhaust all sets of k adjacent feed-forward layers. We choose adjacency for 1) combinatorial ease (ex, 36 GPT-2 layers choose 12 is > 1B) as well as 2) evidence of similarity aligning with adjacency from prior work [1, 2] as well as our own (i.e. Figure 4).

[1] Pires et al. "One wide feedforward is all you need." arXiv preprint arXiv:2309.01826 (2023).

[2] Kornblith, Simon, et al. "Similarity of neural network representations revisited." International conference on machine learning (ICML), 2019.