Title: DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformer architecture

Large language models
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.

DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging

Abstract

The transformer architecture by Vaswani et al. [31] 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 at most a few thousand parameters for large-scale models. 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.

1 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

(a) DenseFormer Architecture.

(b) DWA Weights with Dilation.

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.

2 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 LayerNorm, 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.

3 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:

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:

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.

3.1 Impact on Resources

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 is ∑ d j =1 ( j +1) = d ( d +3) 2 . 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 k x p -DenseFormer. In this notation, the full DenseFormer is a 1 x 1 -DenseFormer.

3.2 Dilated 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:

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.

3

4

5

6

Batches Per Second

16

.

5

17

.

0

17

.

5

18

.

0

18

.

5

Perplexity

48

72

90

104

48

72

90

104

48

72

90

104

4

x

1-DenseFormer

4

x

5-DenseFormer

Transformer

(b) Speed and performance trade-off.

and training. More precisely, a dilation of k reduces the computational overhead induced by DWA modules by a factor of 1 /k .

3.3 Periodic DenseFormer

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 k x p -DenseFormer-can be formalized similarly as a DenseFormer with the following distinction:

Avisual 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 k x p -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.

4 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 .

5

.

0

5

.

5

6

.

0

Batches Per Second (during inference)

17

.

8

18

.

0

18

.

2

18

.

4

18

.

6

Perplexity

1

x

1

2

x

1

4

x

1

8

x

1

12

x

1

4

x

5

4

x

8

4

x

10

4

x

15

4

x

20

k

x

p

-DenseFormer

Transformer

1

.

0

1

.

2

1

.

4

Batches Per Second (during training)

17

.

8

18

.

0

18

.

2

18

.

4

18

.

6

Perplexity

1

x

1

2

x

1

4

x

1

8

x

1

12

x

1

4

x

5

4

x

8

4

x

10

4

x

15

4

x

20

k

x

p

-DenseFormer

Transformer

2

4

6

8

Time (hours)

18

20

22

24

Perplexity

Transformer (41.5k iters)

4

x

5-DenseFormer (40k iters)

Figure 3: Training and inference efficiency of k x p -DenseFormer vs. Transformer. For 48 block models, we compare in (a) the different perplexity/inference speed trade-offs reached by a regular Transformer and k x p -DenseFormers. In the top right corner, the Transformer baseline is the model with the worst perplexity but the fastest at inference. In contrast, the 1 x 1 -DenseFormer in the bottom left corner, is reaching the best perplexity but incurs a cost in inference speed. By varying the dilation k and DWA period p , some k x p -DenseFormer models (e.g. 4 x 5 ) provide most of the perplexity improvement of the original DenseFormer while significantly reducing the time overhead. A similar analysis holds when looking at the training speed in (b) . In (c) , we show the perplexity decreasing during training. The x-axis is time. To compensate for the computational overhead of DenseFormer, we train the Transformer baseline for more iterations, such that the two methods have the same training time budget. We observe how our 4 x 5 -DenseFormer is reaching a better perplexity faster than the baseline.

We perform most of our experiments on the OpenWebText2 dataset [6], an enhanced version of OpenWebTextCorpus [7] with around 17 B tokens. We train all models for 40 k 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 DenseFormer 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 , k x p -DenseFormer perform comparably while significantly boosting speed. Indeed, as can also be seen in Fig. 2b, using 4 x 1 -DenseFormers or 4 x 5 -DenseFormers allows pushing the Pareto frontier on speed and performance trade-off forward.

More efficient during training. We train a 48 block 4 x 5 -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 . 5 k vs. 40 k) to compensate for the time overhead of DenseFormer. In Fig. 3c we

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

(a)

48

block DenseFormer.

0

20

40

60

X

i

0

10

20

30

40

50

60

70

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

(b)

72

block DenseFormer.

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

(c)

48

-block

4

x

1

-DenseFormer.

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.

4.1 Additional Experiments

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 k x p -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.

1

6

12

18

24

30

36

42

48

Depth

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

Cosine similarity

DenseFormer

Transformer

5 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 k x p -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.

6 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.

References

Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language Models are Few-Shot Learners, July 2020. URL http://arxiv. org/abs/2005.14165 . arXiv:2005.14165 [cs].

Charpentier, L. G. G. and Samuel, D. Not all layers are equally as important: Every layer counts BERT. CoRR , abs/2311.02265, 2023.

Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.

ElNokrashy, M., AlKhamissi, B., and Diab, M. Depth-wise attention (dwatt): A layer fusion method for data-efficient classification. arXiv preprint arXiv:2209.15168 , 2022.

Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research , 23(120):1-39, 2022.

Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. OpenWebText2 dataset, as part of 'the Pile: An 800gb dataset of diverse text for language modeling'. arXiv preprint arXiv:2101.00027 , 2020.

