DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging

Authors: matteo  pagliardini
amirkeivan  mohtashami
francois  fleuret
martin  jaggi
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 03:54:32

Abstract: The transformer architecture by Vaswani et al. (2017) is now ubiquitous across application domains, from natural language processing to speech processing and image understanding. We propose DenseFormer, a simple modification to the standard architecture that improves the perplexity of the model without increasing its size---adding a few thousand parameters for large-scale models in the 100B parameters range. Our approach relies on an additional averaging step after each transformer block, which computes a weighted average of current and past representations---we refer to this operation as Depth-Weighted-Average (DWA). The learned DWA weights exhibit coherent patterns of information flow, revealing the strong and structured reuse of activations from distant layers. Experiments demonstrate that DenseFormer is more data efficient, reaching the same perplexity of much deeper transformer models, and that for the same perplexity, these new models outperform transformer baselines in terms of memory efficiency and inference time.

Main paper: 
 introduction: 
 The transformer architecture [31] is the workhorse of modern natural language processing. Recent leaps in the state of the art can be attributed in a large part to efforts scaling this architecture, from millions of parameters [3] to large billion-parameter models [1, 18, 23, 29, 30] . Unfortunately, those larger models come with an increased computational cost, and a large memory footprint. This renders them impractical to use in a wide range of use-cases, limiting who can benefit from them to a handful of big corporations. As an attempt to mitigate this issue, Touvron et al. [29] propose training a smaller model for more steps. However, longer training requires larger datasets which becomes challenging as we are reaching scales where even extremely large datasets fall short of sufficient amounts of data [32] . Furthermore, recent observations suggest that we are reaching the state of diminishing returns where increasing the depth of the model beyond a certain point does not significantly improve performance [21] . Interestingly, a similar state of diminishing returns has been observed in the field of computer vision focused on the training of deep convolutional neural networks. Various solutions were proposed to address this issue, including DenseNets [11] which alleviated the problem by allowing subsequent layers to directly access outputs of earlier layers.In this work, using a similar intuition as DenseNets, we propose the DenseFormer architecture. In particular, instead of only having skip connections from one block to the next, in DenseFormer, a After the first (resp. second) block, the past and current intermediary representations {X 0 , X 1 } (resp. {X 0 , X 1 , X 2 }) are averaged using the first (resp. second) DWA weights [α 1,0 , α 1,1 ] (resp. [α 2,0 , α 2,1 , α 2,2 ]). The DWA weights are supported by red arrows. Those weights are represented in matrix form in (b), for a 12 layers DenseFormer. A DWA module at depth i has i + 1 weights, represented in red. Increasing the dilation sparsifies this matrix, reducing the computational overhead without degrading the perplexity, see Section 3.2 for more details.weighted average of the outputs of all previous blocks is given as the input of the next block. The approach is visually summarized in Fig. 1a .We show that DenseFormers can perform the same as a much deeper standard Transformer model while at the same time being smaller in size, faster, and consuming less memory at inference. More importantly, this is achieved without increasing the amount of required data. As such, DenseFormers are also more data efficient, obtaining much better performance when trained on the same amount of data than a standard model with a similar number of parameters. Our results establish the DenseFormer architecture as an improved version of Transformers for language modeling, encouraging their future use.In addition to providing experimental results on DenseFormer's performance, we also provide additional insights and intuition for their success. Looking at the learned weights of the DWA modules we observe a surprisingly stable pattern in the learned weights that emerges at multiple depths and generalizes across random seeds (see Fig. 4 ). Overall, similar to Huang et al. [11] , we hypothesize that the inter-block connectivity enables the model to more directly re-use early features, without requiring to allocate as much bandwidth to propagate them through multiple layers. Intuitively, this seems to help resolve an ambiguity caused by skip-connections which force the deeper representations to maintain the current token representation while at the same time having to predict the next token. Contributions. Our contributions can be summarized as follows:• Introducing DenseFormer architecture by adding a depth-weighted-average module after each Transformer block. • Demonstrating over different settings (e.g. datasets, batch sizes, sequence lengths) the significantly superior performance of DenseFormer over deeper Transformers, yielding a better speed-performance trade-off during both inference and training. • Providing additional empirically grounded insights and intuition that support the benefits of using DenseFormer. 
related work: 
 While larger models have shown great promise in delivering better capabilities, recent results suggest that the gain from using deeper models faces diminishing returns [22] . Interestingly, the same challenge presented itself earlier for scaling convolutional neural networks [9] . Methods that allow a better flow of information from earlier to later layers such as Residual connections [10] and Highway Networks [27] have been proposed as a successful solution to tackle this challenge, stabilizing training and increasing the threshold where a gain can be observed from increasing depth. Taking it to the extreme, DenseNets [11] demonstrate the benefit of having access to the output of all previous layers. We propose DenseFormer by building on a similar intuition, allowing each block to directly access the output of all previous blocks. Some advantages of attending to representations from earlier layers have already been explored in prior work. Depth-wise Attention [4] suggests adding an attention-like layer before Transformer's final projection layer. This new layer applies attention across the outputs of the Transformer's blocks for the current token (instead of over different tokens as in the standard attention layer). This operation is similar to the weighted averaging block introduced in this work which mixes the outputs of earlier blocks. However, whereas in our proposal the weights are learned during training, Depth-wise Attention computes the weights using the dot product (similar to the attention) which levies a higher overhead. In our experiments, we also show that only a single DWA step before the last layer does not yield the same performance as a full DenseFormer. In another recent relevant work Mohtashami et al. [17] suggest interleaving current and past representations. This allows current tokens to attend to previous representations of themselves and past tokens. Our DenseFormer can be seen as a crude and much more efficient approximation of this mechanism in which we restrict each token to only attend to past representations of themselves, using static (as opposed to dynamic) attention weights. Most relevant to us-in the context of low data training-the concurrent work of Charpentier & Samuel [2] also adds links from the input of each block to representations from earlier layers. An important difference is their use of softmax weights. We conjecture it is important to let the model subtract representations using negative weights (see § 5), and implementing their methods only outperformed the baseline in a low-data regime (see App. B.7). Moreover, they do not provide an efficient implementation of their method.Since its original design by Vaswani et al. [31] , the Transformer architecture used in most applications changed surprisingly little. LLM training is costly, and architecture choices are often conservative. Small variations include changing the activation function [26] , adopting RMSNorm instead of Lay-erNorm, or computing the feed-forward and attention in parallel [20, 29, 33] . More progressive proposals have been made to alleviate computational challenges of the self attention module, such as using kernel methods or other linear approximations [12, 13, 34] , or removing redundant operations without impacting performance [8, 25] . These proposals only affect the internal structure of Transformer blocks. As DenseFormer only adds DWA modules that operate between blocks, we expect that it can be readily used together with these existing proposals.Recent explorations also have shown gains from using multiple language models instead of one. An example is using a mixture of experts, which rely on a routing mechanism [5] to select which expert(s) to use in a given context. Other examples include deploying the same instance of a model in different roles allowing them to debate or provide feedback to each other leading to performance improvements [14, 16] . As these approaches mostly retain the structure of the Transformer architecture and focus on the communication structure between multiple models (or sub-modules), we also expect them to be adaptable to use DenseFormers. 
