Dynamic Layer Tying for Parameter-Efficient Transformers

We thank the reviewer for the comprehensive and constructive feedback.

The related work section has been significantly expanded following the review, and it now includes all the requested lines of work. We now explicitly differentiate our method from pruning methods, Reuse Transformer, and the methods and directions of research pointed to by the reviewer. Lottery ticket hypothesis was mentioned in our submission manuscript as part of the motivation in the introduction and now also in the text describing ablation (ii), which shows that our method does not lead to the lottery ticket property.

(1) Several weight-sharing techniques should be also included in the comparison.

In Ablation (vii), we exactly replicate the Cyclic method from [1], which performs as well as the other two patterns explored in [1].

The Reuse Transformer [2] only shares 10% of the weights, so is not a good baseline.

Subformer [3] shares up to 50% of the weights (much less than what we share) and requires architectural changes that make it unsuitable for our experiments.

Xiaoet al. [4] share attention weights and not layers (and also share much less than we do).

We will add a modern pruning baseline to our experiments. However, the pruning typically reduces much less than what we share and is also a complementary technique to ours (the sharing and pruning can be applied side by side).

[1] Takase, Sho, and Shun Kiyono. "Lessons on parameter sharing across layers in transformers." arXiv preprint arXiv:2104.06022 (2021).

[2] Bhojanapalli, Srinadh, et al. "Leveraging redundancy in attention with reuse transformers." arXiv preprint arXiv:2110.06821 (2021).

[3] Reid, Machel, Edison Marrese-Taylor, and Yutaka Matsuo. "Subformer: Exploring weight sharing for parameter efficiency in generative transformers." arXiv preprint arXiv:2101.00234 (2021).

[4] Xiao, Tong, et al. "Sharing attention weights for fast transformer." arXiv preprint arXiv:1906.11024 (2019).

(2) I would suggest reporting the performance on downstream tasks for a complete comparison. Lower perplexity sometimes does not mean higher downstream performance.

Due to constraints encountered during the discussion phase, we were unable to run additional experiments. We agree that it would be a useful experiment to conduct.

(3) The perplexity of the results seems very high compared to previous results. As far as I know, the perplexity of WikiText-2 is usually around 20 and the perplexity of 1-billion words is around 30. I would suggest finding the proper training recipe or training longer to let the model converge, or the current results may not be very convincing. The current better perplexity in Tables 1 and 2 might only be due to the faster convergence of the model with less trainable parameters.

The perplexity results quoted are achieved by training the models on multiple datasets, not only on wiki/1b. For example, GPT-2 which was trained on multiple datasets without finetuning on wiki2 achieved a PPL of around 18 while the exact same model trained only on WIKI2 is around 50.

Question 1: In Tables 1 and 2, why the training time of the proposed approach is a bit longer compared to the conventional training when training GPT-2 but is much less when training BERT?

This is addressed in the results section: “With respect to training time, the results are mixed. While in Tab. 1 it is demonstrated that our method somewhat slows down the training time, Tab. 2 presents a reduction of almost 50% in runtime. We believe, but have not yet verified, that this is due to the difference in hardware between the two experiments. While GPT-2 experiments run on A100, the BERT experiments run on A6000/A5000.”

Thanks to the authors for the response.

• I checked the related work section. I found one typo "self-attention mechanism,m " (the m is extra) and also the parathesis for citation is missing.

• Also, could you show me some references that also get around 50 perplexity on sole training on Wiki-2?

• I am not satisfied with the explanation of the training time benchmark part, and the authors also don't know why this speed benchmark has mixed results. Now I am even more confused about whether the approach would make the training slower or faster because of this explanation, but it seems like we do know that the approach would make the training slower on either A100 or A6000.

• Currently, I am still not totally convinced by the experimental results as several fixed pattern weight-sharing in [1] are not compared and there are no results on downstream tasks.

[1] Takase, Sho, and Shun Kiyono. "Lessons on parameter sharing across layers in transformers." arXiv preprint arXiv:2104.06022 (2021).

We thank the reviewer for the supportive review and constructive feedback.

1. The method section, while technically sound, suffers from a lack of clarity as to the method being presented. For such a significant contribution, it is important to ensure that one can understand the fundamentals of the method without having to crunch through all of the equations presented and fill many of the gaps with their own imagination. Perhaps something like a functional diagram, a more intuitive algorithm, or just a page that lays out the ingredients one by one, as well as how they are optimized, followed by the algorithm page would serve clarity of communication better.

We have considered this valuable suggestion. We agree that the writing style is layden with equations. Since another reviewer has asked to reduce the redundancy of the method section, we do not wish to add content there. Would it be acceptable if the text in the introduction serves as an overview? We can replace the 6th paragraph of the intro with a textual description, eliminating most of the equations there. We are also looking to create a video that demonstrates the stages of training and its evolution.

2. It's not clear if the authors have tried additional rewards in addition to PPL after updates, or if it was the only one tried.

