How does Architecture Influence the Base Capabilities of Pre-trained Language Models? A Case Study Based on FFN-Wider and MoE Transformers

Authors: xin  lu
yanyan  zhao
bing  qin
liangyu  huo
qing  yang
dongliang  xu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:38:51

Abstract: Pre-trained language models have been proven to possess strong base capabilities, which not only excel in in-distribution language modeling but also show powerful abilities in out-of-distribution language modeling, transfer learning and few-shot learning. Unlike existing work focusing on the influence of scale on base capabilities, our work examines the influence of architecture on those. Specifically, our concern is: How does architecture influence the base capabilities of pre-trained language models? In this work, we attempt to explain and reverse the decline in base capabilities caused by the architecture of FFN-Wider Transformers, seeking to provide some insights. Through analysis, we found the contribution ratio of Multi-Head Attention (a combination function) to pre-trained language modeling is a key factor affecting base capabilities. FFN-Wider Transformers reduce the contribution ratio of this combination function, leading to a decline in base capabilities. We confirmed this by experiments and proposed Combination Enhanced Architecture (CEA) to address the decline in base capabilities of such models. Significantly, we extended our explanation and CEA to Mixture of Experts (MoE) Transformers. We successfully achieved significant improvements in base capabilities on a 14B parameter MoE model, demonstrating the practical application value of our work. This also indicates that our analysis has a certain guiding significance for architecture analysis, architecture improvement and architecture design.

Main paper: 
 base capabilities: 
 Pre-trained language models not only address in-distribution language modeling but also unexpectedly show strong base capabilities. In this work, we focus on the following base capabilities:Out-of-Distribution Language Modeling This reflects the out-of-distribution generalization capability of pre-trained language models. Models that learn more essential language features often outperform others, which is a good measure of the base capabilities of pre-trained language models.Transfer Learning This is a recognized base capability of pre-trained language models. The works proposing GPT [24] and BERT [7] have established the "pre-training and fine-tuning" transfer learning paradigm under the Transformer architecture.Few-shot Learning This is also a recognized base capability. Some works [25, 3] found large-scale pre-trained language models could complete many NLP tasks with no or only a few demonstrations, revealing their remarkable few-shot learning capabilities.In this work, we primarily explore how the architecture of a model influences its base capabilities in pre-trained language models, which requires designing a sound approach for quantitative analysis.Comparing the base capabilities of one architecture model against another is straightforward: we simply compare their performance on a variety of representative tasks. However, changes in base capabilities cannot be directly attributed to changes in architecture alone, as model capabilities are also influenced by factors like scale.This leads to a question: under what conditions can we compare the base capabilities of two different architecture models and be more confident that the differences are due to architecture variations?Our scheme to this question: we compare the models when they use the same pre-training data and objectives, and have similar levels of pre-training performance (i.e., language modeling performance on an in-distribution development set). The underlying rationale: two different architecture models, when trained on the same corpus and objectives, both gain base capabilities by reducing the loss in pre-trained language modeling. If one architecture model achieves greater base capabilities with the same reduction in language modeling loss, it is highly likely that this additional benefit stems from the inductive bias of its architecture. In other words, when pre-training performance is aligned, differences in base capabilities of models are likely reflecting the impact of architecture inductive biases.This scheme is more appropriate for cross-architecture analysis compared to aligning pre-training steps, parameter numbers, or computational load, and we explain it in Appendix A.Unless specified otherwise, all comparisons of base capabilities across architecture models in this work are conducted under the condition of aligned pre-training performance.The experiment with MoE transformers is an exception, because we need to demonstrate practicality. Therefore, in the comparative experiments between the vanilla MoE model and the improved MoE model, not only is pre-training performance aligned, but pre-training steps, parameter numbers, and computational load are also kept consistent.Due to subsequent experiments involving out-of-distribution language modeling evaluation, transfer learning evaluation and few-shot learning evaluation, it is necessary to clarify the selection of pretraining corpus and out-of-distribution test corpus, as well as the selection of downstream tasks. We provide a detailed introduction in Appendix B. We conducted experiments on various models with the intermediate dimension set to 32d, aligning pre-training performances with those of vanilla models. The results are shown in Figure 2 . 
ffn-wider transformers vs. vanilla transformers: 
 We found that, at the same level of pre-training performance, the Transformer models with wider FFN layers exhibit a noticeable decline in performance on most downstream tasks, indicating a deterioration in their base capabilities. This presents us with a good research object.4 Why FFN-Wider Transformers Have Worse Base Capabilities? 