method: 
 Setup & Notations. We consider the standard Transformer architecture. Given a depth d, it consists in a succession of d Transformer blocks B 1 , . . . , B d , each composed of a self-attention module followed by a single hidden layer Multi-Layer-Perceptron (MLP). We name X 0 , . . . , X d the different intermediary representations, with X 0 being the embedded token sequence, and X i for i ≥ 1 being the output of block B i .We summarize the Transformer architecture as follows:X 0 := Embedding(X) ∀i = 1, . . . d, X i := B i (X i-1 ) Transformer(X) := X d .DenseFormer. The only change to the original architecture is the addition of a Depth Weighted Average module (DWA) after each transformer block. A DWA module at depth i performs a weighted average between (i) the output from the current block B i , (ii) the output of all previous blocks B j<i , and (iii) the embedded input X 0 . The weights of the weighted-average for the DWA i module at depth i are α i,0 , . . . , α i,i . A visual summary can be seen in Fig 1a . The elements of the α matrix are the only additional parameters of our method. More formally, our DenseFormer model can be summarized as follows:X 0 := Embedding(X) Y 0 := X 0 ∀i = 1, . . . d, X i := B i (Y i-1 ) ∀i = 1, . . . d, Y i := DWA i ({X 0 , . . . , X i }) = i j=0 α i,j • X j DenseFormer(X) := Y d .In Section 4 we demonstrate that the DenseFormer architecture can outperform the standard Transformer architecture. In particular, it obtains a much stronger performance (in terms of perplexity) than a model of the same depth, matching the performance of a much deeper model which is both slower at inference and larger in size, leading to a much higher memory footprint than DenseFormer. We further demonstrate the importance of the improved inter-block connectivity brought by the DWA modules in Section 5. We do so by comparing our architecture to a variety of baselines with constrained connections and show those do not perform as well as DenseFormers.Initializing the DWA modules. We note that if α i,i is set to 1 while others are set to 0, the DWA module acts as an identity function, reducing DenseFormer to the standard Transformer architecture. Therefore, we start our training from this initialization.Negligible Model Size Overhead. At depth i the DWA module has i + 1 weights. Therefore, for a DenseFormer of depth d, the total number of additional parameters isd j=1 (j + 1) = d(d+3). For typical model depths (less than 100 blocks), this represents at most an order of 10 3 parameters, which is negligible when compared to the full size of the models.Negligible Memory Overhead. We also emphasize that while DWA requires access to the output of blocks and embedded input X 0 , . . . , X d , these values are stored even when using the standard Transformer architecture. During training, the outputs of the blocks are kept in memory to allow backpropagation, while at inference, only the blocks' outputs for the current token need to be stored. Therefore, the total memory overhead of DenseFormer is negligible.Computational Overhead. Computing the output of the DWA modules increases the computational cost since it requires averaging over multiple large tensors of size (batch size × sequence length × hidden dimension). In this work, we provide an efficient implementation of DWA to reduce the overhead and avoid unnecessary data movement. In addition, we introduce two architectural hyperparameters, which allow building a set of DenseFormer variants that approximate the full DenseFormer. These hyperparameters are DWA dilation and DWA periodicity, respectively introduced in Sections 3.2 and 3.3. We refer to a DenseFormer variant with the dilation factor k and the DWA periodicity p as kxp-DenseFormer. In this notation, the full DenseFormer is a 1x1-DenseFormer.In order to further reduce the computational overhead, we introduce a dilation parameter which sparsifies the DWA weights by periodically setting them to 0. In particular, each DWA module is now given the output of every k-th block, where k is called the dilation factor. More formally, given a DWA module at depth i, a dilation factor of k implies DWA i is only computing a weighted average over {X j |j ≤ i, j ≡ i(mod k)}). See Fig. 1b for a visual explanation. Our dilated DenseFormer can be described as:X 0 := Embedding(X) Y 0 := X 0 ∀i = 1, . . . d, X i := B i (Y i-1 ) ∀i = 1, . . . d, Y i := DWA dilated i ({X j |j ≤ i, j ≡ i(mod k)}) DenseFormer(X) := Y d .As shown in Section 4, we observe no noticeable performance degradation for small values of k (e.g. 2 or 4) while the computational overhead is significantly reduced, leading to much faster inference (a) DWA weights with dilation and DWA period. Figure 2 : (a): DWA weights with dilation and DWA period: For a 12 layers DenseFormer, the α weights are sparsified using dilation k and DWA periodicity p (referred to as kxp). Compared to Fig. 1b , only certain rows have some weights other than the upper diagonal weights (which correspond to the regular transformer information flow). Increasing the dilation and period sparsifies the α matrix, reducing the computational overhead without degrading the perplexity, see Sections 3.2 and 3.3 for more details. (b): Speed and performance trade-off: Comparison of speed and performance trade-off between the standard Transformer architecture and DenseFormer. The number of blocks in each architecture is reported next to the data-point. All DenseFormer models on this plot use a dilation factor of 4. We show results using a DWA period of 1 and 5. Comparing perplexities: Considering only the perplexity (y-axis), a 48 block DenseFormer performs similarly as a much deeper 72 block Transformer. Comparing trade-offs: A 48 block 4x5-DenseFormer matches the better perplexity of a 72 block Transformer while being 1.4× faster at inference. and training. More precisely, a dilation of k reduces the computational overhead induced by DWA modules by a factor of 1/k.An alternative method to dilation for sparsifying DWA weights is adding DWA modules to the architecture less frequently. In particular, we can consider only adding the DWA after every p blocks (instead of after every block as in the standard DenseFormer). We refer to p as the DWA period. A standard DenseFormer has period 1. A DenseFormer with dilation k and DWA period p-referred to as kxp-DenseFormer-can be formalized similarly as a DenseFormer with the following distinction:∀i = 1, . . . d, Y i := DWA dilated i ({X j |j ≤ i, j ≡ i(mod k)}) p | i X i p ∤ iA visual representation of the matrix of α weights can be seen in Fig. 2a . By increasing the periodicity p, we further reduce the computational cost of DenseFormer. In Section 4 we evaluate the effect of increasing the period for a 4-dilated DenseFormer on both performance and speed. We observe that using small values larger than 1 for the period can provide a noticeable boost in speed without noticeable performance degradation. Using a period of p reduces the computational overhead by a factor of 1/p. Hence, a kxp-DenseFormer only has 1/kp of the computational overhead of a regular DenseFormer. In Section 5 we also provide results with other sparsity patterns for the DWA weights (α) but show that using dilation and periodicity works most favorably.Interplay between k and p. The ideal value of p can depend on the dilation k used. For instance, using k = p = 4 implies the DWA module after block 4 will look at {X 0 , X 4 }. Its output, Y 4 , will be sent through block 5 to yield X 5 . However, the next DWA module after block 8 will only look at {X 0 , X 4 , X 8 } (and not at X 5 ). This means that Y 4 will have to go through blocks 5, 6, 7, 8 before being accessible by later DWA modules. In contrast, using k = 4 and p = 5 allows the information to propagate much faster since DWA modules always have access to the processed output of the previous DWA module. This interplay can be visualized in Fig. 2a as well as in Appendix B.2.  