We have not tried additional rewards in addition to PPL. This would indeed be interesting. Experience with RL in other domains suggests that combining rewards can be highly beneficial.

3. Figure 1 is hard to read and understand. A more intuitive approach could leverage colour-coding to indicate a particular layers weights and perhaps a sequential figure showing a layer by layer colour coding might be more helpful/clear.

Following the review, in the revised version we have added color coding to the existing graphs. Our initial attempts with a sequential ordering of nodes were not yet attractive enough to share, we will put additional effort into this.

Question 1. Have you tried retraining the architecture that your method learned from scratch? As in, grabbing the connection maps, using that to initialize a new transformer and training from there. It would be interesting to how it compares with the model you trained jointly.

We believe that this is ablation (ii). This architecture when trained from scratch is less effective than a vanilla transformer with the same number of layers as the number of independent layers in it, as reported in ablation (i).

Question 2. How did you try to balance exploration with exploitation? Presumably past actions can have a major effect on future trajectories. Did you employ any random restarts, or perhaps running multiple models in parallel so you can get feedback from a population rather than an individual for your RL updates.

All these are good suggestions. However, our training dynamics are exactly as given in Alg. 1. No random restarts or parallel runs were employed. We think these may be helpful for finetuning (we discuss finetuning in the last paragraph of Sec. 5). Exploration vs. exploitation was balanced using the epsilon greedy policy with a high epsilon at the beginning that favors exploration over exploitation. The schedule for updating $\epsilon$ is given in line 27 of Alg. 1.

Thank you for the supportive review and constructive feedback.

1 The Methods section contains a considerable amount of redundant content, providing a step-by-step explanation of the algorithmic process.

Thank you for this suggestion. The line-by-line explanation has been moved to Appendix A, and we kept in the main paper only the information that was not reported before this explanation.

2 The experimental results only show the comparison of the proposed method and the conventional training method, lacking comparison with other baselines and the related methods.

We will add a modern pruning approach as an additional baseline to the final version. However, we note that such methods seldom prune more than 50% of the weights, which is much less than the reduction in the number of parameters we provide. Also, while pruning is a useful baseline for context, its contribution is orthogonal to ours and the two methods can be combined. We also note that ablation (vii) in our paper is similar to the cyclic sharing pattern of [1] [1] Takase, Sho, and Shun Kiyono. "Lessons on parameter sharing across layers in transformers." arXiv preprint arXiv:2104.06022 (2021).

3 The related work is not adequately summarized. There is only one publication after 2020 shown in the section of Related Work, and there is a lack of research focusing on improving the Transformer neural network architecture.

Thank you for your comment. Following the reviews we have added a considerable amount of related work to Sec. 2 of the manuscript. In the original submission, we have placed a lot of recent work in the first two paragraphs of the introduction, and we have now extended the related work section to better position and differentiate our work.

Question 1: The conclusion "Having this global alignment is crucial for smooth training despite large blocks of weights being copied during the process" seems to conflict with Table 3.(vi).

With regards to this ablation study, we indeed write: ``We conclude that freezing at initialization may not be crucial (more experiments are needed). However, it has a sizable advantage in the peak GPU memory consumption.’’

Following the review, we have revised the explanation in the discussion on how global alignment is obtained to discount the role of freezing. The revised hypothesis is that layer 0 cannot be too specific simply because there is no time for the other layers to co-adapt during the first k iterations.

This is with regard to the role of freezing in obtaining alignment. The alignment itself exists, as we demonstrate in Fig. 4. The necessity of this alignment arises from the fact that without it, every time we would replace a layer (which can never stop occurring due to the exploration constant $\epsilon$), the performance would drastically drop (this theoretical argument appears in the first paragraph of the discussion).

Question 2: 'While in Tab. 2 it is demonstrated that our method somewhat slows down the training time, Tab. 1 presents a reduction of almost 50% in runtime. We believe but have not yet verified, that this is due to the difference in hardware between the two experiments. While GPT-2 experiments run on A100, the BERT experiments run on A6000/A5000.' Is there a citation error present?

We are sorry, but we are not sure that we understand the comment regarding the citation error.

Question 3: The ultimate goal of reducing training parameters is to decrease memory? What is the actual consumption of memory?

The model itself is reduced in size as reported in Tables 1 and 2, i.e., by up to 90%. Following the review, we report in the revised manuscript the memory consumption during training. In the current implementation, which does not actively release memory and was not optimized for memory consumption, the reduction is around 2/3 (Tab. 3).

Thank you for the supportive review and constructive feedback. Following the review, we ran ViT experiments on CIFAR-10. The results are reported in the table below and reinforce the manuscript’s results: with an insignificant drop in accuracy, our model has only 22% of the original model’s parameters and only 7 out of 32 layers are independent.

We realize that CIFAR-10 is only a first computer vision experiment and therefore did not add it to the revised manuscript. Due to personal circumstances, we were not able to run as many experiments in the discussion period as we would normally do.