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.
Strengths
1. **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.
2. **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).
3. **Improvements on Pareto-frontier.** The significant Improvements achieved with DenseFormer compared to the "vanilla" Transformer should interest both the research community and industry applications.
What 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.
Weaknesses
**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
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.