combination and transformation: 
 Transformers consist of MHA and FFN layers. Considering a certain position in a sequence, the updated representation obtained after an MHA layer is a combination of the entire sequence context; whereas the updated representation obtained after an FFN layer is from a transformation of that position representation alone, context-insensitive.When considering models composed of only one MHA layer or one FFN layer, we find that they can both be used for language modeling. If treated as black boxes focusing only on input and output, there is no difference. However, the way they accomplish language modeling is different, or rather, their inductive biases differ: the FFN layer directly maps the previous token to the target token, representing a one-to-one transformation function; while the MHA layer uses the entire sequence to calculate the target token, representing a many-to-one combination function. The latter aligns more closely with the combinability of language.Further considering the multi-layer stacking of MHA and FFN layers in a Transformer model, we easily understand that although all layers contribute to the final language modeling objective, they do so with different inductive biases. Among these, the MHA layer might be the central embodiment of the model architecture's expression of the combinability of language.With this understanding, when we revisit the scenarios where the FFN layer is widened or its capacity is increased, we have reason to suspect this will lead to a change in the actual contribution ratios of the transformation and combination function, thereby affecting the model's expression of the combinability of language, which is ultimately reflected in changes in basic capabilities.Our hypothesis: the actual contribution ratio of the MHA layer (a combination function) is a key factor affecting the model's base capabilities. The FFN-Wider Transformers directly enhance the FFN layer, indirectly reducing the combination function's actual contribution ratio to pre-trained language modeling, thereby leading to a significant decline in base capabilities.To verify our hypothesis, we first need to confirm whether the actual contribution ratio of MHA layers truly decreases, requiring quantitative analysis.The first method is Mutual Information (MI). For positions in the sequence where a predicted token is to be output, we can focus on the representations of these positions after each layer and calculate the MI between these representations and the target tokens. Then, since we have MI before and after passing through any layer, we can calculate the contribution ratios of the MHA and FFN layers.We adopted the mutual information estimate method proposed by [32] , which mainly involves converting the representations into discrete variables through clustering, and then calculating the mutual information. For the definition of MI and more specific details, please refer to Appendix F.We analyzed four small-scale models (H=128): a vanilla BERT, an FFN-Wider BERT, a vanilla GPT and an FFN-Wider GPT. The pre-training performances of the vanilla models and the FFN-Wider models are aligned. We plotted the MI results in Figure 3 , where the blue, orange and grey lines represent the cumulative MI increment contributions of the MHA layer, FFN layer and Block, respectively. It can be seen that in the FFN-Wider models, the MI contribution of the FFN layer is significantly higher than that in the vanilla models, preliminarily validating our hypothesis.Since the MI estimate requires clustering a large number of representations and is costly when the hidden dimension is high, we tried another more direct method called Token Prediction (TP). It involves predicting tokens directly from the hidden representations. The approach is as follows:We still obtained representations from each layer, first dividing them into sets based on the corresponding output token. Then, we normalized all representations within a set and use their mean as the category vector for that token. During this process, we eliminated tokens where the set size was smaller than 50. For the representations of each layer, we then calculated the cosine similarity with the category vectors of tokens in that layer, selecting the token with the highest similarity as the prediction result for that representation. In this way, we could calculate a token prediction accuracy for each layer. The subsequent approach was identical to that of the MI method.We analyzed four small-scale (H=128) and four large-scale (H=768) models concurrently. We plotted the accuracy increment contribution ratio of the FFN layer for each model, as shown in Figure 4 . It can be seen that the contribution of FFN layers in FFN-Wider models remains higher than that in vanilla models, reaffirming our hypothesis. 
