Title: How does Architecture Influence the Base Capabilities of Pre-trained Language Models? A Case Study Based on FFN-Wider and MoE Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: Pre-trained Language Models

Base Capabilities

Mixture of Experts
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.

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

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

1 Introduction

Recent research has discovered that pre-trained language models possess strong base capabilities [24, 7, 3, 22]. They can not only address in-distribution language modeling which is usually their pretraining objective, but also unexpectedly excel in out-of-distribution language modeling , transfer learning , few-shot learning , etc. This has attracted the attention of many researchers.

However, it has also been observed that the cost of pre-training a language model is substantial, with the trial-and-error approach based on empirical improvements proving to be very expensive. Consequently, there is a desire to gain insights into the final performance by analyzing factors like scale and architecture that directly determine the base capabilities of models.

In this process, much attention has been focused on analyzing scale, leading to the formulation of compelling scaling laws [14, 11] that drive the trend of enhancing base capability by increasing parameter numbers, data volume and training tokens. In contrast, the impact of architecture has not

-2.75

-2.74

-2.73

-2.72

-2.71

-2.70

Outer-FFN Width Ratio

Contribution Ratio (Token Prediction)

(b) Outer-FFN Contribution Ratio &

Model Base Capability (OOD)

(c) MoE Models Performance (OOD)

Log-Likelihood on Pile Dev Set

(a) Combination Adjustable Architecture

MHA

LayerNorm

LayerNorm

Wider FFN

MHA

(w/ Direct Pathway)

LayerNorm

LayerNorm

Inner-FFN

Outer-FFN

FFN-Wider Architecture

Combination Adjustable Architecture

LayerNorm

BERT

0%

10%

20%

30%

40%

50%

60%

100.0%

87.5%

75.0%

62.5%

50.0%

37.5%

25.0%

12.5%

0.0%

Contribution Ratio of the Outer-FFN Layer

Out-of-distribution Language Modeling Performance

2.35

2.45

2.55

2.65

2.75

2.85

2.95

3.05

3.15

4

20

36

52

68

84

100

GPT-MoE-14B (64 experts, top-1, activated 1.5B)

GPT-MoE-14B w/ CEA  (64 experts, top-1, activated 1.5B)

Pre-Trained Tokens (B)

Loss on OOD Test Set

GPT

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

contribution ratio of the MHA layer (a combination function) increases, there is a general synchronous improvement in the model's base capabilities .

Subsequently, we identified the optimal width ratio of two parts of FFN for the Combination Adjustable Architecture (CAA) and established it as Combination Enhanced Architecture (CEA), and comparing this architecture with the original FFN-Wider Transformer. We conducted experiments on various scales of BERT and GPT models, with all results robustly supporting our explanation.

Importantly, we also noticed that existing work has observed a similar decline in base capabilities in the Mixture of Experts (MoE) Transformers. We applied our explanation and CEA to MoE Transformers as well . We pre-trained a 14B parameter GPT architecture MoE model and its improved version with CEA on 100B tokens (using CC and C4 from SlimPajama [29]). With the same number of parameters, computational load, and pre-training steps for both versions, the improved version with CEA showed significant improvements in out-of-distribution language modeling and fewshot learning. The performance of the out-of-distribution language modeling (average performance across all domains in SlimPajama except for CC and C4) is shown in Figure 1(c), which can prove the practical application value of our work.

Overall, the actual contribution ratio of the MHA layer (a combination function) is likely a universal factor affecting model's base capabilities, which can provide valuable insights for architecture analysis, improvement and design.

2 Background

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

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

2.2 Evaluation Schemes

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

-3.21

-3.26

-3.40

-3.34

-3.28

-3.22

78.45

76.66

73

75

77

79

-3.88

-3.91

-3.97

-3.94

-3.91

-3.88

45.82

44.60

41

43

45

47

46.09

44.84

41

43

45

47

Pile

(OOD)

GLUE

(FT)

Pile

(OOD)

Multiple

Tasks

(0-shot)

Multiple

Tasks

(1-shot)

-2.16

-2.20

-2.30

-2.26

-2.22

-2.18

83.89

82.60

79

81

83

85

64.41

62.82

59

61

63

65

52.14

51.61

50

51

51

52

-4.73

-4.76

-4.73

GLUE

(FT)

Pile

(OOD)

Pile

(OOD)

SuperGLUE

(FT)

Other Tasks

(FT)

FFN-Wider Model

BERT (H=128)

GPT (H=768)

BERT (H=768)

GPT (H=128)

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.

2.3 Evaluation Tasks

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.

3 FFN-Wider Transformers vs. Vanilla Transformers