results: 
 We demonstrate the effectiveness of DenseFormer through experiments on language modeling tasks. We compare the performance of DenseFormer architectures with the standard Transformer architecture using model size, inference time, training time, and final perplexity (sometimes abbreviated as PPL) as metrics. For each metric, we consider a baseline that performs the same as DenseFormer on that metric. Concretely, we include the following baselines:Same Depth Baseline. A standard architecture with the same depth as the DenseFormer. It roughly has the same number of parameters as DenseFormer given the negligible number of DWA parameters.Same Inference Time Baseline. A standard architecture that has the same inference time as the DenseFormer. Since adding DWAs to the architecture has a computational overhead, this baseline has more layers (i.e. more capacity) than DenseFormer.Same Perplexity Baseline. A standard architecture with roughly the same PPL as the DenseFormer. It usually is much deeper than the DenseFormer, showcasing the benefits of using DWAs.Same Training Time Baseline. A standard architecture that has the same training time as the DenseFormer. Since adding DWAs to the architecture has a computational overhead, this baseline is trained for more iterations than DenseFormer.Skips with Gains. It can be observed that DenseFormer, among other things, allows scaling the output of the previous layer, providing more control than the original skip connections. Therefore, we provide an additional baseline to show this is not the only benefit offered by DenseFormer and emphasize the importance of having direct access to the outputs of earlier layers. In particular, we consider a modified version of the standard architecture where each skip connection also contains a learned scaling factor which is applied to the the values moving through the skip connection before being summed with the output from a different layer (e.g. self-attention).Experimental setup. We use models with 8 heads, each having 64 dimensions, and train them using batches of 400 sequences of length 256. We use rotary positional encoding [28] . We optimize the model using AdamW [15] with β 1 = 0.9 and β 2 = 0.95 with weight decay factor 0.1. We perform most of our experiments on the OpenWebText2 dataset [6] , an enhanced version of OpenWebTextCorpus [7] with around 17B tokens. We train all models for 40k steps, thus keeping the number of data points used in training fixed. We use learning rate 0.001 with a cosine scheduler and do a warmup in the beginning 5% steps.We present the result of training 48 block and 72 block DenseFormers along with baselines of various sizes in Tab. 1. We make the following observations based on these results:Better perplexity than same depth baseline. When comparing with a baseline of the same depth, DenseFormer significantly outperforms the standard architecture. Moreover, as it can be seen in Fig. 2b , the perplexity of a 48 block DenseFormer is only matched by a 72 block Transformer baseline.Faster than a baseline with the same perplexity. The performance of a 48 block DenseFormer is on par with a 72 block standard architecture. Still, the 48 block DenseFormer is much faster at inference (measured in batches per second) than the 72 block standard architecture. Moreover, the number of parameters and memory footprint of the 72 block baseline is 45% larger than the one of the 48 block DenseFormer.Better perplexity than a baseline with the same inference time. Comparing the 48 block Dense-Former without dilation, with a 64 block standard architecture (which has the same inference speed), shows a wide gap between the higher performance of DenseFormer (17.84) and the standard Transformer architecture (17.94) . Considering DenseFormer models with a dilation of 4 and/or a DWA period of 5 would increase this gap further.Weighted skip-connections are insufficient. DenseFormer is changing the flow of information in the model. We can wonder whether it leverages the additional expressivity or whether the performance gains could be explained by making it easier to rescale the contribution of each block. When comparing the 48 block Transformer baseline to the 48 block skip-with-gains baseline, it seems adding tunable weights to each skip connection does not lead to a significant improvement. When compared with the 48 block DenseFormer, this showcases the importance of having direct access to all previous layers.Faster with dilation and DWA period. Finally, our Tab.1 results show that for small dilation factors k and DWA period p, kxp-DenseFormer perform comparably while significantly boosting speed. Indeed, as can also be seen in Fig. 2b , using 4x1-DenseFormers or 4x5-DenseFormers allows pushing the Pareto frontier on speed and performance trade-off forward.More efficient during training. We train a 48 block 4x5-DenseFormer and compare it against a 48 block Transformer baseline trained with the same time budget. The baseline is therefore trained for more iterations (41.5k vs. 40k) to compensate for the time overhead of DenseFormer. In Fig. 3c we visualize the perplexity (approximated on a small subset of the validation set) dropping as a function of the training time. The DenseFormer's perplexity is dropping faster than that of the Transformer baseline. This shows the superior efficiency of DenseFormer during training. While the Transformer is trained for more steps, thus using more data points, it is still outperformed by the DenseFormer. The final perplexities on full validation set reached by the two models can be seen in Tab. 2.We perform additional experiments to show that our results are general and extend to different settings and larger scales. We also study the impact of the dilation factor k and DWA period p on the efficiency of our kxp-DenseFormer architecture.Experiments with longer sequences. Due to computation limitations, we can not run all experiments at a large scale. We however repeat a limited set of experiments with longer sequences of 512 tokens using a smaller batch size of 128. A 48 block Transformer baseline reaches a final perplexity of 18.28 ± 0.03 against 17.73 ± 0.02 for the DenseFormer. This result shows that the gap between the two architectures persists for longer sequences.Effect of Dilation and DWA period. Fig. 3a and Fig. 3b show the impact of using different combinations of dilation k and DWA period p on the final perplexity, training and inference speed.As can be seen, small values of the dilation factor (e.g. up to 4) have a negligible effect on the perplexity. However, increasing the dilation factor further affects the performance more adversely while the gain in both training and inference speed starts to plateau. Increasing the DWA period also provides a similar trade-off, with the perplexity being barely affected for p ≤ 5. From those figures, we conclude that a dilation of 4 and a DWA period of 5 seem to offer the best compromise between speed and perplexity. In Appendix B.2, we provide more detailed results, including showing how increasing dilation yields a more pronounced speed-up for deeper models, making larger dilation factors more effective in those scales.PG-19 experiments. We also show the superior performance of DenseFormer on the PG-19 dataset [24] . At any depth, DenseFormers outperform Transformers in PPL. See App. B.3 for details.  