combination adjustable architecture: 
 Although the FFN layer in FFN-Wider Transformers has a higher contribution ratio, we cannot assert this is the reason. This may merely be a correlation rather than a decisive factor in base capabilities.Therefore, we designed an architecture that can directly intervene in the contribution ratios of the transformation and combination function, named Combination Adjustable Architecture (CAA).The new architecture, as shown in Figure5, differs from the FFN-Wider Transformer in two aspects: one is partially transferring the FFN into the MHA; the other is adding a direct pathway inside the MHA that bypasses the Inner-FFN, which we will introduce in detail below.A wider FFN would increase the contribution ratio of transformation function to pre-trained language modeling. A natural idea: is it possible to devise a variant FFN serves only combination function?We found the primary reason that the FFN does not serve only combination function is due to the residual connections. The representation transformed by the FFN does not necessarily have to go through the MHA combination process; the FFN can bypass the MHA via residual connections and directly transmit the information in the representation to subsequent layers.Therefore, we divided an FFN into two parts with adjustable width ratios: one remains in its original position as a transformation function, known as the Outer-FFN, and the other is relocated within the MHA layer, transformed an Inner-FFN. This design ensures that the representation transformed by the Inner-FFN must undergo the MHA combination process before proceeding, initially achieving the goal of the Inner-FFN serving only for the combination function.Although we considered the impact of residual connections and designed an Inner-FFN, the Inner-FFN still has a hidden pathway to bypass the combination function and directly transmit uncombined information to subsequent layers.The attention mechanism involves weighted summation over values, and typically, the value at the current position also participates. This provides a possibility for the Inner-FFN to directly transmit information. Therefore, we designed a direct pathway in the MHA for current position computation to circumvent the Inner-FFN. Specifically, during the self-attention computation at the In this section, we control the width ratio of the Outer-FFN, reducing it gradually from 100% to 0%,confirming our hypothesis by examining the trends of contribution ratio and base capability.We conducted experiments on large-scale (H=768) BERT and GPT models. Initially, a series of models with different width ratios were pre-trained, ensuring their pre-training performance was aligned. Then, we calculated the contribution ratio of Outer-FFN using Token Prediction (TP) and plotted these trends, as shown in Figures 6(a ) and 6(c).It can be seen that the contribution ratio of Outer-FFN decreases as its width ratio is reduced. Conversely, the contribution ratio of the combination function increases with the reduction in the width ratio of Outer-FFN, demonstrating that controlling the width ratio indeed directly influences the contribution ratio of the combination function.We tested the performance of each model above in out-of-distribution (OOD) language modeling and plotted the trend in Figures 6(b ) and 6(d).For BERT, the OOD performance improves as the Outer-FFN width decreases. The performance of the optimal model has already surpassed the vanilla BERT. Most models not only outperform the FFN-Wider BERT in terms of performance but also require fewer pre-training steps. For GPT, only the model without Outer-FFN showed anomalous results, other models also follow this pattern.After the analyses, a conclusion is drawn from the overall trend: the model's base capabilities generally improve as the contribution ratio of combination function increases. It holds over a wide range of width ratios and proves our hypothesis.However, this overall trend does not mean the contribution of transformation function should be reduced to zero. A minor contribution from transformation function might still be crucial, as indicated by the sole anomalous result in GPT. BERT was unaffected, which is related to bidirectional attention still being able to leak few transformation contribution. More details are in Appendix G.  
combination enhanced architecture: 
 In this section, we determine the Combination Enhanced Architecture (CEA) and pre-train more steps to verify whether it can reverse the decline in base capabilities of FFN-Wider Transformer.We selected the optimal width ratio for Combination Enhanced Architecture (CEA). The Outer-FFN width ratio for FFN-Wider BERT w/ CEA is set to 0%. For FFN-Wider GPT w/ CEA, it is 12.5%.We conducted more pre-training steps on vanilla models, FFN-wider models, FFN-wider models w/ CEA, and aligned the pre-training performances.For BERT, we tested fine-tuning performance on GLUE, SuperGLUE and other datasets. For GPT, we tested zero-shot and one-shot performance on multiple datasets. For both, we tested OOD language modeling on Pile. For small-scale models, we removed some experiments beyond their capabilities.The brief results are shown in Table 1 and 2, and the detailed results are in Appendix H.The results show the FFN-wider models w/ CEA improve in base capabilities, not only surpassing the FFN-wider models in most aspects but also nearly reaching the level of vanilla models. In addition, the new architecture models w/o the direct pathway in MHA show a significant decline in base capabilities. These all confirm our explanation.Additionally, we conducted similar experiments for other width ratios that can align pre-training performance and steps simultaneously, which can indisputably prove improvements come from architecture, as detailed in Appendix I. 
from ffn-wider transformers to moe transformers: 
 The Mixture of Experts (MoE) Transformer [18, 8] is a practical architecture that introduces a mix of experts, which enables expanding the capacity with lower computation expense.However, base capability decline also be observed in MoE Transformers. [8] found MoE models perform lower on SuperGLUE compared to vanilla models when achieving same pre-training level. Similar issue can also be analyzed from the results of [1] .We found the MoE layer can be seen as an enhanced version of the FFN layer. Therefore, we believe the previous explanations could also apply to MoE Transformers. Consequently, we directly transplant our CEA to MoE Transformers, resulting in a new MoE architecture.We selected a 1.3B GPT model as the backbone model, which incorporates Rotary Embedding [30] and RMSNorm [37] on the GPT3 [3] . while retaining the original FFN layer as the Outer-FFN layer. To accommodate FlashAttention-2 [6] , we simplified the direct pathway. Instead of preventing transformation leakage by replacing the key and value at current position, we chose to directly mask the current position to achieve the same effect. The specifications, pre-training procedures and few-shot learning procedures are detailed in Appendix C and E.We conducted pre-training on CC and C4 subsets within SlimPajama [29] , training all models from scratch with 100B tokens. Due to the de-duplication across subsets achieved by SlimPajama, we used the remaining subsets as out-of-distribution sets. We carried out comprehensive out-of-distribution language modeling and few-shot learning tests, with the results presented in Table 3 .From the results, it is evident that our method significantly enhances the base capabilities of the MoE model. This fully demonstrates the effectiveness of our analysis and improvement methods. 