This work focuses on the base capabilities of FFN-Wider Transformers. A typical Transformer model has Feed-Forward Network (FFN) layers. Assuming the hidden dimension is d , the standard intermediate dimension of an FFN layer is 4 d . However, an FFN-Wider Transformer model means the intermediate dimension of its FFN layer exceeds 4 d .

We conducted experiments on various models with the intermediate dimension set to 32 d , aligning pre-training performances with those of vanilla models. The results are shown in Figure 2.

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?

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

4.2 Contribution Ratio Analysis

To verify our hypothesis, we first need to confirm whether the actual contribution ratio of MHA layers truly decreases, requiring quantitative analysis.

4.2.1 Mutual Information

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

4.2.2 Token Prediction

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

-0.400

0.000

0.400

0.800

1.200

1.600

2.000

1

2

3

4

5

6

7

8

9

10

11

12

-0.400

0.000

0.400

0.800

1.200

1.600

2.000

1

2

3

4

5

6

7

8

9

10

11

12

-0.100

0.000

0.100

0.200

0.300

0.400

0.500

0.600

0.700

1

2

3

4

5

6

7

8

9

10

11

12

-0.100

0.000

0.100

0.200

0.300

0.400

0.500

0.600

0.700

1

2

3

4

5

6

7

8

9

10

11

12

Total Increment

Caused by FFN

Total Increment

Caused by FFN

Total Increment

Caused by FFN

Layers

Layers

Layers

Layers

Cumulative MI Increment

Cumulative MI Increment

Cumulative MI Increment

Cumulative MI Increment

(b) FFN-Wider BERT (H=128)

(d) FFN-Wider GPT (H=128)

Total Increment

Caused by FFN

BERT

GPT

BERT

GPT

Cumulative MI Increment of Block

Cumulative MI Increment of MHA

Cumulative MI Increment of FFN

35.19%

28.26%

22.33%

39.46%

53.74%

43.44%

43.00%

60.80%

0%

10%

20%

30%

40%

50%

60%

70%

80%

BERT (H=128)

BERT (H=768)

GPT (H=128)

GPT (H=768)

FFN-Wider Model

FFN Contribution Ratio (Token Prediction)

FFN / Total Accuracy Increment (%)

can be seen that the contribution of FFN layers in FFN-Wider models remains higher than that in vanilla models, reaffirming our hypothesis.

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

5.1 Model Architecture

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 MHAthat bypasses the Inner-FFN, which we will introduce in detail below.

5.1.1 Inner-FFN and Outer-FFN

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

5.1.2 Direct Pathway in MHA

Although we considered the impact of residual connections and designed an Inner-FFN, the InnerFFN 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

Wider FFN

Combination Adjustable

Architecture (CAA)

Inner-FNN and Outer-FFN (Width Ratio is Adjustable)

Direct Pathway in MHA (Single Head)

Outer-FFN

Inner-FFN

MHA

(w/ Direct Pathway)

LayerNorm

LayerNorm

Inner-FFN

Outer-FFN

LayerNorm

keys

queries

Query

Output

Key

Score

Score

Value

values

scores

scores

Derived from MHA Block Input / Strictly Unrelated to Inner-FFN

0%

10%

20%

30%

40%

50%

60%

0%

10%

20%

30%

40%

50%

60%

70%

Outer-FFN Width Ratio

(a) Outer-FFN Contribution Ratio

Outer-FFN Width Ratio

Contribution Ratio (TP)

Contribution Ratio (TP)

(c) Outer-FFN Contribution Ratio

BERT

GPT

20

30

40

50

60

70

80

90

100

-2.77

-2.75

-2.73

-2.71

-2.69

Vanilla BERT

Vanilla BERT

Outer-FFN Width Ratio

Log-Likelihood (Pile

）

Pretraining Steps (k)

(b) Model Base Capability (OOD)

BERT

10

30

50

70

90

110

130

150

-4.14

-4.13

-4.12

-4.11

-4.10

-4.09

Outer-FFN Width Ratio

Log-Likelihood (Pile)

Pretraining Steps (k)

(d) Model Base Capability (OOD)

GPT

current position, the query, key and value from the current position use the input representation that has not been processed by the Inner-FFN. Only the context representations from non-current positions are transformed by the Inner-FFN. This design largely eliminates the possibility of the Inner-FFN bypassing the combination function.

5.2 Width Adjustment Analysis

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.

5.2.1 Trend of Contribution Ratio

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.

5.2.2 Trend of Base Capability

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.

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

6.1 Width Ratio Selection

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

6.2 Further Experiments

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.

7 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]. Then, we added a MoE layer with 64 experts (top-1 activation) before the FFN layer, resulting in a 14B parameter MoE baseline model. Finally, our improved version with CEA involves transforming the MoE layer into an Inner-MoE layer within the MHA,

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.