analyzing the information flow: 
 In this section, we investigate the learned DWA α weights to gain more insight and intuition on the reason behind the superiority of DenseFormer.A stable weight pattern emerges. We start by visualizing the learnt DWA weights for 48 and 72-block DenseFormer models (with a dilation and period of 1) in Fig. 4 . Interestingly, the α weight patterns learned at both depths are very similar:• High weights are on the diagonal (corresponding to the normal information flow as in a standard Transformer) as well as on the immediate previous blocks. • High weights are given to the initial embedding vectors. Those weights are positive in earlier layers while later layers assign a negative weight. • An aggregation block is observed near the final layers where a high weight is given to all previous layers in the block (seen as a high-weight triangle near the diagonal in the lower right corner).Finally, Fig. 4c shows similar patterns persist to some extent when using dilation factors higher than 1. Similar results hold for kxp-DenseFormers as can be seen in Appendix B.1.Other sparsity patterns during training. Alternative to pruning after training, we also consider imposing different sparsity patterns during training. We already have experimented with such sparsity patterns induced through dilation and DWA period. In Fig. 4 , we observe that in many cases the largest weights are given to the few previous blocks as well as the first block. As such, we experiment with a sparsity pattern allowing DWA to access previous k blocks as well as the input embedding vectors calling it "Last K". Furthermore, given the large magnitude of weights on the last layer, we also experiment with only having a single DWA in the architecture which is placed after the last layer, calling it "Connect to Last". Tab. 3 shows the perplexities when using these sparsity patterns and shows that they do not achieve the boost in performance obtained with DenseFormer. This observation further strengthens the importance of small DWA weights both during training and inference.Correlation with Input Embeddings. Based on the special pattern of weights given to the embedding vectors-especially the negative weights given to the input by the last layers (see Fig. 4 )-we hypothesize that the model tries to use the information in the input in earlier layers while removing the influence of the current token as it tries to predict the next token. In order to test this hypothesis we plot the average cosine similarity of each token's vector after each DWA block with its input embedding in Fig. 5 . As expected based on the weight pattern, we observe that the similarity is high in the earlier layers. At the later stage, the model decreases this similarity significantly.We hypothesize this final decrease is due to the model moving from processing the current token to building the next token's representation. In contrast, the similarity drops down very early for a standard transformer and remains low for the rest of the layers.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We provide the code as well as the general command line to run the experiments. One of the datasets, we used: OpenWebText2, was publicly available. Unfortunately this is no longer the case. Given that someone has access to this data, they can use our code for processing it.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