limitations 9 conclusion: 
 This work explores how architecture affects the base capabilities of pre-trained language models. FFN-Wider Transformer is our research object and we try to explain and reverse the decline in base capabilities caused by its architecture. We found the contribution ratio of combination function is a key factor, while FFN-Wider Transformer reduces it, leading to a decline in base capabilities. We solved it by proposing CEA. In addition, we extended our conclusion to MoE Transformers, proving our work can offer guidance for architecture improvement. 
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: We do not release models. 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: [NA] Justification: We do not release models. 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. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: N/A. 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 obtained from people whose asset is used. • 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. 
: 
 received sufficient attention. According to the basic principle of inductive bias in machine learning, model architecture is also a crucial factor affecting base capabilities, and its impact could be equally decisive.Some studies have already noted the significant influence of architecture on the base capabilities. For instance, [31] have observed considerable differences in base capabilities among various Transformer architecture variants. Some variants, though larger in scale and more powerful in pre-training performance than vanilla Transformers, exhibit significantly reduced performance in downstream tasks. This suggests different architecture models vary greatly in converting pre-training performance into base capabilities. Simply increasing scale does not resolve all issues, and exploring the impact of architecture on base capabilities is crucial.However, despite these studies demonstrating the key influences of model architecture on base capabilities, the understanding of the underlying mechanisms of these influences remains limited. In this work, we attempt to explain the base capability change caused by a specific model architecture change and identify the underlying influencing factor, then design experiments to validate our explanation and influencing factor, and finally propose a generalizable enhancement method.Specifically, we first focus on FFN-Wider Transformers. This architecture has wider FFN layers, but we found that the change leads to a significant decrease in base capabilities compared to vanilla Transformers. As shown in Figure 2 , under similar pre-training performance, the FFN-Wider models exhibit a noticeable decline in base capabilities compared to the vanilla models. We believe such a simple change, leading to significant differences in base capabilities, makes for a good object of study in exploring how architecture impacts base capabilities.Then, we attempted to explain this change in base capabilities. Through analysis, we concluded that the MHA (Multi-Head Attention) layer in the Transformer is a combination function, and the FFN (Feed-Forward Network) layer is a transformation function, with the former being a focused expression of the combinability of language. During their contribution to pre-trained language modeling, the actual contribution ratio of the MHA layer, as a combination function, is a key factor affecting the model's base capabilities. The FFN-Wider Transformer models directly enhance the FFN layer, indirectly reducing the combination function's actual contribution ratio to pre-trained language modeling, thereby leading to a significant decline in base capabilities.To validate our explanation for it, we designed a Combination Adjustable Architecture (CAA), as depicted in Figure 1 (a). This architecture bifurcates a wider FFN into two parts with adjustable width ratios: one remains in its original position as a transformation function, known as the Outer-FFN, and the other is relocated within the MHA layer, transformed through a special design into an Inner-FFN that solely enhances the combination function. We controlled the width ratio of the Outer-FFN, reducing it gradually from 100% to 0%. In Figure 1 (b), we observed that the actual contribution ratio of the Outer-FFN progressively decreased, indicating a corresponding increase in the actual contribution ratio of the MHA layer. At the same time, we also observed a gradual improvement in base capabilities. These reveal a key phenomenon that confirms our explanation: as the actual 2 BackgroundSubsequent sections will involve specific analyses and experiments of FFN-Wider and MoE transformers. Therefore, in this section, we introduce the relevant background of base capabilities, as well as the unified evaluation schemes and evaluation tasks. 
 
 Appendix: Aside from the pre-training performance alignment scheme we adopted, there are three other schemes. However, none of these are suitable for our work.Pre-training steps alignment Differences in model's base capabilities might be due to differences in model capacity, not solely from architecture changes. Especially in models with different architectures and parameter scales, models with more parameters have a capacity advantage, rapidly reducing pre-training loss to a low level, showing good metrics on various tasks, and appearing to have strong base capabilities. But this is more due to large model capacity, and not much related to architecture inductive biases.Parameter numbers alignment It imposes restrictions on the model that make it inapplicable to any two different architecture models. More critically, this scheme has clear counterexamples, as it is not suitable for ALBERT models [17] . Under the same parameter numbers, ALBERT models exhibit strong base capabilities, but this is due to the large computational load, not much related to architecture inductive biases.Computational load alignment It imposes restrictions on the model that make it inapplicable to any two different architecture models. More critically, this scheme also has clear counterexamples, as it is not suitable for MoE models [18, 8] . Under the same computational load, MoE models demonstrate strong base capabilities, but this is due to large parameter numbers and model capacity, not much related to architecture inductive biases.Although these schemes have shortcomings, they also have reasonable aspects. To indisputably prove base capability improvements come from our architecture changes, we also conducted experiments that align pre-training performance, pre-training steps, parameter numbers and computational load simultaneously, as detailed in Appendix I.In addition, the experiment with MoE transformers also adopts a stricter setting similar to the above. In the comparative experiments between the vanilla MoE model and the improved MoE model, not only is pre-training performance aligned, but pre-training steps, parameter numbers, and computational load are also kept consistent.All experiments with FFN-Wider transformers follow the unified settings described here:The pre-training corpus is consistent with that of BERT [7] , namely Wikipedia and BooksCorpus. We partition a portion of the corpus to serve as an in-distribution development set, and use it to align pre-training performance.We conduct experiments on four specifications: BERT (H=128), BERT (H=768), GPT (H=128) and GPT (H=768). The specifications and pre-training procedures are detailed in Appendix C.For out-of-distribution language modeling capability, we evaluate on the development set of Pile dataset [9] .For transfer learning capability, we only evaluate BERT models on GLUE [34] , SuperGLUE [33] , HellaSwag [36] , PIQA [2] , OpenBookQA [20] , ARC Easy & Challenge [5] and WinoGrande [27] .The experimental settings are detailed in Appendix D.For few-shot learning capability, we only evaluate GPT models. Limited by the maximum sequence length (128) of our pre-trained models, we conduct 0-shot and 1-shot experiments on HellaSwag, PIQA, OpenBookQA, ARC Easy & Challenge, WinoGrande, Winograd [19] , COPA [26] and StoryCloze [21] . The experimental settings are detailed in Appendix E.All experiments with MoE transformers follow the unified settings described here:The pre-training corpus is the CC and C4 subsets within the SlimPajama dataset [29] . The specifications and pre-training procedures are detailed in Appendix C.For out-of-distribution language modeling capability, due to the de-duplication across subsets achieved by SlimPajama [29] , we used the remaining subsets (Arxiv, Book, Github, Stack and Wiki) as out-of-distribution test sets.For few-shot learning capability, we conduct 0-shot experiments on LAMBADA [23] , and conduct 5-shot experiments on MMLU [10] , OpenBookQA, ARC Easy & Challenge, BoolQ [4] , RACE Middle & High [16] , SIQA [28] , SCIQ [35] , HellaSwag, COPA, PIQA, StoryCloze, WinoGrande and Winograd. The experimental settings are detailed in Appendix E.All experiments were conducted in English only. We concurrently pre-trained both BERT and GPT models, where the architecture design of the BERT model adheres to the work of [7] , and the GPT model structure follows that established by [24] . Both models employ a post layer normalization scheme.We trained small and large scales of both BERT and GPT. The small-scale models have a hidden dimension of 128, 12 layers, and 2 attention heads; the large-scale models possess a hidden dimension of 768, 12 layers, and 12 attention heads. In vanilla BERT and GPT, the FFN intermediate dimension is 4 times the hidden dimension, while the ratio of FFN-Wider models is 32 times. The small-scale vanilla models have 6.3M Parameters and the small-scale other models have 17.3M Parameters. The large-scale vanilla models have 110M Parameters and the large-scale other models have 506M Parameters. All models have a maximum sequence length of 128 and utilize the BERT vocabulary released by [7] , comprising 30,522 tokens. We used PyTorch2 and transformers3 libraries.The BERT model employs a masked language modeling task for pre-training, masking 15% of tokens in the sequence-80% are replaced with [MASK], 10% with random tokens, and 10% remain unchanged. Differing from [7] , we removed the next sentence prediction task, using long and continuous text for pre-training inputs and applying different masking schemes to the same input sequence in various epochs. The GPT model is pre-trained using a language modeling task without additional special settings.The maximum epoch set for all model pre-training is 40, but in practice, it was not reached; midtraining checkpoints were used for alignment and experimentation. All models used the Adam optimizer [15] for pre-training, with a learning rate of 1e-4, β 1 = 0.9, β 2 = 0.999, L2 weight of 0.01, a warm-up over the first 10,000 steps, followed by linear decay. The small-scale models were pre-trained with a batch size of 512 on four Nvidia Tesla V100s, and total GPU days are approximately 55 days; the large-scale models with a batch size of 1024 on four Nvidia Tesla A100s, and total GPU days are approximately 87 days.To fairly compare the base capabilities of different architecture models, all models were pre-trained from scratch. However, as mentioned earlier, due to limited computational resources, our pretraining steps generally fell short of those in the original papers [7, 24, 25, 3] , which may result in discrepancies in downstream task performance compared to the original research.All experiments were conducted in English only.We selected the GPT model as the backbone model (Vanilla GPT 1.3B), which incorporates Rotary Embedding [30] and RMSNorm [37] on the GPT3 [3] . All models have a maximum sequence length of 2,048 and utilize the Mixtral vocabulary released by [13] , comprising 32,000 tokens. We used PyTorch, transformers and fastmoe4 libraries.These models are pre-trained using a language modeling task without additional special settings.We conducted pre-training on CC and C4 subsets within SlimPajama [29] , training all models from scratch with 100B tokens. Due to the de-duplication across subsets achieved by SlimPajama, we used the remaining subsets as out-of-distribution sets. All models used the Adam optimizer [15] for pre-training, with a learning rate of 5e-4, β 1 = 0.9, β 2 = 0.95, L2 weight of 0.01, a warm-up over the first 2,500 steps, followed by linear decay. All models were pre-trained with a batch size of 256 on 8 Nvidia Tesla A100 80G cards, and total GPU days are approximately 280 days.We conducted fine-tuning experiments on BERT models across various datasets. For the small-scale (H=128) models, we only conducted experiments on GLUE, while for the large-scale (H=768) models, we experimented with all datasets.The maximum number of training epochs during fine-tuning was 10, with a batch size of 32. The optimizer was Adam [15] , with a warmup ratio of 0.06, a linearly decaying learning rate, and a weight decay of 0.01. We reported the average performance of multiple runs.For the small-scale (H=128) models, we observed models with a wider FFN might overfit when fully fine-tuned with all parameters. Thus, we performed both full parameter fine-tuning and efficient fine-tuning based on adapters [12] for all models, choosing the better result of the two for reporting.For full parameter fine-tuning, the learning rates were {1e-5, 2e-5, 5e-5}; for efficient fine-tuning, the adapter size was 128, with learning rates of {1e-4, 2e-4, 3e-4}. For the large-scale (H=768) models, we directly conducted full parameter fine-tuning for all models, with learning rates of {1e-5, 2e-5, 5e-5}.In our study on GPT and MoE models, we conducted few-shot learning experiments on multiple datasets. Since small-scale (H=128) models have limited capabilities and struggle with few-shot learning, we focused our experiments only on large-scale (H=768) models for 0-shot and 1-shot learning and MoE models for 0-shot and 5-shot learning.We followed the method of [3] , which involves transforming the classification into comparisons of probability magnitudes.Specifically, the unified format of these datasets involves selecting one option from multiple choices, given a context. We concatenated the context with different options and compared their respective generation probabilities. This comparison could be based on either the probability of the option alone or the entire sequence text, with a choice of normalizing the probabilities by length or not. [3] mentioned a method of normalizing the probability of each option unconditionally, which we also included in our options. We then experimented with all possible combinations of these choices for each dataset, determined the best combination for each, and reported the results of all models under these optimal conditions. For 1-shot experiment, we randomly selected one demonstration each time, repeated the experiment 10 times, and reported the average results. For 5-shot experiment, we randomly selected five demonstrations each time (for MMLU, randomly shuffling the order of demonstrations), repeated the experiment 5 times, and reported the average results.The formal definition of mutual information for two discrete random variables is as follows:EQUATIONwhere X and Y are discrete random variables, P (X,Y ) is the joint probability mass function of X and Y , and P X and P Y are the marginal probability mass functions of X and Y respectively.For estimating the mutual information between intermediate representations and target tokens, we adopted the method proposed by [32] , which primarily involves transforming the representations into discrete variables through clustering and then calculating mutual information.We first randomly sampled 6.94 million input tokens and their corresponding output tokens from a pre-trained development set and collected all intermediate representations of the model at these locations. Then, for the GPT model, we performed clustering directly at each layer and calculated mutual information using the labels obtained from clustering with the output tokens. For the BERT model, we also conducted clustering at each layer; however, since masked language modeling only considers masked positions as the actual output, mutual information was calculated solely at these masked positions. Clustering was done using the mini-batch k-means algorithm, initialized with the k-means++ method, with a batch size of 1024 and a class number set to 2000.While it appears that BERT can reduce the actual contribution ratio of transformation function to zero, in reality, some contributions of transformation function remain uncovered by our method.For BERT, although we have eliminated most of the pathways through which the Inner-FFN directly leaks independent transformation information to subsequent layers, there is still one way this can happen: due to the bidirectional nature of attention, in the MHA of the current block, the entire sequence context has already been infiltrated by information about the current position at lower levels. Therefore, the Inner-FFN can independently enhance this part of the information, which can then leak out again through the weighted summation over values. In contrast, the GPT models employ unidirectional attention, so the context definitely does not contain any information about the current position, hence there is truly no independent transformation without the Outer-FFN.Therefore, while we are fairly certain that reducing the contribution ratio of independent transformation function is a general trend that can enhance base capabilities, completely eliminating the contribution of it could be harmful, as indicated by the trend in GPT; the absence of this phenomenon in BERT's trend is merely due to the likelihood that BERT still retains a small part of independent transformation in the Inner-FFN.In Section 6.2, we present the brief experimental results in Table 1 and 2 , and here we provide the corresponding detailed results.For BERT, out-of-distribution language modeling results on Pile are shown in Table 6 , fine-tuning results on GLUE Benchmark are shown in Table 7 , fine-tuning results on SuperGLUE Benchmark are shown in Table 8 , fine-tuning results on multiple other tasks are shown in Table 9 .For GPT, out-of-distribution language modeling results on Pile are shown in between original FFN-Wider models and new architecture models may still raise question: is the base capability improvement of new architecture models really due to architecture changes?To answer this question, we chose the Outer-FFN width ratios that allow CAA to align with FFN-Wider Transformer for both pre-training performance and pre-training steps, and conducted experiments similar to those in Section 6.2. Moreover, the parameter numbers and computational load of CAA are also aligned with the FFN-Wider Transformer. Therefore, we achieved alignment in four aspects, and under this setting, the improvements of base capabilities can be incontrovertibly attributed to our architectural changes.The Outer-FFN width ratio for FFN-Wider BERT w/ CAA is set to 12.5%, and for FFN-Wider GPT w/ CAA, it is set to 37.5%.The brief results are shown in Table 4 and 5, and the corresponding detailed results are as follows:For BERT, out-of-distribution language modeling results on Pile are shown in Table 12 , fine-tuning results on GLUE Benchmark are shown in Table 13 , fine-tuning results on SuperGLUE Benchmark are shown in Table 14 , fine-tuning results on multiple other tasks are shown in Table 15 .For GPT, out-of-distribution language modeling results on Pile are shown in Table 16 , zero-shot results and one-shot results on multiple datasets are shown in Table 17 .The experimental results show that the base capabilities of new architecture models have still significantly improved, confirming the positive impact of our architecture modifications on base capabilities. 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 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).This work mainly analyzes the models that use language modeling as the pre-training objective, lacking experiments on models with other pre-training objectives. Hence, the conclusions are limited to pre-trained language models. Therefore, the current applicability of our findings is relatively narrow, and we consider conducting more experiments in the future.https://pytorch.org/https://github.com/huggingface/transformershttps://github.com/laekov/fastmoe 


