Summary
This paper examines how the architecture of a transformer model influences its base capabilities, such as out-of-distribution tasks, transfer learning, and few-shot learning. Specifically, it explores the effects of replacing the feed-forward network (FFN) with a wider FFN (FFN-wide) in various parts of the architecture.
Initially, the authors find that replacing the FFN (referred to as the transformation function) after the attention layer (referred to as the combination function) with FFN-wide leads to worse performance. They analyze this impact by measuring the contribution of the combination function using techniques such as mutual information and token prediction. The results indicate that performance deteriorates in most cases when the contribution of the combination function decreases.
The authors then devise multiple architectures with different width ratios of FFN-wide in the combination and transformation functions. They select the best ratio model architecture, termed CEA (Combination-Enhanced Architecture). This CEA architecture is subsequently used in a Mixture of Experts (MoE) model, demonstrating improvements over the non-CEA MoE model.
Strengths
- The paper addresses an interesting problem by measuring the impact of transformer architecture on the downstream performance of the model, which is crucial for optimizing and understanding transformer models.
- Attributing performance to the inner workings of a neural network, particularly the contribution of specific components, and analyzing this contribution using techniques like mutual information is both novel and insightful.
- It is also interesting to see the impact on two different kinds of model like Bert and GPT and see how the optimal width for the two types of models are different.
- The application of CEA to a mixture of experts models definitely demonstrates the practical benefits of this paper.
Weaknesses
- The performance improvements for the different variations of the architectures seem minor, and the absence of standard deviations makes it difficult to assess the robustness of the results.
- For both BERT and GPT, even after a thorough architecture search, the performance is at best similar to the vanilla models, which is not very promising.
- Similarly for the MoE experiment, it will useful to add the standard deviation of the results. Additionally, since the training loss is lower for MoE with CEA compared to vanilla MoE, it is hard to determine whether the improvements are due to the architecture's inductive bias or just the lower training loss. Similar to the rest of the paper, please also report numbers for this experiments by keeps the same training performance level.
Overall, the paper is interesting, but the performance values being so close and the lack of variance measurement do not give me much confidence in the results.
Minor: In Table 1 and Table 2, the results for the proposed approach are bolded rather than the best approach for that experiment.
In most figures, the y-axis does not start at 0, which can give readers an inaccurate representation of the data.
Questions
- What does it mean for the width ratio to be 0% and 100%? Does this imply that the FFN is removed in these cases?
- What does it mean that FFN-wider models with CEA are not able to beat vanilla models in some cases in Table 1 and Table 2? Isn't the vanilla model one of the ratios explored? If not, would it make more sense to explore the architectures on some other dimensions?
- Please indicate whether lower or higher values are better for all the tables and figures in the paper.
Limitations
The authors can also address the limitations of exploring only a narrow subset of architectures and discuss whether any parts of the proposed methodology can be extended to other aspects of transformer architectures. Additionally, all experiments were performed on smaller-scale models, so the results may not generalize to larger models.