future work & conclusion: 
 In this paper, we introduced the DenseFormer architecture. This architecture adds an averaging module called DWA after each block which allows it to directly access the outputs of previous blocks.We established the superiority of this architecture over Transformers in terms of perplexity/speed trade-off through experiments in a variety of settings. Additionally, we provided dilation and DWA periodicity as simple methods to improve speed without significantly hurting performance. Finally, we provided insights about the learned weights, revealing patterns persisting over different depths.As the next steps, finding more efficient implementations of DenseFormer is grounds for future work. One possible direction is finding better sparsity patterns that also can be implemented efficiently. The weights visualization in Fig. 4 suggests such patterns might exist. Furthermore, finding efficient ways to shard DWA across multiple nodes is important to allow large-scale distributed training. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: The hyperparameters for training are discussed in the main paper in § 4 as well as in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Error bars are always given, our results are averaged across three seeds.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We describe the hardware used in App. A, as well as discuss the computational time required extensively in the main paper (see. § 4).Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our work proposes a generic replacement for the Transformer architecture. While we focus on language modeling, it does not target any specific domain of application, therefore there is no direct path to a negative application.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: the paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We cite all datasets and libraries used in our work.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this section, we provide several implementations in Pytorch [19] . The first implementations we propose are very simple and rely on looping over the past representations to compute each DWA output. As a result, those naive implementations are slow when dilation and DWA periodicity are not used. In a second time, we propose a more optimized and faster implementation. As a comparison, for a 4x5-DenseFormer with 48 blocks, the naive implementation takes 674ms per training iteration against 657ms for the later implementation. In Appendix A. The point of this more optimized implementation is to remove the for loops of the naive implementation, and instead use tensor operations. To achieve this, we need to create a tensor containing the previous representations instead of a list. Given the structure of this tensor which will simply accumulate all the past representation during the forward pass, we would want to pre-allocate an accumulator and update this accumulator with a new representation after each block. However, this implies using in-place operations that conflict with Pytorch's autograd. We find a workaround by defining the InPlaceSetSlice class. We implement a helper module DWAModules built on top of this new class: In all our experiments we trained on A100-80GB GPUs. In most of our experiments we trained using data-parallelism, distributing the batches over up to 4 GPUs.All of our models share the following hyperparameters:• Max learning rate of 0.001.For different depth, we tested values in {0.0001, 0.0003, 0.0005, 0.0007, 0.001, 0.002} and found 0.001 to be consistently better.• 2000 warmup steps, chosen between {1000, 2000} • A weight decay of 0.1 • AdamW hyperparameters β 1 = 0.9 and β 2 = 0.95 • We use the GPT-2 tokenizer with a vocab size of 50304 • 12 attention heads of 64 dimensions • An embedding size of 768Small α weights matter. In the visualized weight matrix of Fig. 4 , most of the weights on the inter-block connections are small. This observation raises the question as to whether it is possible to drop most of these connections or not. To answer this question, we plot the perplexity after dropping a portion of the smallest DWA weights and report the results in Fig. 6 . It can be seen that even though a large portion of weights in Fig. 4 are small, dropping beyond 15% of the weights leads to a significant increase in perplexity. Therefore, even though these weights are small, they seemingly play an important role in predicting the next token. The figure shows the performance when the model is trained with no sparsity induced (dilation and period of 1) and the DWA weights are later sparsified based on their magnitude at inference. One can see that the perplexity quickly explodes after only sparsifying 15% of the weights. This observation suggests that even though many of the DWA weights are small (as can be seen in Fig. 4 ) they still play an important role in the output of the model.Visualizing αs with periodicity. Similarly to Fig. 4 , we show in Fig. 7 the DWA weights for 4x3, 4x4, and 4x5-DenseFormers. We observe patterns similar to the 1x1-DenseFormer in Fig. 4 but at a lower resolution. Correlation with input embeddings at 72 blocks. As in Fig. 5 , we analyze the cosine similarity between the output of each DWA module and the initial embedding vectors for models of 72 blocks. The results in Fig. 8 are very consistent with those obtained with shallower models (Fig. 5 ).Evolution of DWA weights during training. In Fig. 9 , we plot the DWA weights of a 48 block DenseFormer during training. We observe how the pattern is learned relatively fast, within the first 5000 iterations. More detailed analysis of dilation. For 48 block models, we study the impact of varying the dilation factor k, we do not vary the DWA period which is set to 1. The results of this experiment are in Fig. 10 . We observe how small dilation coefficients do not significantly deteriorate the perplexity yet increase the inference speed. shows the degradation in perplexity as we increase the dilation factor of kx1-DenseFormer models. A noticeable drop in performance occurs for larger dilation factors, e.g. after k = 4. However, surprisingly, 12-Dilated DenseFormer still outperforms the Transformer baseline. As shown in (b), while the perplexity is not so impacted by dilation, the inference speed is significantly improved. Interestingly, the speed gain also plateaus for larger values of k, e.g. roughly k = 4 for 48 blocks. The gain increases with the depth of the DenseFormer, and the plateau threshold occurs later for deeper models.More detailed analysis of the DWA period. For 48 block models, we study the impact of varying the DWA period p. We do not vary the dilation which is set to 4. In Fig. 11 , we observe the impact of increasing p on the perplexity. Interestingly, the perplexity profile is non-monotonic in p, which exposes the interplay between k, p, and the depth of the model. Moreover, increasing the DWA period further increases the inference speed over increasing the dilation. In this section, we study what would happen if we started training the DWA weights at different training iterations. As seen in Fig. 9 , the DWA weights are rapidly converging to their final values within the first 5000 iterations. Moreover, the initialization of the DWA weights corresponds to the same flow of information as in a normal transformer. This raises the question of whether training the DWA weights during the first training iterations is important, or whether a pre-trained model would still gain from adding the DWA weights later. To answer this question we experiment with training the DWA-weights after N iterations. We do not modify the learning rate scheduler or any hyperparameter besides N . Results in Tab. 5 show a diminishing return as N increases. It seems important to tune the DWA weights from the beginning. A possible hypothesis could be that the iterates commit to a valley in the loss landscape relatively early during training. Once deciding to go to the valley where DWA weights are not used, it is difficult to recover and ultimately benefit from newly added DWA weights. We believe this phenomenon could be mitigated using a better learning rate scheduler. We leave this investigation as future work.In this section, we compare the ranks of matrices learned using DenseFormer and Transformer architectures. Our main result in Fig. 12 is that there is no significant difference in rank between the two approaches.In this section, we revisit experiments from the main paper but use a small batch size of 128 instead of 400 during training. Figure 12 : Ranked singular values averaged across blocks. For 48 block models, we average the singular values for each matrix across blocks (except for the embedding matrix). We observe no significant differences between Transformers and DenseFormers. Results are averaged over 3 seeds.Speed and performance trade-off. In Fig. 13 we show the trade-off between inference speed and perplexity for different numbers of blocks. Similarly to Fig. 2b , DenseFormers reach a better perplexity than much deeper Transformer models. Interestingly, the perplexity gap is larger than when using larger batches (compared to a batch size of 400 used in Fig. 2b ). A 48 block DenseFormer is performing on par with a 90 block Transformer. This might indicate that the DenseFormer is more robust to large gradient noise compared to Transformers. DenseFormers reach better trade-offs in terms of inference speed and perplexity. Those results are expected to improve if we were to train a 4x5-DenseFormer instead of a 4x1-DenseFormer. Detailed results can be seen in Tab. 6.Results with other sparse patterns. In Tab. 7 we reproduce the experiments of Tab. 3 but with a batch size of 128. Similar conclusions follow.Visualizing the DWA weights (trained with a batch size of 128). In Fig. 14 we plot the DWA weights obtained when training with a batch size of 128. The learned patterns are very consistent with the ones of Fig. 4 obtained with a larger batch size. In this section we compare our work with the concurrent work of Charpentier & Samuel [2] . We implement the different variants they suggest and show the results we obtain in Fig. 15 . We observe that while their architecture is more data efficient in the low data regime, DenseFormers catch-up quickly (after about 1.2k iterations). In the non low data regime, ELC-Bert performs similarly as the baseline, while DenseFormer is consistently better. Information capacity bottleneck. We can postulate that the larger the width of our model, the easier it is to pass information sequentially from layer to layer. Therefore, the wider the network, the less we expect to gain from adding inter-layer residual connections. We experiment with varying the width of 24 block models, using the following values: {384, 768, 1536}. Results are summarized in Table 8 .We also train a transformer model with a larger hidden size of 1664. For the larger widths of 1536 and 1664, we tuned the learning rate in {0.001, 0.0005, 0.0003} and found 0.0005 to be best. We use a batch size of 128, a learning rate of 0.001, a sequence length of 256, and train for 40k iterations. We observe-as expected-the gap becomes smaller for wider models, yet does not disappear. A DenseFormer with a smaller hidden size of 1536 performs better than a Transformer with a hidden size of 1664. In general, as the complexity of the task increases, followed by the need for more capacity, it is more compute and memory efficient to use a DenseFormer instead of a Transformer. This is especially relevant when we think of deploying LLM on devices with hardware constraints. Justification: Addressing the limitations of the proposed architecture is an important part of our work. In § 3, we discuss the impact on the resources, which we tackle through introducing sparsity in the DWA weight matrix. Finally, in the conclusion, we mention how future works should focus on understanding how this architecture can be efficiently distributed in the context of large scale training. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Justification: No special consideration.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: Given the overall positive reviews, strong empirical results, and the authors' thorough responses to the reviewers' concerns, I recommend accepting this paper. DenseFormer presents a simple yet effective modification to the Transformer architecture, demonstrating its potential to improve the performance and efficiency of various models.