Decision: This paper explores how and why the FFN-Wider architecture influences base capabilities of LLMs (out-of-distribution modeling, transfer learning, and few-shot learning). Specifically, when the pretraining performance is held constant, FFN-Wider Transformers have worse base capabilities than vanilla ones, but why? The authors identify a property of the architecture (the contribution ratio of MHA layer) that correlates with base capability, and provide further analysis and insights. The authors then propose a specific architecture (with the best ratio) which is shown to improve base capabilities in MoE models.

I recommend to **accept** this paper. Reviewers appreciated that the authors chose to study the influence of architecture (instead of scale) on LLM performance, which is often difficult due to the low interpretability of neural networks in general, but nevertheless the authors provided "novel and insightful" analysis into the topic. It is also significant and impactful that the authors found a better (than vanilla) architecture for MoE models, based on that earlier analysis.

However, reviewers also mentioned that the writing/clarity of the paper could be improved. I am optimistic about this aspect because the authors clarified reviewers' questions in the rebuttals, and I encourage the authors to revise the paper to include the clarifications and generally improve the presentation for the next revision.

Reviews: review: paper_summary: paper_summary: The paper examines the influence of architecture on the base capabilities (OOD, transfer learning, few-short learning) of large language models. The main focus is on FFN-Wider transformers and understanding why they have poorer base capabilities compared to vanilla transformers. The contibution ratio of multihead attention to pretrained language modeling is found to be a key factor affective base capabilities. Based on this observation, a Combination Enhanced Architecture is proposed and also extended to mixture of experts transformers. The analytical study is backed by experimental evaluation successfully achieving significant improvements in base capabilities of the 14B parameter MoE model.
strength_weakness: strength_weakness: - Some of the findings appear to be obvious (perhaps, in hindsight) - as the contribution ratio of the MHA layer increases, the base capabilities also improve. It would have been interesting to also find some non-intuitive influences.  

