Dynamic Layer Tying for Parameter-Efficient Transformers
Authors: Tamir David Hay
Lior Wolf
Venue: ICLR 2024 poster (2024)
Keywords: NLP
Transformer
Reinforcement Learning
Primary area: representation learning for computer vision, audio, language, and other modalities

Abstract: In the pursuit of reducing the number of trainable parameters in deep transformer networks, we employ Reinforcement Learning to dynamically select layers during training and tie them together. Every few iterations, the RL agent is asked whether to train each layer $i$ independently or to copy the weights of a previous layer $j<i$. This facilitates weight sharing, reduces the number of trainable parameters, and also serves as an effective regularization technique. Experimental evaluations validate that our model modestly outperforms the baseline transformer model with regard to perplexity and drastically reduces the number of trainable parameters. In particular, the memory consumption during training is up to one order of magnitude less than the conventional training method.

then the weights of this layer are being optimized independently of other Published as a conference paper at ICLR 2024
DYNAMIC LAYER TYING FOR PARAMETER-EFRICHE TRANSFORMERS
Tamir David- Hay & Lior Wolf
Blavatnik School of Computer Science, Tel Aviv University {davidhay, wolf)0mail. tau. ac.il
ABSTRACT
In the pursuit of reducing the number trainable parameters in deep transformer networks, we employ Reinforcement Learning to dynamically select layers during training and tie them together. Every few iterations, the RL agent asked whether to train each layer i independently to copy the weights of a previous layer j < i. This facilitates weight sharing, reduces the number of trainable parameters, and also serves as effective regularization technique. Experimental evaluations val- idate that our model modestly outperforms the baseline transformer model with regard perplexity and drastically reduces the number of trainable parameters. In particular, the memory consumption during training is up to one order of mag- nitude less than the conventional training method.
INTRODUCTION
The recent work on large language models is based mostly the transformer architecture
of Vaswani et al. (2017). Such models have become increasingly larger and are trained for 100s
of thousands of GPU hours using high-end GPUs (Brown al. 2020 Chowdhery al., 2022 Rae
ce al. 2021 Touvron et al. 2023).
However, is clear that the Transformer architecture (like other deep architectures) is overparame- terized. For example, pruming can be used to reduce the number of FLOPs of transformers during
inference time at least by half, with little effect on accuracy (Kurtic et al., [022] K Kwon al.)
2022
attention heads can be removed post-training with little effect performance MMichel et al.
2019
Voita et al. 2019). The lottery ticket hypothesis holds for transformers (Frankle & Carbin
2018
Chen et al. 2020a b; Prasanna et al. 2020 Movva & Zhao 2020). and, perhaps most relevant to work, layers can be dropped altogether during inference Fan et al. (2019); Sajjad et al. (2020). and
attention scores can be reused Bhojanapalli et al. (2021).
Motivated by the potential to reuse transformer layers, we conducted a preliminary experiment in which we started with transformer of L layers and trained only &/2 layers by sharing the weights between layers i and layer + 2/2 for < 4. The transformer trained this way achieved the same, or somewhat better performance, as the conventional I, layer transformer.
This encouraging preliminary finding raises few questions. First, is there something special about this pattern of repetition? Second, is a factor of two the best we can get? Taken to the extreme, it would be desirable that every layer in the architecture either replicates one of the previous layers or, if needed for the sake of accuracy, have new set of weights.
Our method opts to find such a general pattern. Trying to train only once, we view the repetition pattern as dynamic action that some driver network Q learns from reinforcement during the training of the primary network T. Every few epochs, a new action vector a is obtained based on the Q-function estimation given by Q.
The element 0; E [0, 1, 1..... for i 1 indicates from which layer to copy the weights to layer of T. If a; layers. If a; < then the weights of layer a, are used as the weights of layer This is transductive:
if layer a, j and aj k. then both layers i and share the same weights of layer k.

---

Published as a conference paper at ICLR 2024
After a few training iterations of the primary network T, the reward for the driver network Q is computed by considering the loss obtained on a few training batches. Q is then updated, and a new action vector is recovered. Depending the dynamics of the driver network, the changes in the replication pattern can be rather rapid. Yet, as we show, the training process is stable
Our results indicate that training this way leads layers while maintaining the same level of accuracy, or even slightly better, as the full L layer trans- former. This is achieved with a relatively small Q network, which is only applied during training.
replication of at least 75% of the transformer
Our contributions are: (i) Presenting a novel method for dramatically reducing the number of param- eters in transformer architecture. (ii) Establishing the potential of Reinforcement Learning (RI.) to serve as pivotal mechanism for dynamically optimizing the architectural configurations of trans- formers during training. The impact of RL in this context is considerably more profound than its conventional applications, such as adaptive learning rate tuning Xu et al. (2019). (iii) Demonstrating the use of RL in Neural Architecture Search (NAS) a single training pass, unlike all previous work we are aware of, which follow [aker et al. y [ao et al. (2019); Zoph & Le (2016) and collect multiple training sessions. (iv) Showing that transformers can be trained effectively, despite rapid changes in architecture during the training process.
2 RELATED WORK
Our method changes the architecture of the Transformer network and is, therefore, a Neural Ar- chitecture Search (NAS) method. The promise of the field is to discover architectures that would surpass goman-ddesinndd ones in performance. While most recent contributions rely on techniques such Differentiable Architecture Search (Liu al., 2018a), some of the earlier approaches relied on RL. Zoph& & Le (0177). employ a recurrent neural network (RNN) to generate architectural de- scriptions of neural networks and train it with RI. Baker et al. 2017) employ Q-learning to search for optimal CNN architectures. Cai et al. (2018) uses a controller, trained with the policy gradient method, to search for architectures in more computationally efficient manner. As mentioned, RL NAS methods suggest a fixed architecture and train it from scratch, using the validation score as a reward. The trained network is not changed dynamically during training as we do.
The use of RL for dynamically controlling the training of a deep neural network has focused on learning rate optimization. Controlling the learning rate is often done with fixed schedule, such as step decay
cosine decay, which determines the step size for each iteration of the optimization process (Ruder 2016). Xu et al. (019)) employ proximal policy optimization (PPO) (Schulman et al. 2017) trained across multiple sessions (not single session as in our method). (Subramanian ctt al. 2023) also employ PPO, and use a state vector that includes the training loss of the last epoch, the epoch index, and the number of remaining epochs.
Considerable effort has been dedicated to making transformers more efficient by reducing the quadratic complexity of the self-attention mechanism, e.g., (Child et al., 2019 Ma et al. 2021 With respect to parameter efficiency, network pruning methods (Molchanov et al. 2016 Hassibi et al. 1993; Frankle & Carbin 2018 Liu et al. 2018b) including the transformer pruning methods mentioned above (Kurtic et al. 2022 Kwon et al., (2022) reduce the size of the network by remov- ing or shrinking matrices from the network. Such methods often require further re- re-training, while our method is applied during training. maintaining the training time per epoch and reducing the peak memory consumption. The recent Wanda method (Sun et al., (023) performs straightforward magnitude- based pruning (Han et al., 2015; Gale et al.) 2019 Zhu & Gupta 2018 Liu et al. 2018b) on the trained transformer. Despite its simplicity, it is shown to outperform other pruning alterna- tives. In comparison to our method, the sparsity demonstrated is up to 50% of the weights, while our method is shown to reduce 75% to 87% of the parameters. Our approach, which focuses on reuse, and pruning, which attempts to "reduce", are not mutually exclusive and can be combined.
Other methods that reuse computations or parameters within transformers include the Reuse Trans- former (Bhojanapalli et al.) (021) which, unlike our method, uses specific and fixed pattern of reusing elements and only reuses attention heads. Overall less than 10% of the parameters are shared. Similarly to the Reuse Transformer, the Subformer (Reid et al. 2021) shares the parameters of the middle layers, however, much more extensively, reaching up to 50% reduction in the num- ber of parameters. This requires the addition of auxiliary network elements, which we do not do.
2

---

Published as a conference paper at ICLR 2024
Algorithm Q-learning driven dynamic layer tying
Require: L the number of layers, K the number of training steps of T, k the number of training
steps between the update and evaluation of Q, y the discount factor, and initial exploration probability
1: Initialize the primary model T and the Q-network Q
2: Freeze layers L in T. such that only layer trains at initialization.
3: Initialize 8 = a 0
An all zero vector
4: for step = 0 to
do 5: Sample a mini-batch B from the dataset
6: Perform a training step with T on if' mod(step.k) = then
7:
Every k steps
8:
Obtain an action vector a = (s) Compute s' based for i 0 to I, do if s + then
9:
Eq.
10:
11:
12:
s = then
Untie layer i Tf Copy its weights and update it independently of layer s;
14:
else
Replicate all weights of layer s of T layer of T Tie the weights of layer i to layer s
16:
end if end if end for
18:
19:
20:
Sample mini-batch B from the data-set
21:
steep Compute negative PPL score based on T on T Tpredicated Q(ss( a.)
22:
Eq.
23:
r = Tstep + + max, Q(s')a
24:
L p(roditid r)
25:
update Q using L
s'
27:
c (e 0.95, 0.1}
end
29:
end for
Takase Kiyono (2021) explore three different fixed patterns sharing parameters, reusing 50% to 66% of the layers. The differences in performance between the patterns are small, and our last ablation (ablation vii) is similar to the Cycle pattern. Xiao et al. (2019) share attention weights (and not the parameters for computing these), based on the attention similarity. The number of reduced parameters not reported but the average speedup 1.3 (23% reduction).
Parameter Efficient Fine-Tuning (PEFT) often target specific layers or modules, e.g., only the top layers (Gheini et al. 202T), only the bias parameters (Zaken et al. 202T), or selecting based on scores (Sung et al.) 2021] Vucetic et al., 2022). Additive PEFT methods introduce additional trainable parameters that can be added to the attention and feed-forward layers of transformers (Houlsby ct al. 2019). LoRA (Hu et al. (022) adds low-rank matrices to the weight matrices. PEFT methods substantially reduce the number of trainable parameters, but are applicable for finetuing (after the full model has been trained), while method is for training from scratch. See Sec. [ for future work on finetuning.
METHOD
We aim to train a transformer T with L layers from scratch. All elements transformer layer, including the key, query, and value projections, and the linear layers considered single set of training parameters. The set of parameters for layer can be either independent from all layers j < i, tied to the set of parameters of some layer j i.
The state vector s E N indicates, at each location i = 0,1,...,, L 1, the layer with the lowest
index that has the same tied weights. Therefore, Vi [0.... L 1] < 8; i. If

---

Published as a conference paper at ICLR 2024
it indicates that layer i does not have its parameters tied with any of the previous layers. By this definition, it always holds that So = O.
The action space is similar, except that the action vector N can point to any previous layer that has its weights tied with layer i, not necessarily the one with the lowest index j S
To obtain from a, one employ the following recursion
a.
8:
Sa:
Otherwise
The Q-function of a Markov Decision Process represents the expected cumulative future reward for taking a particular action in a particular state 8, while following certain policy π (Sutton & Barto 2018). Similarly to previous work that employs deep Q-learning(Mnih et al., 2013). we employ greedy policy obtained interpolating between a random policy and one obtained by maximizing, at a given state, the Q-function over the available actions.
{ arg max, Q(s, a) uniformly sampled at probability
at probability 1-
π (s))
where Q is the network we learn in order to approximate the Q-function. Its implementation takes as input and returns a vector of Q-values for each index i, indicating the Q-value obtained for each action j 0, i.
((s, a) = = [(s))iaa
(3)
where indexing occurs first for the vector of Q-values per each index and then for an element in this vector. Therefore, the input and output domains of the approximated Q-function are Q R 4-1
(L+2)(1-)) This reflects the fact that so is fixed and that for every layer = 1,2,.... L - the network Q needs to assign values to i + different actions. The optimal action-value function Q* obeys an important identity known as the Bellman equation
Qs(s.a) [r + y max Q" (s',
(4)
a²
We run the policy π based on Q to obtain a new action a after every & training steps of the primary network T. k is relatively small and such actions are taken frequently. At initialization, only layer i = trained; all other layers are fixed at their initial values. Then, after training steps, and every l training steps afterwards, we perform the following set of actions: (i) obtain new action a = rTs). (ii) extract the new state s' based on a, as in Eq. [ (iii) replicate the weights of each layer to be the same as 8; and tie these weights, (iv) compute reward for 'I based on the negative perplexity score as computed on a random training batch, (v) update Q based on the expected reward vs. the computed one, using the Bellman equation, (vi) reduce the exploration factor by fixed factor of 0.95, but always keeping it above a constant of 0.1, and, finally, (vi) run k more training steps for and repeat.
The method is depicted in Alg. I and a line- line-by-line description is provided in Appendix A A few implementation details are worth noting. First, in line 2, the replication pattern of the first k steps (where k << K) is determined to be such that layer 0 trains and the other layers are kept fixed at their initialization values. Then, every steps we obtain new action a, using the e e-greedy policy in Eq. [2] see lines 7-8. Second, we note that a layer that changes state can, based on the condition in line 12, either (i) shift from being untied or tied to one layer to being tied to a new layer, or (ii) shift to being trained independently. In the first type of shift, a new set of weights would be copied, which may change the transformer much more quickly than through gradient steps. In the second type of shift, the weights are not changed immediately. However, they begin to drift between layers that were previously tied. Third, the exact schedule for modifying the value of e is given in line 27.
4 EXPERIMENTS
In our experiments, two architectures were used: (i) GPT-2 with 48 decoder blocks, each with 16 attention heads. The hidden dimension for each block was set 1600, and (ii) BERT, which consists
4

---

Published as a conference paper at ICLR 2024
Table 1: Metric scores for the GPT-2 architecture
Training set
Metric
Method
Wiki-2
Wiki- 103
Lambada 1-billion
Conventional training
53.57
22.32
94.96
88.35
Perplexity
Our method
49.37
22.35
93.84
72.35
Conventional training
1.6B
1.6B
1.6B
1.6B
Number of trainable parameters
Our method mean over training
171M
151M
166M
218M
Our method at end of training
264M
142M
326M
203M
Conventional training
48
48
48
48
Number of independent layers
Our method mean over training
4.395
2.309
3.547
4.486
Our method at end training
7
10
Conventional training Our method
148.5
3609.5
26376.5
15440
Training time per epoch (seconds)
165
4010.5
29307.2
17155.5
Table 2: Metric scores for the BERT architecture
Training set
Metric
Method
Wiki-2
Wiki-103 Lambada 1-billion
Conventional training
70.15
154.2
202.70
>100
Perplexity
Our method
69.27
132.6
156.30
215.50
Conventional training
376M
376M
376M
376M
Number of trainable parameters
Our method mean over training
52M
52M
52M
57M
Our method at end of training
46M
46M
67M
60M
Conventional training
12
12
12
12
Number of independent layers
Our method mean over training
2.36
2.83
1.88
2.45
Our method at end of training
3
Conventional training Our method
51.2
1244.6
5324.1
9095.3
Training time per epoch (seconds)
26.5
644.5
2757.1
4704.4
of 12 decoder blocks with hidden size of 768 and attention heads at each layer. In all of our experiments, Q is an MLP with one hidden layer with 128 units and the ReLU activation function.
We ran all experiments for K 300 epochs, batch size of 16, and l: 15 with a separate validation set used to select the best model. The hyper-parameters used were: the transformer learning rate is set to 0.0001 and Q's learning rate was set to 0.001, = 0.99, the initial exploration probability is set to value is set to 0.1. Our experiments ran on 2-4 A100 GPUs for the GPT-2 based architecture and 1-4 A6000/A5000 GPUs for the BERT architecture.
1.0 (explore), and as depicted in Alg. T the e-decay factor: 0.95, and the minimal
Datasets this study, we employ four widely used datasets to evaluate the performance of our method for language modeling tasks. All datasets tokens using GPT-2's 2ss tokenizer, which has vocabulary of 50, 257 tokens. WikiText-22 (Wiki2) is large language modeling corpus that consists of 2 million tokens. It derived from snapshot of verified Good and Featured articles on Wikipedia. The dataset is widely used for training language models and serves as standard benchmark for evaluating various NLP algorithms. WikiText-103 (Wiki103) is an extension of the WikiText-2 dataset, containing more than 100 million tokens. It is also sourced from Wikipedia articles and considered to one of the most comprehensive datasets for training large-scale language models. LAMBADA is designed to test the capabilities of language models in predicting the final word of sentence, given all the preceding words in that sentence. The dataset contains approximately 10,000 examples, each sequence of sentences extracted from books. The task is challenging as often requires understanding the broader context provided by the preceding sentences. The 1 Billion Words dataset is a corpus of text containing approximately billion tokens. sourced from news articles. provides a diverse range of vocabulary and sentence structures, making it ideal for training robust language models.
pre-processed by converting the text into
5

---

Published as a conference paper at ICLR 2024
(b)
Figure I: The replication map for the GPT-2 architecture post-training for (a) Wiki-2, (b) Wiki-103. The lowest-index layer in each group of layers that share weights is connected to itself.
Results In Table ! we present a comprehensive evaluation of our proposed method against con- ventional training the GPT-2 architecture across multiple datasets: Wiki-2, Wiki-103, Lambada, and 1-billion. Our method consistently outperforms the baseline in terms of perplexity, with the most significant gains observed in the 1-billion words dataset, where reduce the perplexity from 88.35 to 72.35. Additionally, our method exhibits a significant reduction in the number of trainable parameters, with a mean over training other datasets, compared to the baseline's 1.6B. Although the conventional method outperformed our method on Wiki-103, the gap is marginal.
low as 151M for Wiki-103, and not much higher on the
Table 2 showcases the results for the BERT architecture, presenting similar trends. Our method outperforms the conventional training across all datasets. Notably, in the 1-billion dataset, the per- plexity is reduced drastically, from over 1000 in conventional training to 215.50 in our method. The number of trainable parameters also sees substantial decrease, with mean during training of 52M-57M, compared to the conventional 376M. In both architectures, number of independent layers (or, equivalently, the number of groups of identical layers) is rather low during training and is somewhat higher in the final model. Especially in BERT, that even for large datasets the number of independent layers small. In our ablation study below check whether one can simply train much less layers.
observe that the mean
observe
With respect to training time, the results are mixed. While in Tab. it is demonstrated that our method somewhat slows down the training time, Tab. [ presents reduction of almost 50% in time. We believe, but have not yet verified, that this is due to the difference in hardware between the two experiments (GPT-2 runs on A100, the BERT runs on A6000/A5000).
The status at the end of the training is shown in Fig, I A line is drawn between every layer index and the layer it replicates 8;. layer for which the state vector satisfies 8; is connected to itself. As shown, there are seven such layers for Wiki-2 and six for Wiki-103, matching the statistics report in Tab, 2 The dominance of layer zero clear, see Sec. 5 for discussion of this property and its implications.
Training dynamics The training process takes place under the guidance of a policy that trained from scratch. This policy can change the layer topology drastically and it is, therefore, interesting to explore the training dynamics. First, it is not clear whether any changes are made at all to the topol- ogy throughout the training process, It could be the case that after certain period of exploration, the policy is to keep the state fixed from one step to the next. As Fig. 2 demonstrates, this is not the case. We distinguish two types of state-change events, detailed in Sec. In the first, which we called "tied events", a layer replicates layer it did not replicate previously. In the second type, termed "untied events", layer obtains new state of
i and is trained independently of

---

Published as a conference paper at ICLR 2024
ges Changes
Figure 2: The number of change state events per type for training GPT-2 on Wiki-2
Number Times Each Layer was Trainable (Mapped to Itself) Over Time
3000
2500
2000
Count
500
0 6 9,00000
Layer Number
Figure 3: The number of steps in which each decoder block was trainable
Table 3: Memory consumption during training (GPT-2; batch size of 16: sequence length of 256).
Statistics
Conventional training
Our method
Peak memory
12,566. 66 MB 10,223.08 MB
4,514.31 MB
Average memory consumption
3, MB
previous layers, which had replicated another layer j < i. Evidently, both types of events continue to occur throughout the training process and their frequency does not diminish.
The memory consumption during training is result of the training dynamics. Tab. 3 depicts the peak and average memory consumption during the training of GPT-2. Our memory consumption is lower by 65% in peak consumption and 68% on average. This difference is obtained without any attempt to optimize memory usage during training or to release unused memory. and does not reflect in full the drop in the size of the model.
One may wonder if all layers have the same chance of being untied. We note that since the explo- ration factor is at least 0.1 throughout training, with the exception of layer O, all layers are expected to be tied to other layers at one point or another. As can be seen Fig. 3 this is indeed the case.
It can also be observed that the lower layers are more likely to have an untied status of 8; = i (other layers with index j > i may still have 8, = i and train together in a tied way). This makes sense due to the increasing number of replication options that higher layers have. However, we note from Fig. bb that the layers with 8; == can be relatively evenly distributed at the end of training.
Ablation study Ablation experiments were conducted on the Wiki-2 dataset with the GPT-2 ar- chitecture. Since much of the ablations focus on validating that the success of the method does not arise from avoiding overfitting by reducing the network capacity, Shakespeare dataset. This dataset has parts of Shakespeare's 's plays, sonnets, and other writings. is small, with 250K tokens and the ablation uses a 12-layer GPT-2 like model.
also run ablations on the small
7

---

Published as a conference paper at ICLR 2024
Table 4: Perplexity scores for the ablation study for the Wiki-2 and the Shakespeare datasets.
Architecture
W
Shakespeare
(i) Vanilla transformer L #independent layers in ours
54.64
167.3
(ii Training all epochs with the final architecture
59.35
185.3
(iii) Applying the recorded dynamic permuted indices
65.21
172.2
(iv) Applying the recorded dynamics on the indices without Q
50.05
161.3
(v) Fully dynamic without weight tying
235.8
202.9
(vi) All layers are trainable at initialization
50.18
159.3
(vii) "Cycle": Connecting layer to layer 1 +
51.93
176.6
(viii) "Cycle Rev": Connecting layer to layer L
52.83
180.0
(ix) "Sequeenc". Connecting pairs of consecutive layers
54.01
173.9
Our full method
49.37
161.1
Since the model obtained with our method has about sixth of the number of parameters in the original model, we need to explore whether the full model capacity is required at all. To validate this, we designed few ablations: (i) transformer in which the number of layers I. is the number of independent layers obtained by our method, and (ii) training from scratch a static transformer architecture that has the same weight-tying structure as our method's final architecture.
As can be seen in Tab. 4 both these transformers are far behind our full method's results and also behind the conventional training results. The second ablation implies that our method is not suitable for finding "lottery tickets", i........... pruned architectures for training from scratch (Frankle & Carbin,
2018 Chen al. 2020a b. Prasanna et al. 2020 Movva & Zhao 2020).
Another related ablation (iii) checks whether the dynamic status changes can be made arbitrarily, by recording the state vector s during the course of training, and applying a permuted version of it (s) when changing the status of a layer to copy another layer or to be tied, where is fixed permutation operator that is applied element-wise.
The results of this ablation demonstrate that the layer identity is important and that a significant degradation occurs in the model's performance when the same dynamics are applied to a different set of layers. As sanity check, we also (iv) run the recorded set of states on another training session (without performing Q-learning). As can be seen, this obtains results that are similar but slightly worse than those of the full unablated method.
The necessity of weight tying is demonstrated by ablation (v), in which weight replication occurs as in the full method, but weight tying does not take place. This leads to very unstable training and very high perplexity score.
We also explore (vi) the effect of freezing all layers except for layer 0 initialization by freeing all layers to train (removing line 2 of Alg. [ This somewhat outperforms the full method on the shakespeare dataset but is less successful Wiki-2. We conclude that freezing at initialization may not be crucial (more experiments are needed). However, has a sizable advantage in the peak GPU memory consumption.
We also provide results for (vii) using half the layers and tying every layer i = 2, L/2 to layer L/2-+1 i. This cuts the number of trained layers by much smaller fraction than our own method and is given as a reference since it was outlined as motivation in Sec. [ As mentioned, this improves perplexity over the conventional training. but not nearly as much as our full method.
As mentioned in Sec. g ablation (vii) is the Cyclic pattern of (Takase & Kiyono 2021). The two other patterns there are provided for completeness as ablations (vili) and (ix). As can be seen, these patterns, which reuse only 50% of the layers, are not as effective as our method.
5 DISCUSSION AND LIMITATION
Replacing the weights of an entire layer with those of another is a drastic change to the network. Yet, as shown in Fig. 2 (blue graph), such changes occur throughout training. This ability to perform this change without causing temporal setback to the training process is not trivial. since even functionally equivalent layers can be expressed in multiple ways, by permuting the attention heads or the outputs of the feed-forward network. However, permutation the ferw-fowward network would
8

---

Published as a conference paper at ICLR 2024
LOVIE SET Mastc
LINE SANESPPIIE wty MARTC
Figure 4: (a) Pearson corre- lations between the weights of the feed-forward networks of the untied layers (Wiki- 2; GPT-2 architecture). The colorbar range is [0.93,1] (b) As a reference, the cor- relations between the same layers in the conventionally trained GPT-2 model. The value range is [O,1].
(b)
drastically modify the token embedding the next layer observes, and would cause the network's performance to degrade unless the other layers co-adapt.
We attribute the fact that such setbacks occur to the way the training process initializes. Layer O trains in a way that cannot be too specific, due to the randomly initialized filters downstream, which require time to co-adapt. Then, layer O is replicated and multiple copies of it trained simultancously. Other layers are also copied and their copies begin to train. However, given that layer zero is valid replication source for all layers, and given that the exploration constant is initialized at high value, layer zero is dominant. This domination, as can be seen in Fig. L is maintained until the end of training.
We posit that all layers are exposed directly or through a replication chain to the information of layer 0, and that spreads a specific order of attention heads and embeddings that are maintained across layers. Having this global alignment is crucial for smooth training despite large blocks of weights being copied during the process. Support for this hypothesis can be seen Fig. [a), which depicts the Pearson correlations between the weights of the feed-forward networks of the independent transformer layers trained with our method. The minimal correlation is 0.93, For reference, the correlation between the same layers in the conventional training (some of the 48 untied layers) is shown in panel (b). The inter-layer correlations are close to zero, as expected by the arbitrary permutation argument.
Our research is focused on training transformer models from the ground up. contrasting with the extensive body of work that primarily concentrates the fine-tuning of pre-trained transformers. (Devlin et al. 2018 Liu et al.) (2019; Dodge et al.) 2020; Raffel et al.) 2020) Brown et al. 2020 He al.] 2021). is unclear whether method that starts with one trainable layer and then gradually explores options to untie some layers can be applied in such case, especially since, as shown in Sec. the number of independent layers remains small throughout training. An alternative that makes sense, but which is left for future work, is to apply the dynamic weight tying to the low- rank updates (LoRA) of Hu et al. (202T). One can also try to apply RI. methods that employ backtracking (Dary et al. 2022), use alternative search strategies, such as CAB (Zhang 1998) or MCTS [Chaslot et al. 2008), changing one state index time.
The evaluation of our work is limited to transformers in the language domain. However, transformers are ubiquitous. preliminary computer vision experiment reinforcing our conclusions can found in Appendix B Finally, transformers are often finetuned downstream tasks. Preliminatry results on the GLUE set of benchmarks Wang et al. (2018) are presented in Appendix demonstrating that the tied models can be effectively trained for downstream tasks.
CONCLUSIONS
We present method that is, as far as we can ascertain, the most dynamic form of Neural Architec- ture Search presented. During the training process itself, a deep Q-learning network drives a layer replication process, which ends up with over 90% of the parameters being in layers that completely replicate an earlier layer. This order of magnitude reduction in the number of parameters is achieved without sacrificing the perplexity score and, in cases, also leads to improvement in this metric. These surprising findings are further explored by visualizing the dynamics of the training process and the crucial components of the method demonstrated in ablation study.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by a grant from the Tel Aviv University Center for Al and Data Science (TAD).
REFERENCES
Bowen Baker, Otkrist Gupta, Nikhil Naik, and Ramesh Raskar. Designing neural network architec- tures using reinforcement learning. ICLR, 2017.
Srinadh Bhojanapalli, Ayan Chakrabarti, Andreas Veit, Michal Lukasik, Himanshu Jain, Frederick Liu, Yin-Wer Wen Chang, and Sanjiv Kumar. Leveraging redundancy in attention with reuse trans- formers. arXiv preprint arXiv:2110.0682), 2021.
Tom Brown, Benjamin Mann, Nick Ryder. Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Han Cai, Ligeng Zhu, and Song Han. Efficient neural architecture search via parameter sharing. ICLR, 2018.
Guillaume M Jb Chaslot, Mark HM Winands, H Jaap van den Herik, Jos WHM Uiterwijk, and Bruno Bouzy. Progressive strategies for monte-carlo tree search. New Mathematics and Natural Computation, 4(03))343-357) 2008
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis for pre-trained BERT networks. arXiv preprint arXiv:2007. 12223, 2020a.
Xiaohan Chen, Yu Cheng, Shuohang Wang, Zhe Gan, Zhangyang Wang, and Jingjing Liu. Earlybert: Efficient bert training via early-bird lottery tickets. arxiv preprint trXiv:2101.0000 2020b.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1000//044//// 10509, 2019.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023111 2022.
Franck Dary, Maxime Petit, and Alexis Nasr. Dependency parsing with backtracking using deep reinforcement learning. Transactions of the Association for Computational Linguistics, 10:888- 903, 2022.
Jacob Devlin, Ming- Weig Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810004800. 2018.
Jesse Dodge, Gabriel Ilharco, Roy Schwartz, Ali Farhadi, Hannaneh Hajishirzi, and Noah Smith. Fine-tuning pretrained language models: Weight initializations, data orders, and early stopping. arXiv preprint wwxii:2002206300 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL hetps: w/openzeview.netee..
cNTTy
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1909. 1556, 2019.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: 8803.0363, 2018.
10

---

Published as a conference paper at ICLR 2024
Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.0957: 2019.
Yang Gao, Hong Yang. Peng Zhang, Chuan Zhou, and Yue Hu. Graphnas: Graph neural architecture search with reinforcement learning. arXiv preprint arXiv: 1904,0981, 2019.
Mozhdeh Gheini. Xiang Ren, and Jonathan May. Cross-atiention is all you need: Adapting pre- trained transformers for machine translation. arXiv preprint arXiv:2104.08771, 2021.
Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural networks with pruming, trained quantization and huffman coding. arXiv preprint arXiv: 2015.
Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In JEEE international conference on neural networks, pp. 293-299. IEEE, 1993,
Pengcheng He, Xiaodong I Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhanced bert with disentangled attention. In International Conference on Learning Representations, 2021.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly, Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning, pp. 2790-2799. PMLR, 2019.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation large language models. International Conference on Learming Representations, 2021.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen, LoRA: Low-rank adaptation large language models. In International Con- ference on Learning Representations, 2022. URL
net forum?
id=nzeVKeeFYFg
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Eldar Kurtic, Daniel Campos, Tuan Nguyen, Elias Frantar, Mark Kurtz, Benjamin Fineran, Michael Goin, and Dan Alistarh. The optimal bert surgeon: Scalable and accurate second-orde -oddrr pruning for large language models. arXiv preprint 2rXiv:2203.07259 2022.
Woosuk Kwon, Sehoon Kim, Michael W Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gho- lami. A fast post-training pruning framework for transformers. Advances in Neural Information Processing Systems, 35:24101-241 16, 2022.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. In International Conference on Learning Representations, 2018a.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arxiv: 1692, 2019.
Zhuang Liu, Mingjee Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. Rethinking the value of network pruning. arXiv preprint o 1810.881000000 2018b.
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettle- moyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Sys- fems, 34:2441-2453, 2021.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? arXiv preprint arxiv: 10650, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 132.5533/ 2013.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. arXiv preprint arXiv: 16/ 1, 1,06400, 2016.
11

---

Published as a conference paper at ICLR 2024
Rajiv Movva and Jason Zhao. Dissecting lottery ticket transformers: Structural and behavioral study of sparse neural machine translation. In Proceedings of the Third BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP, pp. 193-203, 2020.
Sai Prasanna, Anna Rogers, and Anna Rumshisky. When BERT plays the lottery, all tickets winning. arXiv preprint arXiv:2005.0056/ 2020.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring. Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:21 1444, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, al. Exploring the limits of transfer learning with unified text-to-text transformer. J. Mach. Learn. Res., (1(14)):1-61 2020.
Machel Reid, Edison Marrese- Marrese-Taylor, and Yutaka Matsuo. Subformer: Exploring weight sharing for parameter efficiency in generative transformers. In Findings of the Association for Computational Linguistics: EMNLP 2021, pp. 4081-4090, 2021.
Sebastian Ruder. An overview of gradient descent optimization algorithms.
arXiv preprint
1609.04600 2016.
Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. On the effect of dropping layers of pre-trained transformer models. arXiv preprint vrXiv:2004.03844 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
Shreyas Subramanian, Vignesh Ganapathiraman, and Aly El Gamal. Learned learning rate sched- ules for deep neural network training using reinforcement learning, 2023. URL https:
/ //openrevie ew. mettfforu? d=00hwu VaCs
Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306. 1695, 2023.
Yi-Lin Sung, Varun Nair, and Colin A Raffel. Training neural networks with fixed sparse masks. Advances in Neural Information Processing Systems, 34:24193-2423 2021.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Sho Takase and Shun Kiyono. Lessons parameter sharing across layers in transformers. arXiv preprint ar arXiv:2104 104:000040602,, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiw:2302.13971 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems. 30, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest be pruned. arXiv preprint 1rXiv: 2019.
Danilo Vucetic, Mohammadreza Tayaranian, Maryam Ziaeefard, James Clark, Brett H Meyer, and Warren J Gross. Efficient fine-tuning of bert models the edge. In 2022 IEEE International Symposium on Circuits and Systems (ISCAS), pp. 1838-1842. IEEE, 2022.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE; A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 1r04.v:8004.07444 2018
Tong Xiao, Yinqiao Li, Jingbo Zhu, Zhengtao Yu, and Tongran Liu. Sharing attention weights for fast transformer. arXiv preprint arXiv: /906.1 1024, 2019.
12

---

Published as a conference paper at ICLR 2024
Zhen Xu, Andrew M Dai, Jonas Kemp, and Luke Metz. Learning adaptive learning rate schedule. arXiv preprint arXiv:1909.09712 2019.
Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efffcient fine-tuning for transformer-based masked language-models. arXiv preprint arXiv:2106.10199 2021.
Weixiong Zhang. Complete anytime beam search. AAAI/IAAI, pp. 425 435-430, 1998.
Michael Zhu and Suyog Gupta. To prune, or not to prune: Exploring the efficacy of pruning for model compression. In 6th International Conference Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Workshop Track Proceedings. OpenReview.net, 2018.
Barret Zoph and Quoc Le. Neural architecture search with reinforcement learning. In International Conference on Learning Representations, 2016.
Barret Zoph and Quoc Le. Neural architecture search with reinforcement learning. In International Conference on Leurning Reprexentations, 2017. URL /tR/:/tss///ppeereieee.eee forum?
id=r1Ue8Hcxq
13

---

Published as a conference paper at ICLR 2024
Table 5: The results of applying method to ViT CIFAR-10
Metric
ViT
Our
Accuracy
0.999
0.995
# trainable params (mean)
630M
80M
# trainable params (end of training)
630M
139M
# trainable layers (mean)
32
5.5
# trainable layers (end of training)
32
7
A A LINE-BY-LINE EXPLANATION OF THE METHOD.
The method is depicted in Alg. In line the replication pattern of the first & steps is determined to be such that layer 0 trains and the other layers are kept fixed at their initialization values. The matehttn action vector and stage vector s are set to the all-zero vector, see line
The method then iterates over the data set and performs regular training step on T, see lines 5-6.
Every & steps (where k << K) obtain a new action a, using the e greeddy policy in Eq. 2 see lines 7-8. In line 9, we compute the state s' based on the obtained action a according to Eq. This state acted upon by replicating and tying the weights lines 12,13. The condition in line ensures that this happens only when this exact replication did not occur in the previous state 8.
We note that a layer that changes state can, based on the condition line 12, either (i) shift from be- ing untied or tied to one layer to being tied to new layer, or (ii) shift to being trained independently. In the first type of shift, a new set of weights would be copied, which may change the transformer much more quickly than through gradient steps. In the second type of shift, the weights are not changed immediately. However, they begin to drift between layers that were previously tied.
In lines 20-21 random mini- matiibbthhhaa is used to estimate the perplexity (PPL) score of T. The reward r is set similar fashion to the reward of DQN (Mnih et al. (2013) as the sum of the evaluation score and the discounted prediction of the next state ((x') (lines 22, 23). We do so by using Bellman (Eq. 4h as an iterative update: 2++1 (8, a) = + max, Qi+1 (s',
The MSE loss is then used to update network Q in lines 24, 25. As mentioned, after every training step of the Q-network, the value of is modified to balance exploration vs. exploitation, see line 27.
B PRELIMINARY COMPUTER VISION EXPERIMENTS.
Since transformers are ubiquitous, evaluating our method only for transformers in the language domain constitutes limitation. As preliminary computer vision experiment, we have ap- plied our method to the Vision Transformer (ViT) (Dosovitskiy et al,, 2021 on the CIFAR-10 dataset (Krizhevsky et al., 2009).
The results are reported in Table 5
As can be seen, similarly to the NLP experiments, 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 at the end of training.
C PRELIMINARY DOWNSTREAM TASKS EXPERIMENTS.
In the domain of NLP, transformers often trained for a causal language modeling task large corpus and are then fine-tuned on a smaller dataset for a specific task such as sentiment analysis, questions answering, named entity recognition.
As a preliminary downstream task experiment, we have taken our GPT-2 based model which was trained using method the 1-billion word dataset and trained it multiple GLUE tasks Wang ct al. (2018). During training on the new tasks, the tied layers were kept as the final state of the model and the language modeling head was replaced with trainable head suited for each task.
14

---

Published as a conference paper at ICLR 2024
Table 6: The results of finetuning the GPT-2 model trained the 1-billion word dataset multiple classification benchmarks. The number of trainable parameters and number of independent layers are the end of training on the 1-billion word dataset and the subsequent finetuning, in which the tying of the layers is fixed.
Metric
Conventional
Our
SST-2 (Accuracy)
0,811
0.799
Cola (Accuracy)
0.691
0.691
QNLI (Accuracy)
0.608
0.599
MRPC (Accuracy)
0.697
0.697
RTE (Accuracy)
0.527
0.541
# trainable params
1.5B
235M
# trainable layers
5
As the vanilla baseline, also trained conventional GPT-2 model the 1-billion word dataset and then finetuned all layers.
The results reported in Table 6 As can be seen, method leads to minimal drop in the given metrics compared to the conventional method, while having only 12% of the trainable parameters and only 5 out of 48 layers untied.
15


Reviews:
[{"review_id": "KyoSs6wxvz", "reviewer": "Reviewer_ccf2", "summary": "In this paper, the authors focus on the problem of reducing the number of parameters in transformers using Reinforcement Learning, specifically a variant of Q-Learning, to dynamically select layers in transformers and tie them together. They evalutate their approach on language models such as GPT-2 and BERT. They demonstrate the performance of their approach both in terms of final model performance and the amount of parameters reduced. They also conduct an ablation analysis of their approach.", "questions": "1. The paper could be strengthened by evaluation on different data domains or perhaps even other neural architectures entirely. How would this approach perform on a vision transformer?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Large transformer models have demonstrated excellent performance, but they are often very expensive in terms of computational resources. The paper addresses a very timely problem: ensuring transformer models can be applied in a practical setting without excessive cost.\n\n2. The proposed method, parameter tying via Q-Learning-is intuitive and sensible approach for this problem.\n\n3. Quantitative results both on the axis of performance and computational resources is convincing.", "weaknesses": "1. The approach is very specific in its focus. The architecture in question are only transformers, and evaluation is only on transformers in the language domain. It is unclear if these results hold for transformers in other data domains (such as vision transformers). It is also unclear if this approach (specifically parameter tying) could work for other neural architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "978Fbd0cfJ", "reviewer": "Reviewer_xrb9", "summary": "This paper proposes a method to reduce the number of trainable parameters of Transformers by dynamically tying the weights of different layers using reinforcement learning during the training process of Transformer. Experimental results indicate that compared to the conventional training method, the proposed approach effectively reduces the number of trainable parameters and slightly improves performance (perplexity score).", "questions": "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).\n\n2 '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?\n\n3 The ultimate goal of reducing training parameters is to decrease memory? What is the actual consumption of memory?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "1 By employing a simple deep reinforcement learning algorithm, the number of trainable parameters of GPT-2 and BERT has been significantly reduced. The algorithm is easy to implement and holds the potential for application in other Transformer-based neural networks.\n\n2 The algorithm has been thoroughly analyzed, providing a detailed explanation of the reasons for its effectiveness.", "weaknesses": "1 The Methods section contains a considerable amount of redundant content, providing a step-by-step explanation of the algorithmic process.\n\n2 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.\n\n3 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0zmtAxoDHm", "reviewer": "Reviewer_Mmpd", "summary": "The authors propose a new means for essentially doing architecture search on a transformer model. The available variation for the model are which layers should be tied together, and which should remain independent. \n\nThey employ an RL algorithm to do the selection of what should be connected to what. The reward signal comes straight from the perplexity of a model after an action has been taken. The results are very exciting and significant. It is very possible that the authors have just opened the floodgates to a new research path on which many will walk in the near future. \n\nThe resulting models demonstrate that a transformer model with 48 layers, may require between 7-10 independent layer weights to perform as well, or better than conventionally trained 48 layer models. This opens up a lot of questions related to modern artificial neural circuitry and may allow for interesting modularization research to take place.", "questions": "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. \n2. 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.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. Clear, concise and precise writing. \n2. The idea attempted itself is intriguing. \n3. The idea is executed very well, resulting in excellent outcomes. \n4. The results are very promising.\n5. The results open up a lot of new questions, as well as motivate the need for research in the pathway the authors have opened up.\n6. The paper represents what could be a seminal moment in architecture search methods, as well as understanding modularity in NNs.", "weaknesses": "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 understanding 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. \n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ocKEAKPstZ", "reviewer": "Reviewer_zEuz", "summary": "The paper proposes to share the weights across layers when training Transformer architecture from scratch. The proposed method utilizes a network Q, which is trained with Q-learning, to assign the sharing policy. Experiments on several text datasets show that the proposed method uses less trainable parameters, reaches comparable performance as the conventional training, and sometimes reaches better training speed.", "questions": "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?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper studies the weight-sharing scheme of training Transformers, and this may be used to reduce the training cost of developing large models.", "weaknesses": "There are two main weaknesses of the paper in my opinion.\n\n1. The related work section is not complete and misses many relevant works and topics. I think several lines of research are very related to this paper but they are missing in the related work: \n\n(1) network pruning (also lottery ticket hypothesis), \n\n(2) weight sharing in Transformers, \n\n(3) parameter-efficient fine-tuning methods (PEFT methods). \n\nNetwork pruning is mentioned a bit in the introduction, but the paper doesn't mention the pros and cons of this paper against network pruning, making the contribution of this paper unclear. I also found that there are several shared-weight Transformers papers [1, 2, 3] that are not cited in the paper, and they should be discussed. The parameter-efficient methods (LoRA, Adapters, Prefix Tuning, Prompt-tuning, etc) save a significant amount of parameters (such as 1 - 5%) and are very relevant to this work.\n\n[1] Subformer: Exploring Weight Sharing for Parameter Efficiency in Generative Transformers\n\n[2] Sharing Attention Weights for Fast Transformer\n\n[3] Lessons on Parameter Sharing across Layers in Transformers\n\n2. Some weaknesses in experiments:\n\n(1) Several weight-sharing techniques should be also included in the comparison.\n\n(2) I would suggest reporting the performance on downstream tasks for a complete comparison. Lower perplexity sometimes does not mean higher downstream performance. \n\n(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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)