Gokaslan, A. and Cohen, V. Openwebtext corpus. http://Skylion007.github.io/ OpenWebTextCorpus , 2019.

He, B. and Hofmann, T. Simplifying transformer blocks. CoRR , abs/2311.01906, 2023.

He, K. and Sun, J. Convolutional neural networks at constrained time cost. In Proceedings of the IEEE conference on computer vision and pattern recognition , pp. 5353-5360, 2015.

He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pp. 770-778, 2016.

Huang, G., Liu, Z., van der Maaten, L., and Weinberger, K. Q. Densely connected convolutional networks. In CVPR , pp. 2261-2269. IEEE Computer Society, 2017.

Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML) , pp. 5294-5303, 2020. URL https://fleuret.org/papers/ katharopoulos-et-al-icml2020.pdf .

Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. In ICLR . OpenReview.net, 2020.

Liang, T., He, Z., Jiao, W., Wang, X., Wang, Y., Wang, R., Yang, Y., Tu, Z., and Shi, S. Encouraging divergent thinking in large language models through multi-agent debate. arXiv preprint arXiv:2305.19118 , 2023.

Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Madaan, A., Tandon, N., Gupta, P., Hallinan, S., Gao, L., Wiegreffe, S., Alon, U., Dziri, N., Prabhumoye, S., Yang, Y., et al. Self-refine: Iterative refinement with self-feedback. arXiv preprint arXiv:2303.17651 , 2023.

Mohtashami, A., Pagliardini, M., and Jaggi, M. Cotformer: More tokens with attention make up for less depth. In Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ NeurIPS 2023) , 2023.

OpenAI. GPT-4 Technical Report, March 2023. URL http://arxiv.org/abs/2303.08774 . arXiv:2303.08774 [cs].

Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A. Automatic differentiation in pytorch. 2017.

Penedo, G., Malartic, Q., Hesslow, D., Cojocaru, R., Cappelli, A., Alobeidli, H., Pannier, B., Almazrouei, E., and Launay, J. The refinedweb dataset for falcon LLM: outperforming curated corpora with web data, and web data only. CoRR , abs/2306.01116, 2023.

Petty, J., van Steenkiste, S., Dasgupta, I., Sha, F., Garrette, D., and Linzen, T. The impact of depth and width on transformer language model generalization. CoRR , abs/2310.19956, 2023.

Petty, J., van Steenkiste, S., Dasgupta, I., Sha, F., Garrette, D., and Linzen, T. The impact of depth and width on transformer language model generalization. arXiv preprint arXiv:2310.19956 , 2023.

Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners, 2019.

Rae, J. W., Potapenko, A., Jayakumar, S. M., Hillier, C., and Lillicrap, T. P. Compressive transformers for long-range sequence modelling. In ICLR . OpenReview.net, 2020.

Shazeer, N. Fast transformer decoding: One write-head is all you need. CoRR , abs/1911.02150, 2019.

Shazeer, N. GLU variants improve transformer. CoRR , abs/2002.05202, 2020.

Srivastava, R. K., Greff, K., and Schmidhuber, J. Highway networks. arXiv preprint arXiv:1505.00387 , 2015.

Su, J., Lu, Y., Pan, S., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. CoRR , abs/2104.09864, 2021.

Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models, 2023.

Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V ., Goyal, N., Hartshorn, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y ., Mao, Y ., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I., Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subramanian, R., Tan, X. E., Tang, B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y ., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models, 2023.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. CoRR , abs/1706.03762, 2017. URL http: //arxiv.org/abs/1706.03762 .

Villalobos, P., Sevilla, J., Heim, L., Besiroglu, T., Hobbhahn, M., and Ho, A. Will we run out of data? an analysis of the limits of scaling datasets in machine learning. arXiv preprint arXiv:2211.04325 , 2022.

Wang, B. and Komatsuzaki, A. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax , May 2021.

Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. CoRR , abs/2006.04768, 2020.

Contents

A Implementation details

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 4 x 5 -DenseFormer with 48 blocks, the naive implementation takes 674 ms per training iteration against 657 ms for the later implementation. In Appendix A.2, we provide a denseformer python package to help turn Transformers into DenseFormers in only 3 simple steps.

TL;DR: If you want to implement your own DenseFormer, simply check the denseformer python package: https://github.com/epfml/DenseFormer

A.1 Naive Pytorch implementation