- Overall the paper's analysis is much narrow in scope than the title of the paper. Nonetheless, it is interesting and useful study. 

- The language and presentation of the paper can be improved. It will be useful to have the grammar issues fixed in the paper before the final version.
questions: questions: - Some choices of parameters such as chosing intermediate dimension to 32d instead of say 8 or 16d can be better explained. 

- Was the width adjustment limited to Outer-FFN?
limitations: limitations: The reviewer does not except any negative societal impact of this work.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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: 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.
paper_summary: paper_summary: The paper studies the contribution raion of FFN and MHA layer in transformers and its effect on out-of-distribution performance. It finds that the wider FFN layer will increase its contribution ratio and lower the OOD performance. Lastly, the paper proposes a new architecture that moves part of the FFN layer to MHA and shows the CEA can improve the baseline models' OOD ability.
strength_weakness: strength_weakness: - The paper does not give a convincing explanation of why the study aligns pre-trained performance with different parameter scales. A larger scale model may not fully converge when it has a similar training loss to a smaller model. It is also not known whether the proposed CEA will harm in-domain performance.

- The MHA module in the vanilla transformer also has linear layers for Q, K, V, and after self-attention. I think they can be seen as inner-FFNs that will transform the combination function and make it not necessary to introduce an extra inner-FFN layer.