Reviews: review: paper_summary: paper_summary: The paper proposes to construct current transformer block input by weighted average of all inputs from previous transformer blocks. The weights are static and learned during training process. To reduce computation complexity, the method comes with a dialated version controlled by modulo and division relations. Experiments are conducted to observe the changes in model size, train/test speed, PPL. Models are trained from scratch on OpenWebText2. Results show better PPL, or similar PPL with smaller/faster model.
strength_weakness: strength_weakness: not a rigorous study regarding experiment comparison with related work, dialation mechanism design
questions: questions: 1. the abstract is a little misleading, L5 says "100B parameters", but experimetal models are at the level of millions of parameters. 
2. the proposed method should consider a generalization into triangular/pyramid structure. what happens if there is DWA on top of DWA and again on top of DWA? 
3. the sparsity design by modulo and division is rule based (with no intuition on hyperparameter setting), not learned from training. is it possible a simple regularization can be much better?
4. the whole design is meant for traning models from scratch. it becomes not useful given many well trained transformers are released.
limitations: limitations: N/A
paper_summary: paper_summary: The paper introduced DenseFormer -- a simple and effective architecture that boost transformer language model's performance by adding trainable residual connections to all the previous layers. The paper discusses intuitions behind the architecture (that it enhances information flow between earlier and latter layers) and design choices for quality-cost tradeoff. Empirical results on the language modeling tasks validates the effectiveness of the architecture. The paper also presents empirical observations of the residual weights which exhibit consistent patterns that might be tied to the effectiveness of the architecture.
strength_weakness: strength_weakness: A minor weakness is the robustness of the architecture across tasks: the paper focuses on the language modeling task, while transformers have been used for vision, speech and general time series data. A discussion on the robustness of Denseformer on other tasks would strength the paper.
questions: questions: Regarding the memory consumption of DenseFormer, the paper mentioned that it has "Negligible Memory Overhead" since the previous layer's output are stored in the KV-cache. This seems to be at odd with the most common practice, in which the actual key and value vectors are stored in the KV-cache, not the layers' outputs. The reviewer assumes that the authors are referring to a difference way of caching -- caching the layers' output and re-compute key, value projections at inference time, which will be consistent with the negligible memory overhead claim, but introduces additional computation. A clarification on this point would be helpful.
limitations: limitations: NA
paper_summary: paper_summary: The work introduces DenseFormer, a variant of Transformer architecture using special connections between blocks. Each transformer block in DenseFormer may look at a weighted average of all previous blocks' outputs (instead of the simple output of the previous block). Authors run extensive experiments showing the performance improvement of such a design and show that this simple change beats the baseline model on the Pareto-frontier of model quality and speed. Authors also run good ablations and showcase different variants of DenseFormer.
strength_weakness: strength_weakness: **Non-standard model shape (issue essentially resolved during the rebuttal).** By a model shape, I mean the relative dmodel and nlayers. I am extremely worried that results would hold up on standard model shapes. The work uses a dmodel of 768, with the smallest models being 48 layers deep and the biggest model being 90 layers. This is extremely deep and thin.