8 Limitations

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.

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.

Acknowledgments

This work was supported by the Fundamental Research Funds for the Central Universities (project number: 2022FRFK060002), the National Key RD Program of China via grant 2021YFF0901602 and the National Natural Science Foundation of China (NSFC) via grant 62176078.

References

M. Artetxe, S. Bhosale, N. Goyal, T. Mihaylov, M. Ott, S. Shleifer, X. V. Lin, J. Du, S. Iyer, R. Pasunuru, G. Anantharaman, X. Li, S. Chen, H. Akin, M. Baines, L. Martin, X. Zhou, P. S. Koura, B. O'Horo, J. Wang, L. Zettlemoyer, M. Diab, Z. Kozareva, and V. Stoyanov. Efficient large scale language modeling with mixtures of experts. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 11699-11732, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics.

Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence , volume 34, pages 7432-7439, 2020.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1877-1901. Curran Associates, Inc., 2020.

C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 2924-2936, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv , abs/1803.05457, 2018.

T. Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. 2023.

J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

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

L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020.

D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding, 2021.

J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, O. Vinyals, J. Rae, and L. Sifre. An empirical analysis of compute-optimal large language model training. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 30016-30030. Curran Associates, Inc., 2022.

N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, Q. De Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly. Parameter-efficient transfer learning for NLP. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 2790-2799. PMLR, 09-15 Jun 2019.

A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. de las Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, L. R. Lavaud, M.-A. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed. Mistral 7b, 2023.

J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models, 2020.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In Y. Bengio and Y. LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015.

G. Lai, Q. Xie, H. Liu, Y. Yang, and E. Hovy. RACE: Large-scale ReAding comprehension dataset from examinations. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing , pages 785-794, Copenhagen, Denmark, Sept. 2017. Association for Computational Linguistics.

Z. Lan, M. Chen, S. Goodman, K. Gimpel, P. Sharma, and R. Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations , 2020.

D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. {GS}hard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations , 2021.

H. Levesque, E. Davis, and L. Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning , 2012.

T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing , pages 2381-2391, Brussels, Belgium, Oct.-Nov. 2018. Association for Computational Linguistics.

N. Mostafazadeh, N. Chambers, X. He, D. Parikh, D. Batra, L. Vanderwende, P. Kohli, and J. Allen. A corpus and evaluation framework for deeper understanding of commonsense stories. arXiv preprint arXiv:1604.01696 , 2016.

OpenAI. Gpt-4 technical report, 2023.

D. Paperno, G. Kruszewski, A. Lazaridou, N. Q. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 1525-1534, Berlin, Germany, Aug. 2016. Association for Computational Linguistics.

A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. 2018.

A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

M. Roemmele, C. A. Bejan, and A. S. Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series , 2011.

K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM , 64(9):99-106, 2021.

M. Sap, H. Rashkin, D. Chen, R. Le Bras, and Y. Choi. Social IQa: Commonsense reasoning about social interactions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pages 4463-4473, Hong Kong, China, Nov. 2019. Association for Computational Linguistics.

Z. Shen, T. Tao, L. Ma, W. Neiswanger, Z. Liu, H. Wang, B. Tan, J. Hestness, N. Vassilieva, D. Soboleva, and E. Xing. Slimpajama-dc: Understanding data combinations for llm training, 2024.

J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.

Y. Tay, M. Dehghani, S. Abnar, H. Chung, W. Fedus, J. Rao, S. Narang, V. Tran, D. Yogatama, and D. Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling? In H. Bouamor, J. Pino, and K. Bali, editors, Findings of the Association for Computational Linguistics: EMNLP 2023 , pages 12342-12364, Singapore, Dec. 2023. Association for Computational Linguistics.

E. Voita, R. Sennrich, and I. Titov. The bottom-up evolution of representations in the transformer: A study with machine translation and language modeling objectives. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pages 4396-4406, Hong Kong, China, Nov. 2019. Association for Computational Linguistics.

A. Wang, Y. Pruksachatkun, N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , pages 353-355, Brussels, Belgium, Nov. 2018. Association for Computational Linguistics.

J. Welbl, N. F. Liu, and M. Gardner. Crowdsourcing multiple choice science questions, 2017.

R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 4791-4800, Florence, Italy, July 2019. Association for Computational Linguistics.

B. Zhang and R. Sennrich. Root Mean Square Layer Normalization. In Advances in Neural Information Processing Systems 32 , Vancouver, Canada, 2019.

A Other Evaluation Schemes

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.

B Evaluation Tasks