- Tables 1 and 2 show that BERT w/ CEA performs worse than vanilla in GLUE and SGLUE and GPT w/ CEA perform worse in all datasets.
questions: questions: See weaknesses 1 and 2.

Others:

- Figure 6 lacks a legend to explain the figure elements, including the blue line, orange bar, blue and orange dotted lines.
limitations: limitations: The authors addressed the limitations.
paper_summary: paper_summary: This study investigates the influence of architecture on pre-trained language models' base capabilities. It reveals that the contribution ratio of Multi-Head Attention to pre-trained language modeling affects base capabilities. FFN-Wider Transformers reduce this ratio, leading to a decline in base capabilities. The researchers propose a Combination Enhanced Architecture (CEA) to address this issue. They also extend this to Mixture of Experts (MoE) transformers, achieving significant improvements in base capabilities.
strength_weakness: strength_weakness: - The description of CEA is not detailed enough.
questions: questions: In general, I think the discovery of the paper is interesting, and I don't have many questions. Just some minor issues: 

- What is the definition of mutual information? Give some introduction.
- How is the pre-training performance of the models?
- How is the structure of CEA? It is not clear to me after reading Section 6.
limitations: limitations: The authors address it in Section 8.

Comments: title: Thank you
comment: Thank you for the rebuttal! It addresses my concerns and I will keep my rating.
Thank you very much for reviewing our work! 

We have addressed some of the concerns you raised in our rebuttal. Specifically, for the issues you were particularly concerned about, such as the fairness of the MoE experiment comparisons and the standard deviation of the experimental results, we have provided more detailed results (in the global response PDF above). We strongly believe that these additional results will effectively address your concerns, especially regarding the fairness of the MoE experiment comparisons, which seems to have led to some misunderstanding. 

If time permits, we kindly ask you to review our rebuttal. We look forward to your feedback. Thank you!
Thanks for the response, it addresses most of my concerns. 

I still have some concerns about the results of BERT/GPT w/ CEA. I understand that the FFN-Wider performs worse than the vanilla model, but I think w/ CEA should at least achieve similar results to the vanilla model. Otherwise, there is no reason to use CEA rather than vanilla. Maybe authors can show that the CEA can outperform vanilla with the same parameter scale and not align pre-training performance.

However, the results of MoE are convincing. I think this is a good paper and will raise my rating.
Thank you for addressing my concerns. Having carefully read all other reviews and your responses, I will keep my positive score.