For reference I will use GPT3 model shapes, as they are reasonably tuned and still used by the community today ( see Table 2.1 on page 8 of https://arxiv.org/pdf/2005.14165 ). The only model with a dmodel of 768 has just 12 layers, 4x shallower than **the shallowest** model used by the authors. This ratio 768/12 was also used by BERT-base. The common ratio used by the community, shared by GPT3 6.7B, llama7B, Mistral, is 4096 dmodel with 32 layers - still shallower than any models used by the authors, and 6x as wide! Chinchilla 80B and LLama 80B use 8192/80, with depth comparable to the authors' models. Only when we get to "the GPT3", 175B, we finally get to the model deeper than any of the authors' models, but with a whopping 12288 dmodel.

The above issue concerns me because one possible interpretation of DenseFormer improvements is essentially increasing the information throughput/capacity of the residual stream, as the layers can communicate directly without using the limited residual stream. With those hyperparameters artificially bottlenecking the throughput of the baseline model's residual stream (compared to the optimal hyperparameters), every technique that effectively circumvents the bottleneck gains an unfair advantage, and the results may not reflect the improvement in the real-world models.

I am unsure how big of an impact this nonstandard architecture has (maybe it doesn't have any impact?), but it leads me to trust the results less. While I trust that running wider models isn't feasible with computational constraints, I would suggest the authors run models of "standard" shape, especially when it comes to the baselines. For example, seeing the results of a smaller model, like 768/12 (and/or less than 12 layers while keeping the dmodel of 768), would showcase the impact of changing the width/depth ratio of the model on the relative improvement achieved by DenseFormer. Those experiments would also be relatively cheap (I'm not suggesting authors pre-train 7B models, of course).

The saving grace for now is that the code exists and everything is easy to reproduce by the community, so I'm leaning accept even with this issue standing. However, running those experiments would be perfect.

**Learning rate, hyperparameter tuning.** Those are less important than the depth/width ratio, but still important. Learning rate (and other parameters, if tuned). How was it chosen? Was it tuned for the baseline? Was it tuned at all? This, again, is important - especially for those deep models, which are generally harder to train and less stable.

**Minor**
1. Number of heads. In line 198, "We use models with 8 heads", while in line 602, "12 attention heads". One of those lines is incorrect (I assume 12 heads is used, though).
2. I needed to see the code to see the details about the residual connection inside the block and where the layer norm is (is Pre-LN used, in particular). Adding a "Block" implementation to the paper's naive implementation code would be nice to make it easier for future readers to check this. That said, the code is very nicely written.
3. It would be worth adding that the memory overhead (line 130) is only negligible if no activation checkpointing is used (and it is commonly used). Also, again, for the inference, the overhead is negligible if no optimizations are used to shrink the KV cache (and again, things like grouped attention are commonly used).
questions: questions: See weaknesses. The primary concern is the architecture shape, and the secondary is hyperparameter tuning. With this primary issue resolved, this would be outstanding work.

Another question is, can we get a more detailed analysis of how exactly the technique differs from previous work, apart from being applied to Transformers? My understanding is that it has multiple novel parts that were adapted specifically, but more detailed comparison would be welcome.
limitations: limitations: Properly addressed, except for the architecture shape (see weaknesses).
paper_summary: paper_summary: The authors propose to use a weighted average of previous layer outputs as an input for the successive layers. The weights are learned parameters and can also take negative values. The weighting module can be coarse: it is enough to insert it only in some layers, and it is enough to attend to a subset of previous layers, which significantly lowers the overhead resulting from the weighting. The authors show perplexity gains on language modelling on OpenWebText2. They analyze the learned weights, and uncover interesting patterns.
strength_weakness: strength_weakness: The evaluation could be stronger: It would be nice to evaluate the zero-shot performance on some downstream tasks (e.g., BLiMP, CBT, PIQA). This should be easy to implement and does not require additional training.

The evaluated transformers are narrow ($d_{model}=768$) and very deep (48 layers for the 378 and 72 for the 548M parameter models). It is unclear how the gains would transfer to a more traditional setup (e.g a 350M GPT-3 has 24 layers). The effects might be smaller with fewer layers if the main role of the DWA module is to provide additional shortcuts for the gradients. The unusually narrow $d_{model}$ can provide an additional advantage for DWA because there is less space to store everything in the residual directly, and the shortcuts can help with that. There are two ways to improve this: either show the gains with a more standard aspect ratio or show that the narrow+deep DWA model is better than a more traditionally shaped DWA model (both of these should be paramer-matched).
questions: questions: In line 133, the authors write, "the outputs are stored in KV cache to facilitate decoding.". Depending on the implementation, this is not what is typically done. To avoid recomputation, the K and V values are stored instead of the residual. However, I agree with the authors that this is not an issue since only the state of a single column has to be stored, which is typically negligible.
limitations: limitations: Already discussed in the weaknesses section. The authors should add some discussion on the limitations of the evaluation.

Comments: title: Thank you for the valuable inputs
comment: We are pleased that the reviewer found our response helpful and sincerely thank them for their decision to increase their score.
We thank reviewers 1YUF and EGgo for their responses and for their support. We also thank the AC for his comment.  We agree with these responses and hope they address reviewer Ljsg's concerns. We remain available to answer any further questions.

Concerning the mention of 100B parameter models in the abstract, we will change it in our next revision.

We also agree with reviewer 1YUF in that comparing scaling laws for DenseFormers and Transformers would be a convincing way to verify that our results generalize to larger model sizes. However, deriving those scaling laws requires running experiments at a scale that needs more computation than is available to us.
We thank the reviewer for increasing their score. We will add the new width experiment to our next revision, along with a description of the hyperparameters training procedure.  

* **Experimental details for the width experiment.** As we changed the width of the models and reduced the depth, we had to re-tune the learning rate. For both Transformers and DenseFormers alike, we ended up using a learning rate of $0.001$ for the widths of $384$ and $768$, and a learning rate of $0.0005$ for models of width $1536$. In the interest of time, we decided to reduce the batch size from $400$ to $128$. As a result, the models are trained on less tokens, which explain the gap in perplexity. Remaining hyperparameters are the same as for our other experiments, i.e. we use $40k$ iterations and a sequence length of $256$.  

* **Concerning the tuning of the learning rate for DenseFormer and Transformer models.** We extensively tuned the learning rates for our Transformer models for all depths. We then mostly applied those learning rates to our DenseFormer models. On a small set of experiments, we verified that the added DWA modules do not seem to be affecting the optimal learning rate value. 

We thank again the reviewer for his valuable feedback, and stay at his disposal in case any further clarification is needed.
I understand reviewer Ljsg's point that results useful for small models might not directly apply to billion-scale or trillion-scale models. However, we should also recognize research that advances the foundational understanding of models and their scaling. Dense connectivity could potentially lead to higher compute per example, which in turn could improve data efficiency. I would appreciate seeing more theoretical and empirical evidence to support this claim.
I would like to thank the authors for the time and effort invested in their rebuttal.

I appreciate their new experiment investigating more standard model shapes and am glad to see that the gains still hold. Thus, I am increasing the score.
I am also responding to reviewers Ljsg and 1YUF as a reviewer.

Specifically, I strongly disagree with reviewer Ljsg's point 4. With such reasoning, we should stop all research that is not fine-tuning and prompting and hope that, somehow, all the remaining shortcomings of the models will be fixed by themselves. This is obviously the wrong direction. In my opinion, the most interesting research is the one that targets fundamental limitations of the models. These typically require retraining of the model. Maybe such research is not directly useful for very applied research but is the most useful for advancing the field.

Regarding points 2 and 3, I think that the fact that high dilations and periods retrain almost all gains of the smaller ones indicates that more dense connectivity or a learned connectivity pattern is not necessary. They would increase complexity and slow down the model.
I am responding to a reviewer as another reviewer. I agree with Reviewer Ljsg in some respects. For example, **(point 1)** I also think the mention of the 100B model in the abstract is misleading, and the same thing should be phrased differently w/o implying the method was tested on such models (replacing the sentence with, e.g., "adding less than a percent of total model params").

I also partially agree with Reviewer Ljsg that **(point 4)** pretraining models are costly, and it is a clear limitation of the technique that it requires pretraining. For sure, I'd like to see experiments on Transformer-to-DenseFormer fine-tuning of the Llama or other models, probably as future work. With that said, I find the arguments about "the cost of training the entire model from scratch" or "the huge training cost hinders many scientific explorations" - **those arguments apply not only to this paper but to all papers on improving LLM pretraining** - so I would not hold that argument against authors. I'd like to see more research on more efficient and better pretraining, after all. I'd suggest to the authors, e.g., extrapolating the performance of the model to larger scales (like comparing scaling laws for DenseFormer vs for Transformer), but such an extrapolation also would require more experiments.

Regarding **(point 2)**, it seems to me that adding DWA on top of DWA would not change the model's expressiveness. Reading Reviewer Ljsg's suggestion, I don't know how the proposed technique would be meaningfully different from what is already tested in the paper.

Regarding **(point 3)**, while I agree with Reviewer Ljsg that regularization-based sparsity would be scientifically interesting, I also agree with the authors that the existing static sparsity pattern seems to capture the majority of the gains. Given this, it seems to me that any kind of regularization-based sparsity would be impractical because it introduces additional complexity both in terms of training (auxiliary losses) and in hardware/performance optimizations. The computational budget for such experiments could be better spent elsewhere (e.g., larger models or more experiments for determining scaling laws).
I want to thank the authors for their response, and **I have raised my score (from 6 to 7)** - as I assume a reasonable response will be given to the remaining questions regarding the training setup (see below).

**On model shapes.** Thank you for providing experimental results. I'm glad to see that you also included a pretty standard architecture 1536/24 (e.g., GPT-Large) and a wider range of model parameters (881M vs. 62M, so a 14x difference, instead of around 2x in the original work). While ideally, I would recommend going for more standard shapes of models, especially with a reasonable width-to-depth ratio, I agree the presented results show that DenseFormer should show gain in practice as well.

Those experiments show considerably worse perplexity for a given model size than experiments in the paper. What was their training length? I assume this must have been the critical difference—or was it something else? It looks like 16k steps or so instead of 40k. Or something else was changed as well?

**On other things.** The rebuttal answered my questions. Just to confirm - the learning rates tuning was done separately for the Transformer baseline and separately for DenseFormer. Am I understanding this correctly?

**Change in score from 6 to 7.** Assuming that the above results (with an explanation of differences in the training setup or explanation of worse perplexity in those wider models), a brief note of hyperparameter tuning like the one in the rebuttal, and the expanded comparison to DenseNet like the one in the rebuttal, will be included in the paper or the appendix for the benefit of readers - I raise my score (from 6 to 7), as I'd like this work to be seen at NeurIPS.
1. parameter amount scale can greatly affect model performance. current experimental results are useful for million level models. however, it may be more useful, less useful, useless or possibly performance-hurting for billion level models. the hypothesis is inclined to being less useful or useless, as large models are harder and much more challenging to improve. 
2. the question is looking to see the potential or the upper bound of performance improvements from using DWA. it is straightforward to test by recursively adding DWA on top of DWA to form a pyramid structure, which should not be an issue considering DWA parameter amount are quite small. 
3. rule based sparsity design is easy to test and start. given sparsity design is an important part of the paper, it takes 2 sections (3.2 & 3.3) of method writing. it is expected that authors would naturally explore other simple design options, such as regularization based methods. 
4. the proposed method seems to only require adding few parameters, but at the cost of training the entire model from scratch. in fact, such limitation is directly or indirectly related to above comment 1/2/3, because the huge training cost hinders many scientific explorations. 

Overall, it is challegening and expensive for authors or for the community to verify its effectiveness on billion level models, which are common in the community. to be clear, this is not a question/requirement on author's experimental equipment setup, but a question on the applicability, generalization, and usefulness (on billion level models) of the method.