Naive 1 x 1 -DenseFormer implementation. A naive implementation in Pytorch would consist of storing the output of each block during the forward, and feeding those representations to DWA modules after each block. A pseudocode would look like the following:

Naive k x p -DenseFormer implementation. Here we introduce a naive implementation with dilation and DWA-frequency. The DWA module in the following implementation would be the same as for the 1 x 1 -DenseFormer naive implementation.

A.2 More Optimized Pytorch implementation

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:

denseformer python package. The above module is available in the denseformer package which can be installed through the following link: https://github.com/epfml/DenseFormer . It provides the DWAModules class which orchestrates all the DWA logic given a number of blocks, a dilation factor, and a DWA period. After installing the package, a Transformer can be turned into a DenseFormer in three simple steps:

Listing 4: Faster k x p -DenseFormer implementation using the denseformer package.

A.3 Hyperparameters and Hardware Used

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 768

B Additional Results

B.1 Information Flow

Small α 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.

0

10

20

30

40

50

Sparsity (% of

α

dropped)

17

.

5

18

.

0

18

.

5

19

.

0

19

.

5

20

.

0

Perplexity

DenseFormer

Transformer

Visualizing α s with periodicity. Similarly to Fig. 4, we show in Fig. 7 the DWA weights for 4 x 3 , 4 x 4 , and 4 x 5 -DenseFormers. We observe patterns similar to the 1 x 1 -DenseFormer in Fig. 4 but at a lower resolution.

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

(a)

48

-block

4

x

3

-DenseFormer.

(b)

48

-block

4

x

4

-DenseFormer.

(c)

48

-block

4

x

5

-DenseFormer.

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.

1

6

12

18

24

30

36

42

48

54

60

66

72

Depth

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

Cosine similarity

DenseFormer

Transformer

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

B.2 Analysis of Dilation and DWA Period

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.

1

2

4

6

8

12

Dilation Factor

17

.

8

18

.

0

18

.

2

18

.

4

18

.

6

Perplexity

DenseFormer

Transformer

1

2

4

6

8

12

Dilation Factor

1

.

0

1

.

1

1

.

2

Rel. BPS Improvement

DenseFormer (48 Blocks)

DenseFormer (72 Blocks)

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.

1

3 4 5 6 7 8 9 10

15

20

DWA period

p

17

.

8

18

.

0

18

.

2

18

.

4

18

.

6

Perplexity

DenseFormer 4

×

p

Transformer

(a) Perplexity of 48 block 4 x p -DenseFormer on OWT2

1

3 4 5 6 7 8 9 10

15

20

DWA period

p

1

.

0

1

.

1

1

.

2

Rel. BPS Improvement

B.3 PG-19 Experiments

PG-19 experiments. Tab. 4 shows the performance of DenseFormer against standard Transformer on PG-19, which consists of a large collection of full-length books from Project Gutenberg [24]. We trained both architectures for 48 k steps and used a batch size of 128 instead of 400 . All the other training parameters are kept the same as for OWT2 experiments. On this dataset, we can clearly see the superior performance of DenseFormer.

B.4 Delaying the Training of DWA Weights

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.

B.5 Rank Analysis

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.

B.6 Experiments with a batch size of 128

In this section, we revisit experiments from the main paper but use a small batch size of 128 instead of 400 during training.

Model

N

Perplexity

Baseline Transformer

-

18.61 (0.02)

4

x

5

-DenseFormer

0

17.87 (0.02)

4

x

5

-DenseFormer

1k

17.99

4

x

5

-DenseFormer

2k

18.07

4

x

5

-DenseFormer

4k

18.13

4

x

5

-DenseFormer

6k

18.17

4

x

5

-DenseFormer

10k

18.23

4

x

5

-DenseFormer

20k

18.33

4

x

5

-DenseFormer

30k

18.40

0

250

500

750

Ranked singular values

0

5

10

Intensity

Transformer

DenseFormer

(a) QKV attn matrix

0

200

400

600

800

Ranked singular values

0

2

4

Intensity

Transformer

DenseFormer

(b) Output attn matrix

0

250

500

750

Ranked singular values

0

5

10

15

Intensity

Transformer

DenseFormer

(c) MLP up proj.

0

200

400

600

800

Ranked singular values

0

.

0

2

.

5

5

.

0

7

.

5

Intensity

Transformer

DenseFormer

(d) MLP down proj.

0

250

500

750

Ranked singular values

0

20

40

Intensity

Transformer

DenseFormer

(e) Embeddings

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 4 x 5 -DenseFormer instead of a 4 x 1 -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.

3

4

5

6

Batches Per Second

22

.

0

22

.

5

23

.

0

23

.

5

Perplexity