B.1 FFN-Wider Transformers

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.

B.2 MoE Transformers

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.

C Model Specifications and Pre-training Procedures

C.1 FFN-Wider Transformers

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 PyTorch 2 and transformers 3 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.

C.2 MoE Transformers

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]. Then, we added a MoE layer before the FFN layer, resulting in our MoE baseline model (Vanilla MoE 14B). Finally, our improved version with CEA (MoE 14B w/ CEA) involves transforming the MoE layer into an Inner-MoE layer within the MHA, 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 Vanilla GPT 1.3B has a hidden dimension of 4,096, 6 layers and 32 attention heads; the Vanilla MoE 14B has a hidden dimension of 4,096, 6 layers, 32 attention heads and 6 MoE layers with 64 experts (top-1 activation, intermediate dimension is 4,096) before the corresponding FFN layers; the MoE 14B w/ CEA is similar to the Vanilla MoE 14B, just the position of the parameters has changed. 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 fastmoe 4 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.

D Fine-tuning Settings

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

E Few-shot Learning Settings

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.

F Mutual Information Estimate

The formal definition of mutual information for two discrete random variables is as follows:

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

G Additional Transformation Function Contribution in BERT

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.

H Detailed Results (Main Setting: Pre-training Performance Alignment)

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 Table 10, zero-shot results and one-shot results on multiple datasets are shown in Table 11.

I Detailed Results (Extra Setting: Pre-training Performance Alignment & Pre-training Steps Alignment)

The main experiments in this work are based on the pre-training performance alignment scheme. Although we have introduced the rationale behind this scheme, the difference in pre-training steps

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

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: In Abstract and Section 1.

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: In Section 8.

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: Not involving proof.

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 Appendix C, D, E and F.

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: [No]

Justification: There are some restrictions that make it inconvenient to publish directly.

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: In Appendix C, D, E and F.

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: We use multiple rounds of randomization to ensure the stability of the results.

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]

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

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: This work is a foundational research.

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

13. 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 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: N/A.

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: N/A.

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

STRENGTHS:
+ The paper shifts the focus from the commonly studied impact of scale on pre-trained language models to the impact of architecture.
+ The exploration of the influence of using wider FFN layer is interesting, and should impact further researches.
+ The proposal of CEA as a solution to the decline in base capabilities shows a proactive approach to solving the identified problem.
+ The findings are confirmed through extensive experiments, adding credibility to the claims.

WEAKNESSES:
- The description of CEA is not detailed enough.

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:
The authors address it in Section 8.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-08

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

STRENGTHS:
- The paper analyzes the contribution ratio with new methods that evaluate Mutual Information and token prediction accuracy of each layer's output.

- The proposed CEA method has improved the MoE transformer in various tasks and datasets with the same parameter scale, showing its effectiveness in increasing OOD ability.

WEAKNESSES:
- 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:
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:
The authors addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nInitially, 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.\n\nThe 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.\n\nSTRENGTHS:\n- 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.\n\n- 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.\n\n- 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. \n\n- The application of CEA to a mixture of experts models definitely demonstrates the practical benefits of this paper.\n\nWEAKNESSES:\n- 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. \n\n- 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.\n\n- 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. \n\nOverall, 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. \n\nMinor: In Table 1 and Table 2, the results for the proposed approach are bolded rather than the best approach for that experiment.\n\nIn most figures, the y-axis does not start at 0, which can give readers an inaccurate representation of the data.\n\nQUESTIONS:\n- What does it mean for the width ratio to be 0% and 100%? Does this imply that the FFN is removed in these cases?\n\n- 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?\n\n- Please indicate whether lower or higher values are better for all the tables and figures in the paper.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n+ While most existing study of large models has focussed on the impact of scale, this work is a significant effort in understanding the influence of architecture. \n\n+ The empirical findings are supported by a more deeper interpretation of the underlying mechanisms of these influences. \n\n+ The proposed CAA wherein the wider FFN is split into adjustable two parts - outer-FFN which stays as a transformation function, and an inner-FFN which is relocated within the MHA layer for enhancing the combination function is interesting in its own right. \n\n+ The extension to MoE transformers and the experimental demonstration with 14B parameter GPT architecture MoE model is very interesting.\n\nWEAKNESSES:\n- 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.  \n\n- Overall the paper's analysis is much narrow in scope than the title of the paper. Nonetheless, it is interesting and useful study. \n\n- 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.\n\nQUESTIONS:\n- Some choices of parameters such as chosing intermediate dimension to 32d instead of say 8 or 16d can be better explained. \n\n- Was the width adjustment limited to Outer-FFN?\n\nLIMITATIONS:\nThe reviewer does not except any negative societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}