48

72

84

90

48

72

84

90

DenseFormer

Transformer

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

(a)

48

block DenseFormer.

0

20

40

60

X

i

0

10

20

30

40

50

60

70

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

(b)

72

block DenseFormer.

0

10

20

30

40

X

i

0

10

20

30

40

DWA Index

-

1

.

000

-

0

.

368

-

0

.

135

-

0

.

050

0

.

000

0

.

050

0

.

135

0

.

368

1

.

000

DWA weights (symlog scale)

(c)

48

-block

4

x

1

-DenseFormer.

B.7 Comparison with ELC-BERT

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.

0

5

10

15

20

25

30

35

40

Steps (k)

18

20

22

24

26

28

30

Validation (Subset) Perplexity

Transformer

DenseFormer

ELC Bert

600

700

800

900

1000

1100

1200

1300

1400

Steps (k)

50

100

150

200

250

Validation (Subset) Perplexity

Transformer

DenseFormer

ELC Bert

(a) ELC-BERT vs. DenseFormer

(b) ELC-BERT vs. DenseFormer in low data regime

B.8 Effect of varying the width

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 40 k 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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The superior performances of DenseFormer are reflected in Fig. 2b and Fig. 3. Those figures show that (i) during training, DenseFormers are reaching a lower perplexity faster than a transformer baseline, as well as (ii) during inference, DenseFormer provide a better perplexity for a given inference time budget. See § 4 for more detailed results.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: This work is empirical in nature.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In addition to providing the source code, we give great details on how to implement our approach in App. A. The method is detailed at length in the main paper, see § 3.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.

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.

5. Open access to data and code

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.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

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.

7. 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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

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).

9. 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 ?

Answer: [Yes]

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).

10. 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).

11. 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.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [No]

Justification: We share source code and guidelines to reproduce our experiments, but we do not share any other asset.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe method is easy to use, improves perplexity, and shows interesting and consistent weight patterns. The overhead is minimal, less than for depth-wise attention, and the perplexity in a large data regime is better than that of the concurrent DenseNets.\n\nWEAKNESSES:\nThe 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.\n\nThe 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).\n\nQUESTIONS:\nIn 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.\n\nLIMITATIONS:\nAlready discussed in the weaknesses section. The authors should add some discussion on the limitations of the evaluation.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. **The idea.** The idea behind the work is simple, in a very positive way, with a good story. The implementation is straightforward (even if optimizations are useful), the methods aren't costly during training, and what costs it incurs it makes up for in performance. Moreover, I believe the technique introduced could not only improve Transformer-based language models but Transformers in different modalities as well, and even other kinds of residual networks apart from Transformer - not only \"classical\" CNN-based ResNets, but also new architectures like Mamba.\n2. **Experimental setup.** Experiments are extensive, and results show consistent improvement. Every metric that I'd want to see is present and measured, like inference throughput, training speed, and comparison to models equal in inference/training time (both through an increase in size and through an increase in training iterations).\n3. **Improvements on Pareto-frontier.** The significant Improvements achieved with DenseFormer compared to the \"vanilla\" Transformer should interest both the research community and industry applications.\n\nWhat is more, **the authors provide excellent, easy-to-understand code**, so even when I had doubts about how things were implemented, I could easily check the source code; thank you so much.\n\nWEAKNESSES:\n**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.\n\nFor 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.\n\nThe 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.\n\nI 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).\n\nThe 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.\n\n**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.\n\n**Minor**\n1. 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).\n2. 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.\n3. 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).\n\nQUESTIONS:\nSee weaknesses. The primary concern is the architecture shape, and the secondary is hyperparameter tuning. With this primary issue resolved, this would be outstanding work.\n\nAnother 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.\n\nLIMITATIONS:\nProperly addressed, except for the architecture shape (see weaknesses).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOriginality: while the idea of residual connection or trainable residual connection have been proposed, it is still surprising that a simple trainable combination of residual streams can be both effective and has consistent weight patterns. The design of this architecture and the findings of the effectiveness and weight pattern are novel.\n\nClarity: the paper presents the idea in a clear and easy-to-follow fashion. Various design choices are discussed regarding quality-cost trade-off; and important issues such as whether the architecture remains effective if introduced in latter training stages are discussed.\n\nSignificance: the simplicity and effectiveness of the method is quite appealing for consideration in practical language models.\n\nWEAKNESSES:\nA 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.\n\nQUESTIONS:\nRegarding 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.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: 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.

STRENGTHS:
the method is well defined, and experimental results are good

WEAKNESSES:
not a rigorous study regarding experiment comparison with related work, dialation mechanism design

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:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-25