Unraveling Feature Extraction Mechanisms in Neural Networks
Authors: Xiaobing Sun
Jiaxi Li
Wei Lu
Venue: EMNLP 2023 Main (2023)
Keywords: Interpretability
Infinite-width
Feature extraction
Learning dynamics
Neural tangent kernel
TL;DR: We propose a theoretical approach to delve into the feature extraction mechanisms behind neural models.

Abstract: The underlying mechanism of neural networks in capturing precise knowledge has been the subject of consistent research efforts. In this work, we propose a theoretical approach based on Neural Tangent Kernels (NTKs) to investigate such mechanisms. Specifically, considering the infinite network width, we hypothesize the learning dynamics of target models may intuitively unravel the features they acquire from training data, deepening our insights into their internal mechanisms. We apply our approach to several fundamental models and reveal how these models leverage statistical features during gradient descent and how they are integrated into final decisions. We also discovered that the choice of activation function can affect feature extraction. For instance, the use of the ReLU activation function could potentially introduce a bias in features, providing a plausible explanation for its replacement with alternative functions in recent pre-trained language models. Additionally, we find that while self-attention and CNN models may exhibit limitations in learning n-grams, multiplication-based models seem to excel in this area. We verify these theoretical findings through experiments and find that they can be applied to analyze language modeling tasks, which can be regarded as a special variant of classification. Our work may offer insights into the roles and capacities of fundamental modules within deep neural networks including large language models.

Unraveling Feature Extraction Mechanisms in Neural Networks
Xiaobing Sun Jiaxi Li, Wei Lu StatNLP Research Group Singapore University of Technology and Design (xiaobing_sun, jiaxi_ 1i_ Jemymaiimmmllll sutd. edu. sg, wei 1i__uudsudd.. edu.
Label
Abstract
(oottec -)
The underlying mechanism of neural networks in capturing precise knowledge has been the subject consistent research efforts. this work, we propose a theoretical approach based on Neural Tangent Kernels (NTKs) investi- gate such mechanisms. Specifically, consider- ing the infinite network width, hypothesize the learning dynamics of target models may in- tuitively unravel the features they acquire from training data, deepening our insights into their internal mechanisms. We apply approach to several fundamental models and reveal how these models leverage statistical features during gradient descent and how they are integrated into final decisions. We also discovered that the choice of activation function can affect fea- ture extraction. For instance, the use of the ReL. activation function could potentially in- troduce bias features, providing a plausible explanation for its replacement with alternative functions in recent pre-trained language models. Additionally, we find that while self-attention and CNN models may exhibit limitations in learning n n-grams, multiplication- based models seem excel this area. We verify these the- oretical findings through experiments and find that they can be applied analyze language modeling tasks, which can be regarded as spe- cial variant of classification. Our contributions offer insights into the roles and capacities of fundamental components within large language models, thereby aiding the broader understand- ing of these complex systems.
Inot,
Jalayps
(for
always
for
better
Figure 1:
Example of CO- ro-occurrence features between
tokens and labels (self-attention model).
may result in interesting findings, but those models still remain black-boxes to us. Another line seeks reveal internal mechanisms of neural models ing mathematical tools (Levy and Goldberg, 2014; Saxe et al., 2013; Arora al., 2018; Bhojanapalli et al., 2020; Merrill et al., 2020; Dong et al., 2021; Tian et al., 2023), which can be more straightfor- ward and insightful. However, few of them have specifically focused on the feature extraction of neural NLP models.
When applying neural models to downstream NLP tasks in practice, we often notice some mod- ules perform better than others specific tasks, while some exhibit similar behaviors. We may won- der what mechanisms are behind such differences and similarities between those modules. By acquir- ing deeper insights into the roles of those modules in a complex model with respect to feature extrac- tion, we will able to select or even design more suitable models for downstream tasks.
In this work. we propose novel theoretical ap- proach to understanding the mechanisms, through which fundamental models (often used as mod- ules in complex models) acquire features during gradient descent in text classification tasks. The evolution of model output can be described as Iearn- ing dynamics involving NTKs (Jacot al., 2018; Arora et al., 2019), which are typically used study various properties of neural networks, includ- ing convergence and generalization. While these representations can be complex in practice, when the width of the network approaches infinity, they tend to converge to less complex representations and remain asymptotically constant (Jacot et al., 201 8), allowing us to intuitively interpret the learn- ing dynamics and identify the relevant features cap-
1 Introduction
Neural networks have become indispensable across a variety of natural language processing (NLP) tasks. There has been growing interest in under- standing their successes and interpreting their char- acteristies. One line of works attempts iden- tify possible features captured by them for NLP tasks (Li et al., 2016; Linzen et al., 2016: Jacovi et al., 2018 8; Hewitt and Manning, 2019; Vulić et al., 2020). They mainly develop empirical methods to verify hypotheses regarding the semantic and syn- tactic features encoded in the output. Such works

---

tured by the model.
useful features during training from a theoretical perspective.
We applied our approach to several fundamental models, including a multi-layer perceptron (MLP), a convolutional neural network (CNN), linear Re- current Neural Network (L-RNN), self-attention (SA) model (Vaswani et al., 2017), and matrix- vector (MV) model (Mitchell and Lapata. 2009) and exhibit the MLP, CNN, and SA models may behave similarly in capturing token-label features, while the MV and L-RNN extract different types of features. Our contributions include:
Infinite-width Neural Networks Researchers found that there could be interesting patterns when the neural network's 's width approaches infinity, Lee et al. (2018) linked infinitely wide deep networks to Gaussian Processes. A recent line of work (Jacot et al., 2018; 8; Bietti and Mairal, 2019; Nguyen et al., 2021; Loo et al., 2022) proposed that as the net- work width approaches infinity, the dynamics can be characterized by the NTK, which converges to a kernel determined at initialization and remains con- stant. This conclusion holds for fully- fully-connected neural networks, CNNs (Arora et al., 2019) and RNNs (Emami et al., 2021; Alemohammad et al., 2021). Later, Yang and Littwin (2021) showed that such properties of NTKs can be applied to a randomly initialized neural network of any archi- tecture. Very limited studies have delved into the analysis of feature extraction in neural NLP mod- els. We will investigate the internal mechanisms of neural NLP models under extreme conditions.
We propose an approach to theoretically inves- tigate feature extraction mechanisms for funda- mental neural models.
We identify significant factors such as the choice of activation and unveil the limitations of these models, e.g.. both the CNN and SA models may not effectively capture meaningful n-gram information beyond individual tokens.
Our experiments validate the theoretical find- ings and reveal their relevance to advanced architectures such as Transformers (Vaswani et al., 2017).
3 Analysis
Our intention through this work is to provide new insights into the core components of complex models. By doing so, we aim to contribute to the understanding of the behaviors exhibited by state- of-the-art large language models and facilitate the development of enhanced model designs'.
We use learning dynamics to describe the updates of neural models during training with the aim of identifying potentially useful properties. For the of presentation and discussion, focus on binary text c
Model Deseription Assume we have a training dataset denoted by D. consisting of m labeled instances. Let X' and y represent all the sen- tences and labels in the training dataset, respec- tively.
2 Related Work
Probing features for NLP models Probing lin- guistic features is an important topic for verifying the interpretability of neural NLP models. Li et al. (2016) employed visualization approach to de- tect linguistic features such negation captured by the hidden states of L.STMs. Linzen et al. (2016) examined the ability of LSTMs to capture syntac- tic knowledge using number agreement in English verj dependencies. Jacovi et al. (2018) studied whether the CNN models could capture n n-gram features. Vulié et al. (2020) presented sys- tematic analysis probe possible knowledge that the pre-trained language models could implicitly capture. Chen et al. (2020) proposed an algorithm to detect hierarchical feature interaction for text classifiers. Empirically, such work reveals that neural NLP models can capture useful and inter- pretable features for downstream tasks. Our work seeks to explain how neural NLP models capture
X' is an instance consisting of se- quence of tokens, and ) is the corresponding label. The vocabulary size is [V]. Consider a binary text classification model, where y E (-1, +1}. The model output, denoted as s(t) E IR at time t is
s(t) - frr; 0:),
(1)
where O, (a vector) is the concatenation of all the parameters, which are functions of time l. We refer to the model output st)) as the label score at time L. This score is used for classification decisions, positive if stt) > 0 and negative otherwise.
Learning Dynamics The evolution of label score can be described by learning dynamics, which may indicate interesting properties. Let f((2)) E IR" represent the concatenation of all the outputs of training instances at time t, and E y
'Our code is available at Shharrsun-vvyager//ueee
AAalyysis for multi-class classification can be found in Appendix A.

---

is the desired label. Given test input x', the corre- sponding label score s1()) follows the dynamics
deterministic trend can be predicted for test input x' that closely resembles the training instances of specific type. For example, suppose O (a', x)) exhibits significantly large gain when x' is similar to x at particular y. and the dynamics will likely receive significant gains in desired direction dur- ing training, thus enabling us to predict the trend of the label score.
(1) = Vesi ( sf = -,(x'' s(aYe
(2)
where (x(2 is the NTK at time t and L is the empirical loss defined as
We thereby propose the following approach to investigate a target model and verify our afore- mentioned hypotheses: 1) redefining the target model following the settings proposed by Jacot et al. (2018); Yang and Littwin (2021), which guar- antees the convergence of NTKs; 2) obtaining the converged NTK (a', x)) and the learning dy- namies under the infinite-width condition; 3) per- forming analysis on the learning dynamics of basic units and revealing possible features.
1
L
--- m
Σ lug)(
(3)
(E
where g is the sigmoid function. For simplicity, we will omit the time stamp in our subsequent notations. The dynamics s' will obey
s' Eg(-ys (z) )y0(e',,) m (Z,))ED
(4)
where g(x) is the label score for the training in- stance z. Obtaining closed-form solutions for the differential equation in Equation 4 is challenge. We thereby consider an extreme scenario with the infinite network width, suggested by Lee et al. (2018).
Interpreting Fundamental Models
We investigate an MLP model, a CNN model, an SA model, an MV model, and an L-RNN model, respectively. Details and proofs for the lemmas and theorems can be found in Appendix A.
Notation Let e E RIM be the one-hot vector for token e, [(x) be the instance length, W e Rdox(V] be the weight of the embedding layer, and E Rdout be the final layer weight. W E Rdoot Xden is the weight of the hidden layer in the MLP model. We E Reou is the kernel weight corresponding to the k-th token in the slid- ing window in the CNN model. For simplicity, we let dout = din = d. Assume all the parame- ters are initialized with Gaussian distributions in our subsequent analysis, i.c., Wij ~ (O o We 1j ~ (O, of), and vy ~ (N, o and WGJ ~ (0,, o for the sake of NTK convergence.
Infinite- Width When the network width ap- proaches infinity, the NTK will converge and stay constant during training (Jacot et al., 2018; Arora et al., 2019; Yang and Littwin, 2021). Therefore, the learning dynamics can be written as follows,
8' s' 1 Eg(-ys(a) D0 m. (D.8))C
(5)
where 2x ,x) refers the converged NTK de- termined initialization. This convergence may allow to simplify the representations of the learn- ing dynamics and offer more intuitive insights to analyze its evolution over time.
4.1 MLP
There can be certain interesting properties (re- garding the trend of the label scores) harnessed by the interaction y (x', x). where controls the direction and O (x', (x', 2)) may indicate the relation- ship between x' and r. Certain hypotheses can be drawn from these properties. First, the converged NTK Oz (a', 2)) may intuitively represent the in- teraction between the test input x' and the training instance This could extend to the interaction between the basic units (tokens n-grams) from x' and v, as the semantic meaning of an instance can be deconstructed into the combination of the meanings of its basic units (Mitchell and Lapata, 2008; Socher et al., 2012), 2). Second, if O (a', depends on the similarity between x' and x, a more
Following Wiegreffe and Pinter (2019), given in- stance T, the output of MLP is defined as
(x) W W)e)).
8
(6)
Vd
The label score will be used for making clas- sification decisions. & is the element-wiss ReLU function. ej is the one-hot vector for token ej. It is not straightforward analyze directly, which can be viewed as the sum of token-level label scores. Instead, as basic units are tokens in this model, we focus on the label score of every single token and understand how they contribute to the instance- level label score. When the test input x' is simply

---

token e, we can get the corresponding NTK with the infinite network width.
them in final classification decisions, i.e., tokens solely appearing in positive/negative instances will likely contribute in the direction of predicting positive/negative label.
Lemma 4.1. When + oc. the NTK between the
token and instance in the MLP model converges to
((0]
p(a)
4.2 CNN
2) = e e
D
44
(7)
We consider the I-dimensional CNN, with kernel size, stride size, and padding size set to K, 1, and K respectively. For each sliding window ej comprising K consecutive tokens, the correspond- ing feature Cj
[m-1]0202 j=1
2eo 0203300200 and μ
where p
2x
2n
Note that, for two tokens ej and ek. their one- hot vectors satisfy e. es 0 if e # Cki ej ex =
IRd can be represented
1 if ej = cg-- The dot-product E e ej can be interpreted as the frequency of é appearing in instance x.
K
C, c-
3 W k=1
W ej+k-1,
(9)
Theorem 4.2. The learning dynamics of token e's label score obey
where W is the kernel weight corresponding to the k-th token in the sliding window.
The label score of an instance is computed as
g
P
Σ
ge--yss - )yw((,
(e)
m (T.))ED
Σ Φ(c)); (K-I)
(10)
(8)
μ
Σ
a (ys))))))) )y)(x)
+ m
where -(K 1) means the position for the left- most padding token. The first and last K - 1 padding tokens in instance are represented by zero vectors. & is the element- ReLU function.
(DD
((=)
where w(e
2j=1 ej. which depends on the training data and will not change over time.
The non-linearity of the sigmoid function g( uyy makes it challenge to obtain closed- form solution for the dynamics.
For brevity, we will denote Σ (K 1) by Σj. Let us focus on a single sliding window and study the learning dynamics of its label score.
However, we can predict trends for the label scores in special cases. Note that the polarity of the first term in Equation 8 will depend on yw(e, (,, x) in each training instance. For instance, consider a
Lemma 4.3. Consider a sliding window e consist- ing of tokens 61, 62, ........,, CK, when d - oo the NTK between e and instance converges
token that only appears in positive instances, i.e., (ee 2) > when y = +1; w(e, x) = 0 when this case, the first term remains posi- tive and incrementally contributes to the label score sf throughout the training process. The opposite trend occurs for tokens solely appearing in negative instances. If the impact of the second term is mini- mal, the label scores of these two types of tokens will be significantly positive or negative after suf- ficient updates. The final classification decisions are made based on the linear combination of the label scores for the constituent tokens. The second term in Equation 8 is unaffected by w(e, x) and is shared by all the tokens e at each update. It can be interpreted as an induced feature bias. Particularly, when this term is sufficiently large, it may cause an imbalance between the tokens co-occurring with the positive label and those co-occurring with the negative label, rendering one type of tokens more influential than the other for classification.
(x,(c, = Σ el)
y
K
H HweleccjleTTe ez-k-11 (11)
P
k=1
where
K
K
Welc, (e)(c,)))= Σ ek #=1 1=1
ejtk-1, -=0
We means the number of shared tokens between and Cj regardless of positions. F and H3 are connotonicalllyinceessiii and non-negative func- tions depending on oa
The first term in (x, x) captures the token similarity between sliding windows e and ej re- gardless of token positions. In the second term, Σj H [[e(c, cj)le, Cj+k-1 can be viewed as the weighted frequency of token eş in instance x, and when og is sufficiently large, the converged NTK is majorly influenced by the sum of the weighted frequencies of the tokens in appearing in x.
Theorem 4.2 may explain how the MLP model leverages the statistical co-occurrence features be- tween e and y as shown in Figure 1, and integrate
Thhei definitions can be found in Appendix A.2.

---

a² a Theorem 4.4. The dynamics of the label score of the test sliding window obey
where P: (P3) is the positional embedding at po- sition ()) and will be fixed during training. The instance label score will be computed
K
Σ
9( ys (Ew(
((m) ((r)
((m)
b= (xxD)ED
-
h. h
T W (16)
Σ (L,))D
g( yolo)) y Σ FFwelc
m
which can be viewed as the weighted sum of token- level label scores if we define such a score for each token as 8. v22 W We consider the case where the test input is also simply token e. Lemma 4.5. When d oo, the NTK between the token and the instance will converge 2 (e, 2), which obeys
(12)
where w(ek, .) Huclecllee
Theorem 4.4 indicates that with sufficiently
large ou, the learning dynamics for window may mainly depend on the linear combination of the weighted learning dynamics of its constituent to- kens. Similar analysis can be performed on the label score of the sliding window. This may not exactly encode n- n-grams, which are inherently sen- sitive to order and can extend beyond their con- stituent elements. Instead, for each window, more akin the composition model based on vec- for addition as described in the work of Mitchell and Lapata (2009). The second term in Equation 12 may not be zero even if shares tokens with suggesting there be an induced feature bias similar to the one in the MLP model.
(l0) p(r) E B(ozj)e ej,
(c) to
>
i11=1
(17)
where B(axt) the expectation of aij.
Theorem 4.6. The learning dynamics of the label score of a token e obey
P
Σ
g( ys (x) (yw
(18)
m
(a.))CD
When e only shares tokens with either positive or negative instances, regardless of position, the corresponding label score will receive relatively large gains in one direction during updates. This means the CNN model also captures co-oceurrence features between tokens and labels. Importantly, single token can also be viewed as a sliding window, padded with additional tokens, thereby leading to conclusions about the trend of label scores that mirror those drawn from the MLP model.
where w(e, x)) = E i= E B(ai))e e, and p
Theorem 4.6 shows the learning dynamics of to- ken e's label score also depends on the weighted sum of the frequencies of appearing in æ. The learning dynamics of a single token's label score will likely resemble that in the MLP model, cap- turing the co-occurrence features between tokens and labels despite the weights. This model may not experience an induced bias, compared to the MLP model as discussed in Theorem 4.2. This will be further explored in our experiments.
4.3 SA
We employ a fundamental self-attention module, analogous to the component found in Transform- ers. The representation of the i-th output in the instance will be computed as a weighted sum of token representations as follows,
4.4 MV
We consider the matrix-vector representation as applied in adjective-noun composition (Baroni and Zamparelli, 2010) and recursive neural networks (Socher al., 2012). It models each word pair through matrix-vector multiplication. The label score of an instance is defined as 1
((x)
h;
W'
(13)
where Ouj is the weight produced by a sofimax function follows,
s= Σ dvd
M(e;) eit1
(19)
exp(ai))
where M(e)) diag Moag(War (WWe)) (diag converts vector into a diagonal matrix.) and 1,2, (x) 1.
aij
(14)
[() E exp(ait)
We define the attention score ai from position i
to j as
Lemma 4.7. Given a bigram consisting of two
aij- (W + P.) (W + P2)
tokens C. with the infinite network width the NTK will converge to
(15)

---

Model
Feature
Bias
Activation
Definition
NTK
MI.P
token-lahel
Yes
ReI.U
va 20 O( w)s))
Bu((a)) = a E E'a, Σ5-i
CNN
token-lahel
Yes
ReE.
3a 2x--1-1;000 ((g)
Bulna) Eg Fiullie)) + Modeellee
SA loken-label
No
Eim Wey
Wee
Oule, Ea
MV
bigram-lahel
No
Σ, za (ej
L.RNN
loken label position
123311 -WWe; Bu) apn-
Table
Co-occurrence features captured by target models. "k" in L-RNN refers the distance from the last tokens. p (f) refers to the non-negative coefficient determined at initialization for each model. and x refer to a token and instance, respectively. "-" means not applicable.
20 (CaCb: x) (o² 30 o E e ea eeeeene
Data
Train
Valid
Test
Len
SST
6,920
872
1,821
16,174
18
(20)
SSTwsub
98,794
872
1,821
17,404
It is worth highlighting that the interaction ej e is different from the interaction re- sulting from the aforementioned models. When Cj and 6 = Cj+1 (i.e., Co = cyoj+1), the NTK will gain a relatively large value, implying the ability to capture co-occurrence knowledge be- tween bigrams and labels.
IMDB
39,877
5,016
5,107
146,582
270
Agnews 110,000
10,000
7,600
85,568
36
Table 2: Dataset statistics. "Train", "Valid", and "Test" refer the training, validation, and test sets, respec- tively. "IVI" refers to the vocabulary size and "I.en" refers the average training instance length.
which consists of movie reviews with relatively longer texts. The statistics are listed in Table 2. In addition, Penn Tree Bank (PTB) (Marcus et al., 1993), WIKITEXT2 (Wiki2), and Shake- speare dataset are considered for language model- ing, a special classification variant..
Theorem 4.8. The dynamics of the label score of the test bigram CaCb obey
gab
P m
Σ
g( g(-ys(*)) (yw(ee)))
(21)
(T.
where P= = (o? 30 20 and a(cacb,
Ej ej eeee11e.
Setup We randomly initialize all the parameters with Gaussian distributions. Unless specified other- wise, the variances of the parameters are set as 0.01. While our analysis is based on vanilla gradient de- scent, training models using SGD optimizers with small learning rate can be challenging. Therefore, Adagrad (Duchi et al., 201 1) optimizers are used in practice. The network width d is set as 64. To verify the features learned by the models, we ex- tract corresponding co-occurrence pairs for each model from training data. Specifically, for the MLP, CNN, and SA models, calculate the loken-label frequencies from the training data. For example, if token co-occurs three times more frequently with the positive label (+) than with the negative label (-), i.e., freg(e,t)
Here, w(eees z) can be viewed as the frequency of bigram eats seen in instance Specifically, when bigram co-occurs with a positive (negative) label, it will receive a positive (negative) gain dur- ing gradient descent.
We provide the analysis of the L-RNN model in Appendix A. The features captured by different architectures are listed in Table
5 Experiments
We conduct experiments to verify our aforemen- tioned analysis in the following aspects: a) verify the features acquired by our models; b) explore fac- tors that may affect feature extraction; c) examine the limitations of those models.
freq(e.-) 2 3, we will extract this (e, +) pair". For the MV and L-RNN models, we calcu- late bigram-label and laken-label-positeee frequen- cies respectively, and extract co-occurrence pair in similar way. For simplicity, tokens (bigrams) co- occurring with the positive/negative label will be referred to as positivee/neaatiee tokens (bigrams).
Datasets We consider the following datasets: SST, instances with "positive" and "negative" la- bels are extracted from the original Stanford Sen- timent Treebank dataset (Socher al., 2013). We also extract instances with sub- sub-phrases (along with labels) under the name "SSTwsub". Agnews, the AG-news dataset, which consists of titles and de- scription fields of news articles from 4 classes, "World", "Sports", "Business" and "Sci/Tech". IMDB, the binary IMDB dataset (Maas et al, 2011)
"More details are listed Appendix C.
'The Adam (Kingma and Ba, 201 4) optimizer is also con- sidered in Appendix C. 'The conditions in the theoretical analysis relaxed in experiments.

---

tilen
-18
-13
15
-24
-23
(a) MLP, C7
(b) CNN, de
(c) SA. dg 0.1
pon 128mn
Inlon
tooon
meg :00en
2.5
14g I2ken
Reg Iokan
0.5
0.0
-05
1.00
760
140
100
430
500
162
200
100
530
Ezech
Epech
Spach
(d) MLP, 0. 0.001
CNN. 0.001
(0) SA, O 0.001
pos 12kmn
Token
poc tokmn
pPn
rig
3.0
1.00
740
300
602
503
100
200
300
400
500
100
250
300
660
505
Ezech
Epoch
Ezach
(g) MLP (tanh), a. 0.001
(h) MLP (GeLU), a. 0.001
(i) SA+MLP,
0.001
Figure 2: Evolution label scores the extracted tokens from SST over epochs. "pos token" and "neg token" refer to "positive tokens" and "negative tokens" respectively.
5.1 Feature Extraction
72
We illustrate the label scores for the extracted co- occurrence pairs to examine the features predicted by our approach. It can be seen from Figures 2a, 2b, and 2c that the label scores for tokens in the extracted co-occurring pairs evolve as expected over epochs the MLP, CNN, and SA models. The label scores of the tokens co-occurring ma- jorly with the positive label consistently receive positive gains during training, whereas those of the tokens co-occurring majorly with the negative label experience negative gains, thus playing opposite roles in final classification decisions. Similar pat- terns can be observed on IMDB in Appendix C. We also extract bigrams co-occurring majorly with either the positive or negative label from SSTwsub and calculate their label scores using a trained MV model, which exhibits the capability of capturing the co-occurrence between bigrams and labels as shown in Figure
5,0
22-
p-bigram
n-bigram
Figure Distribution of the label scores for extracted bigrams from SSTwsub. "p" refers to positive and "'" refers negative.
be assigned relatively large scores in the dimen- sion corresponding the class shown in Figures 4a, 4b, 4c, and 4d. These observations support our analysis of the feature extraction mechanisms within our target models.
In addition, extend our experiments to lan- guage modeling tasks, which be viewed as variant of multi-class classification tasks, with the label space equivalent to the vocabulary size. Inter- estingly, we observe similar token-labe patterns on Transformer-based language models incorporating self-attention modules despite their complexity.
Our analysis on the binary classification tasks can be extended to the multi-class classification sce- nario on Agnews of four-class. The label scores for the tokens associated with specific class would

---

(a) class I, Agnews (b) class 2, Agnews
(c) class 3, Agnews
(d) class 4, Agnews
Figure 4: Label scores for extracted tokens from Agnews, dataset with four classes. SA model. d 64.
both word and character levels. Particularly, we find that nanoGPT, a light-weight implementation of GPT, can capture the co-occurrence features between context characters and target characters the character-level Shakespeare dataset, and re- flect them in the label scores shown in Figure 5. Given a context character, the model's output is more likely to assign higher scores to target charac- ters that predominantly co-occur with this context character in the training data, thereby making those target characters more likely to be predicted. This implies the significance of a large dataset may be (partially) ascribed to rich co-occurrence informa- tion between tokens. Further details can be found in Appendix C.
ReLU
tanh
GeLU
SiLU
11
II
valid
78.4
68.0
77.3
78.0
78.3
77.8
77.8
SST
lest
80.0
67.3
78.9
78.9
79.9
79.8
79.7
79.9
valid
90.0
90.7
Agnews
test
91.0
90.2
90.6
90.1
90.6
89.6
90.6
89.6
91.5
91.7
IMDR
1est
89.9
89.5
89.9
90.2
91.2
91.4
91.2
91.4
Table 3: Average accuracy (%) SST with scaled
variances (I: 0. = 0.1 and II: a. = 0.001) and different activation functions. 3 trials for each run. MLP model.
vation functions such as tanh, GeLU, and SiLU, which are alternatives to ReLUU Figures 2g and 2h show that these alternatives are more robust than ReLU in the MLP model. This also suggests that while non-linen non-linear activations may not significantly alter the nature of learned features during training, they can affect the balance of the extracted features. Figure 2f shows the SA model is also robust to the change in initialization. However, incorporating an MLP with ReLU activation after the SA model reintroduces bias, as can be observed in Figure 2i, suggesting a possible reason why ReLU was replaced in models such as BERT (Devlin et al., 2019). GPT3 (Brown al., 2020), and LLaMA (Touvron et al., 2023), despite its presence in the original Transformer architecture (Vaswani et al., 2017).
Laoel Score
-2-
8
1
Models' Limitations We aim to examine whether the CNN and SA models have a limita- tion in encoding n-grams in situations beyond con- stituent tokens' semantic meanings. We choose negation phenomena as our testbed, where nega- tion token can (partially) reverse the meanings of both positive and negative phrases, task that is challenging achieve by linear combination. We run experiments on the SSTwsub dataset with la- beled sub-phrases, which contains rich negation phenomena, i.e., phrases with their negation ex- pressions achieved by prepending negation tokens such as not and never. We extract positive and negative adjectives and create their corresponding negation expressions by prepending the negation word not. Figure shows that the SA model can
Figure 5: Distribution of the label scores for target characters majorly (blue) and rarely co-occurring with each extracted context character. nanoGPT. Shakespeare Dataset.
Induced Bias Our approach also indicates that factors such as activation and initial weight vari- ances could affect feature extraction. We down- scale the variances for the final layer weight vectors initialization and compare the learning curves of the extracted tokens' label scores from models with different activations. As can be seen from Figures 2d and 2e, smaller initialization of the final layer weight variance can lead to large fea- ture bias, rendering negative tokens less significant than positive ones in the MLP and CNN models. This may not be a desirable situation, as Table suggests a performance decline for the MLP model with ReLU Furthermore, we compare other acti-
"The discussion of tanh and visualization of SiLU can be found in Appendix B and Appendix C, respectively.

---

-10
(lp-adj [-In-adj MV. negation
SA. pos adj
(b) SA, neg adj
Figure 6: Label scores for the extracted positive adjectives (pos/p adj and negative adjectives (neg/n adj), as well as their negation expressions. SSTwsub. "[-]" refers the negation operation.
capture negation phenomena for positive adjectives but does not perform well for negative adjectives shown in Figure 6b. Specifically, prepending negation word to the negative adjectives does not alleviate their negativity as expected but leads to the contrary. Based on our analysis, the polarity a negation expression relies largely on the lin- ear combination of the tokens' polarity in the SA model. As both the negation word nor and nega- tive adjectives are assigned negative scores, their linear combination will still be negative. This is not a desirable case and not surprising, as recent studies (Liu et al., 2021; Dong et al., 2021; Orvi- eto et al., 2023) have challenged the necessity of self-attention modules. Similar patterns can also be observed on extracted phrases with negation words, on the CNN model, and even the Trans- former model in Appendix C. Conversely, the MV model demonstrates the efficacy of capturing such negation for negative adjectives, as shown in Figure 6c, demonstrating that the multiplication mecha- nism may play a more effective role in composing semantic meanings.
tations in effectively composing n-grams beyond tokens, deficiency highlighted in negation cases. This observation points towards potential neces- sity for alternative models that are adept at handling tasks involving complex n-gram features. This ob- servation aligns with studies by Bhattamishra et al. (2020); Hahn (2020); Yao et al. (2021); Chiang and Cholak (2022), which underscore the constraints of self-attention modules, despite their practical successes. Contrarily, the MV model, based on matrix-vectis multiplication, can better capture such negation evident in both analytical and ob- servational perspectives. This model emerges promising alternative for tasks that hinge the interpretation of n-grams. Regarding activation functions, our findings indicate that the utilization of ReLU does not significantly impact the nature of features learned but can introduce feature bias. Consequently, we suggest the exploration of al- ternative activation functions to mitigate this bias, enhaneing the model's performance and reliability in diverse applications.
Conclusions
5.2 Discussion
Our experimental results verify theoretical anal- ysis of the feature extraction mechanisms employed by fundamental models during their training pro- cess. These findings are consistent even with net- work widths as small as d = 64, a scenario in which the infinite-width hypothesis is not fully real- ized. This observed pattern underscores the robust- ness and generalizability of our model, a conclu- sion that aligns with the insights presented by Arora al. (2019). They suggest that as network width expands, the NTK closely approximates the compu- tation under infinite width conditions while keeping the error within established bounds. In our study, we noted that both CNN and Self-Attentin models predominantly rely on the linear combination of token-label features. However, they exhibit limi-
We propose theoretical approach to delve into the feature extraction mechanisms behind neural models. By focusing on the learning dynamics of neural models under extreme conditions, we can shed light on useful features acquired from training data. We apply our approach to several fundamen- tal models for text classification and explain how these models acquire features during gradient de- scent. Mcanwhile, our approach allows us to reveal significant factors for feature extraction. For ex- ample, inappropriate choice of activation functions may induce a feature bias. Furthermore, we may also infer the limitations of a model based on the features it acquires, thereby aiding in the selection (or design) of an appropriate model for specific downstream tasks. Despite the infinite-width hy- pothesis, the patterns observed are remarkable with finite widths. Our future directions include analyz- ing more complex neural architectures.
"The negation word not appears more frequently in neg- ative instances (2086 times compared to 813 in positive in- stances).

---

Limitations
Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. 2020. Low-rank bottleneck in multi-head attention models. In Pro- ceedings of ICML.
Despite the findings on the aforementioned funda- mental models, applying our approach to analyze complex models like Transformers, which incor- porate numerous layers, non-linear activation func- tions, and normalizations, presents challenges due to the increased complexity. These factors con- tribute to more intricate learning dynamics, making it less straightforward to gain comprehensive in- sights into the model's behavior. We would like to investigate and formulate them in future directions.
Alberto Bictti and Julien Mairal. 2019. On the inductive bias of neural tangent kernels. In Proceedings of NeurIPS.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam. Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vool Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. 2020. Language models are few-shot leamers. CoRR, abs/2005.14163
Acknowledgements
We would like to thank the anonymous reviewers, our reta-reviewer, and senior area chairs for their constructive comments and support on this work. This research/project is supported by Ministry of Education, Singapore, under its Tier Programme (The Award No.: MOET320200004), and Min- istry of Education, Singapore, under its Academic Research Fund (AcRF) Tier Programme (MOE AcRF Tier 2 Award No. MOE-T2EP20122-0011 Any opinions, findings and conclusions recom- mendations expressed in this material are those of the authors and do not reflect the views of the Ministry of Education, Singapore.
Zhihong Chen, Yan Song, Tsung- Hui Chang, and Xi- ang Wan. 2020. Generating radiology reports via memory-driver transformer. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1439-1449, Online. Association for Computational Linguistics.
David Chiang and Peter Cholak. 2022. Overcoming theoretical limitation of self-attention. In Proceed- ings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pa- pers), pages 7654-7664, Dublin, Ireland. Association for Computational Linguistics.
References
Sina Alemohammad, Zichao Wang, Randall Balestriero, and Richard Baraniuk. 2021. The recurrent neural tangent kernel. In Proceedings of ICLR.
Youngmin Cho and Lawrence Saul. 2009, Kernel meth- ods for deep learning. In Proceedings of NeurIPS.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. 2018. On the optimization deep networks: Implicit ac- celeration by overparameterization. In Proceedings of TCML.
Jacob Devlin. Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language under- standing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, Volume (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota. Association for Computational Linguistics.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. 2019. On exact computation with infinitely wide neural net. In Proceedings of NeurIPS.
Marco Baroni and Roberto Zamparelli. 2010. Nouns are vectors, adjectives are matrices: Representing adjective-noun constructions in semantic space. In Proceedings of the 2010 Conference on Empirical Methods in Natural Language Processing, pages 1183-1 193, Cambridge, MA. Association for Com- putational Linguistics.
Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. 2021. Attention is not all you need: pure attention loses rank doubly exponentially with depth. In Proceedings of ICML.
John Duchi, Elad Hazan, and Yoram Singer. 201 1. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(ul):2121-2159
Satwik Bhattamishra, Kabir Abuja, and Navin Goyal. 2020. On the Ability and Limitations of Transform- ers to Recognize Formal Languages. In Proceed- ings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7096-7116, Online. Association for Computational Linguistics.
Melikasadat Emami, Mojtaba Sahraee-Ardakan Parthe Pandit, Sundeep Rangan, and Alyson K. Fletcher. 2021. Implicit bias of linear rnns. In Proceedings of ICML.

---

Albert Gu, Isys Johnson, Karan Goel, Khaled Kamal Saab. Tri Dao, Atri Rudra, and Christopher Re, 2021. Combining recurrent, convolutional, and continuous- time models with linear state space layers. In Pro- ceedings of NeurlPS.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. 2011. I. Learning word vectors for sentiment analysis. In Proceedings of ACL.
Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. 1993. Building a large annotated cor- pus of English: The Penn Treebank. Computational Linguistics, 19323333--330
Michael Hahn. 2020. Theoretical limitations of self-
attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 1515 171.
John Hewitt and Christopher D. Manning. 2019. A structural probe for finding syntax in word represen- tations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Tech- nologies, Volume 1 (Long and Short Papers), pages 4129-4138, Minneapolis, Minnesota. Association for Computational Linguistics.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture mod- cls. CoRR,
William Merrill, Gail Weiss, Yoav Goldberg, Roy Schwartz, Noah A. Smith, and Eran Yahav. 2020. A formal hierarchy of RNN architectures. In Pro- ceedings of ACI.
Arthur Jacot, Franck Gabriel, and Clement Hongler. 2018. Neural tangent kernel: Convergence and gen- cralization in neural networks. In Proceedings of NeurlPS.
Jeff Mitchell and Mirella Lapata. 2008. Vector-based models of semantic composition. In Proceedings of ACL-08: HLT, pages 236-244, Columbus, Ohio. Association for Computational Linguistics.
Alon Jacovi, Oren Sar Shalom, and Yoav Goldberg. 2018. Understanding convolutional neural networks for text classification. Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and In- terpreting Neural Networks for NLP, pages 56-65, Brussels, Belgium, Association for Computational Linguistics.
Jeff Mitchell and Mirella Lapata. 2009. Language mod- els based on semantic composition. In Proceedings of the 2009 Conference on Empirical Methods in Natural Language Processing, pages 430-439, Sin- gapore. Association for Computational Linguistics.
Timothy Nguyen, Roman Novak, Lechao Xiao, and Jae- hoon Lee. 2021. Dataset distillation with infinitely wide convolutional networks. Proceedings of NeurIPS.
Diederik P. Kingma and Jimmy Ba. 2014. Adam: method for stochastic optimization. In Proceedings of ICLR.
Jachoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. 2018. Deep neural networks as gaussian processes. In Proceedings of ICLR.
Antonio Orvieto, Samuel L. Smith, Albert Gu, Anushan Femando, Caglar Gulcehre, Razvan Pascanu, and Soham De. 2023. Resurrecting recurrent neural net- works for long sequences.
Omer Levy and Yoav Goldberg. 2014. Neural word embedding as implicit matrix factorization. In Pro- ceedings of NeurIPS.
Andrew M Saxe, James L McClelland, and Surya Gan- guli. 2013. Exact solutions to the nonlinear dynamics of learning deep linear neural networks. In Pro- ceedings of ICLR.
Jiwei Li, Xinlei Chen, Eduard Hovy, and Dan Jurafsky. 2016. Visualizing and understanding neural models in NLP. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. pages 681 681-691, San Diego, Califomia. Association for Computational Linguistics.
Richard Socher, Brody Huval, Christopher D. Manning. and Andrew Y. Ng. 2012. Semantic compositional- ity through recursive matrix-vector spaces. In Pro- ceedings of the 2012 Joint Conference on Empiri- cal Methods in Natural Language Processing and Computational Natural Language Learning, pages 1201-1211, Jeju Island, Korea. Association for Com- putational Linguistics.
Tal Linzen, Emmanuel Dupoux, and Yoav Goldberg. 2016. Assessing the ability of LSTMs to learn syntax- sensitive dependencies. Transactions of the Associa- tion for Computational Linguistics, 4:521-535.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher Manning, Andrew Y Ng, and Christopher Potts. 2013. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of EMNLP.
Hanxiao Liu, Zihang Dai, David So, and Quoe V Le. 2021. Pay attention to MLPs. In Proceedings of NeurlPS.
Noel Loo, Ramin Hasani, Alexander Amini, and Daniela Rus. 2022. Efficient dataset distillation us- ing random feature approximation. In Proceedings of NewrlPS.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. 2023. Scan and snap: Understanding training dy- namics and token composition in I-layer transformer.

---

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux. Timothée Lacroix, Baptiste Rozière, Naman Goyal, Fric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models.
A Learning Dynamics of Models with Infinite-width
A.1 MLP model
The representation for the instance is defined as
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, ukasz Kaiser, and Illia Polosukhin. 2017. Attention all you need. In Proceedings of NeurIPS.
((x) h Σ Φ(
(22)
j=1
Ivan Vulié, Edoardo Maria Ponti, Robert Litschko. Goran Glavaš, and Anna Korhonen. 2020. Probing pretrained language models for lexical semantics. In Proceedings of the 2020 Conference Empirical Methods Natural Language Processing (EMNLP). pages 7222-7240, Online. Association for Computa- tional Linguistics.
and the label score of x is computed as follows,
8 u
W W ej), Vd
(23)
Sarah Wiegreffe and Yuval Pinter. 2019. Attention is not not explanation. In Proceedings of the 2019 Confer- ence on Empirical Methods in Natural Language Pro- cessing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 11 11-20, Hong Kong, China. Association for Computational Linguistics.
where ((11 is the instance length, W E Rdout is the weight of the hidden layer, W E RdaxxIV is the weight of the embedding layer, and v E Rdouu is the final layer weight. For simplicity, we let dout = dn
Φ is the element-wise ReLU function. ej the one-hot vector for token ej. The gradients of the parameters can be computed as follows,
Greg Yang and Etai Littwin. 2021. Tensor programs iib: Architectural universality neural tangent kernel training dynamics. In Proceedings of ICML.
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. 2021. Self-attention works can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Lan- guage Processing ( (Volume 1: Long Papers), pages 3770-3785, Online. Association for Computational Linguistics.
((x)
Ös
Σ
plc)),
Ov
Vd
g(0)
Ōs
Dj
(W
eT)
(24)
OW
(()
Ōs
W Dju
ej
dWe
d
where
Cj W W
Vd D; ad(c;)
(25)
dej
Note that D, is diagonal matrix with elements being either Is or Os.
Given test input x', the learning dynamics of the label score s' will be
8'
Σ (x.0)ED
(1-ys(t)yy0((s
(26)
m
With the gradients, we can obtain the NTK

---

(x', æ) for this MLP model follows,
We will give the proof along with the proof of Lemma 4.1, where the test input simply a token
1 Re0 (6))
(() = d ΣΣ Φ (c))(c))
Proof. We only need to compute the NTK when the network width grows to infinity, as the NTK's s convergence during training has been proved in the work of Jacot et al. (2018); Yang and littwin (2021).
1l05s paa
v ve (We) We
Elj=1
1(0s (()
For the first part in Equation 27, the dot-product between two activation outputs can be written as,
+ d
(WW " Duve ej, ej,
1=
(27)
& (c)))(c)) Σ Hcw)(())))
(29)
where
c; = W 1 W
vd
where r refers to the row index and
(28)
D: 0o(c)
Cir Wr 1 W
dc.
Vd
(30)
The label score of an instance can be viewed as the sum of the label scores of all the tokens and the NTK can be viewed the sum of the interaction between each token pair from the test input z' and the training instance x.
Cjr = W
Weej.
Vd
As elements in W and We follow Gaussian distri- butions, when the network width d -> co, Cir and Cjr will also be Gaussian distributed respectively and they follow a Gaussian process (Lee et al., 2018). Based on the work of Cho and Saul (2009), the covariance K(o(cr)) (((r))) o(C,))) (regardless of r) will be calculated as
A.1.1 NTKs under the Infinite-width
It's delicate to analyze such an NTK directly in practice as the NTK will vary over time. How- ever, the previous work discussed in the literature has proved that the NTK will converge and stay constant during training under the indinite-width condition. Now let us consider the infinite-width scenario (Lee al., 2018) and obtain the NTK subsequently. We first give the NTK between two instances and then give the NTK between an input token and instance.
(((()))) 0.00 [sin any - any) cos onyl,
(31)
where ouj = cos e,Tej. ej.
Then, we arrive at
lim d d
(ci)((c)) = Bld
NTK between instances Assume we initialize parameters following Gaussian distributions, i.e.. Wi ~ (0 ).), W (0, o²), and v, ~ NO, o
(32)
0:000 (sin only) cos azy]]
2n
For the second part in Equation 27, let us look at e, (Wc) W first. Let M = (W*) We We and the elements of M can be computed as
Lemma A.1. When the network width approaches infinity, O(r',an x)) converges to deterministic NTK (x' .r) during training which obeys
d
M Σ W W
(33)
(x)
(6r)) ((x)
2 [sin aijo aij + (x aci)) cos ang]
where We We rj
NN0 o?).
When d xc,
2n
d
lim do000 d
M ij = lim
w W = E w
((')) (()
d d
Σ
000000 + +0 2 5
e, Te
r=1
2
0 if
0,
otherwise
where 0zj = cos e, ej.
(34)

---

We can thereby arrive at
Let us look the token-label features learned in the dynamics. As previously mentioned, the instance label score can be viewed as the sum of token label scores. Consider the scenario where the test input x' is simply a token the NTK 2 (e, obeys
lim Jot (Wo) T ee = = ej. yoc
(35)
Let us look at v D,D If e ej 0, which means D; + Dj. we will arrive at
(() pe
(()
1 v D d e (w")
lim d0o00 d2
(x, - j=1
T ej
Σ 5=1
14
(40)
lim
1
Din
Dir D jm
(36)
d+x00 d
where the dot-product preted as the frequency of e appearing in instance T. We can thereby arrive at Lemma 4.1.
دj= == eTe, ej will be inter-
((r))
r=1
0,
where Dirr and jer diagonal elements in D, and Dj respectively. of is the r-th element in v. If e, e, # O, which means D: Dj, we will have
A.1.2 Features Encoded in Gradient Descent
With Lemma 4.1 and Equation 5, we can get Theo- rem 4.2. Under the infinite-width, the dynamics of token e's label score obey
lim - D.D (w") e" "W", d+030 d
d
8 = 1 Σ
(37)
g(- (10
2 e
m (x,))D
Similarly, get the third part in Equation
(41)
Σ (Dr
g( ((-ys(a) ((-y)(=)))111))
27,
(e') ((r
lim ΣΣ 5=
D.)
DiueTTe,
d tos d
(38)
where w(e,a) is the frequency of token in in- stance w(e,r) depends on the training data and will not change over time. We cannot give a closed-form solution for this ODE due to the non- linearity of the sigmoid function ş( tyyy... However, as g( ((yys(z) is non-negative, there can be certain interesting trends for the label scores.
= 2 e
Plugging the above equations in Equation 27, we will arrive at Lemma A.
NTK between a token and an instance Next, we give the proof of Lemma 4.1 based on Lemma A.1.
Note that the first term A Equation 41 will depend on this token's term-frequencies w(c, x)) in each training instance. The second term B depends on the entire training set and is shared by all the tokens For example, if does not appear in instance I, w(e, z) will be 0.
Proof. Note that as e, and ey are one-hot vectors, their dot-product e._ e, satisfies that, e,_ e, = when e; = Gj or otherwise. Therefore, 0;j- = cos e, ej, aij can be E or And the kernel can be further written as
2
A.1.3 Bias Induced in Gradient Descent
((al) ((x)
(a'h) (())
Let us look at the term B in Equation 41, which can be viewed as an induced feature bias shared by all tokens. It is affected by the variances and the instance lengths. Suppose the term B is sufficiently large; in this case, the positive tokens and the nega- tive tokens will be affected by this bias in different directions during training. For example, if term B is positively large, it will positively contribute to the learning dynamics of positive tokens and make their label scores much larger than 0 after sufficient updates. However, the negative tokens will have
(z0 = ΣΣ Σ
pe ej +
(39)
(TT 100002
2.0 + 02
where p
and /
2
0202
2T This means the converged kernel e (4', x) will keep being non-negative during training, and the direction of the dynamics will depend on the label y in Equation 26.

---

weakened learning dynamics and end up with label scores close to O.
where
In Equation 41, both w(e, x) and ((x) are deter- mined by the training instances. Therefore, the factor p and jl defined Equation 39 can affect the influence of this induced bias. It can be inferred that, significantly large variance By can make p much larger than nl, thus reducing the influence of the bias.
Cj
W
W
ejtk-1
(45)
Vd
dΦ(c,)
Dj dcg
(46)
For input x', the dynamics its label score will obey
A.2 CNN Model
8/6
Σ (z,y)ED
vgl uo(-vss)6(ar,z))
(47)
We consider the I-dimension CNN, which is com- monly used in NLP tasks. The kernel size, stride size, and padding size will be set as K, 1, and K respectively.
m
where
For each sliding window that consists of K consecutive tokens, the corresponding feature ej E Bdou can be represented as
( -E
j (c))(())
T de
Σ
D,D uk- - (Wa) " e-tk-1++
Gj.k
9j - Σ WC W Weetk-1 k=1
(42)
B
Vd
(WE) D, Dj -Bith--1
where W E Rdou is the kernel weight cor- responding to the k-th token in the sliding win- dow, W E dg xV is the embedding matrix, and e is the one-hot vector for token c. We also let din = dout = d.
(48)
where Σigh - EE and ENjAN DE K
Given an input x, the label score will be calcu- lated as
A.2.1 NTK under the Infinite-width
It should be highlighted that, even when the net- work width approaches infinity, it may not be easy to describe the converged NTK with an explicit closed-form expression due to the integrals used in obtaining expectations. However, we will show that the converged NTK can be written as functions of parameter variances and similarity between slid- ing windows, and the functions are affected by the similarity between sliding windows.
(e)
v
Σ
((C)),
(43)
(K-1)
where -KK-1) 1K-11 1) means the position for the left- most padding token. The first and last K tokens in an instance are padding ones represented by zero vectors. Φ is the element-win wise ReLU func-
Lemma A.2. Assume we initialize parameters following Gaussian distributions, i.e., W (0, W ~ ~N(0, and ~ N' a?) (W5 refers to an element in WG). When the net- work width approaches infinity, given two instances x' and x, the NTK Oa', for the CNN model converges to
tion.
For brevity, we will denote Σ ーj=-(K-1)
(()) (K-1) by Eg
The gradients will be computed as
K
0s
Σ
Σ
W
W 23in
0o
Vd
Vd
J
(x)) = ΣΣ F(we(cr
0s
Dju
Σ
(Weeik T)
aw: 8
1
0. H (wele, (, w2)en)))
Øs
(WG)
Dju
ΣΣ
e 1
dWe
d
Σ Σ
o 0.2 2 H (Wo(o, (alci,,)))) wi)) cj),
(44)

---

where 6s and ej are sliding windows starting from the i-th token and the j-th token in instances x' and respectively. Functions we and w2 defined as
W is the r-th row of matrix W.. Elements in We and We follow Gaussian distributions and are L.I (Independent and identically distributed) random variables. With the infinite network width, given token c, elements in W 1 W can also vaa
Welci, cj) =
Σ
e,
-- Σ
ej+k-1
be viewed as L.D. and follow a Gaussian process (Lee al., 2018). Let w W Va get that w N (0, Similarly, we can obtain that C₁ and Cjr follow Gaussian distributions.
W=1
k=1
Wee We can
K
w2(ce, ej) Σ
e
eitk-1ejthee
k=1
When the network width approaches infinity, the multiplication will be viewed as the expectation as follows.
Functions F and are defined as
F(n)= =
E
We
Σ 1=++1
we
wet Σ
lim d 100 d (c))(c)) B(o(()))((;))) (51)
1=++1
H(n)=
Then, we can arrive at
K
n
E(D(
wi+
Σ i=111
w;) D
wi)
lim d
d (ci)(c))))
=
1=
n
where (0 < n < K) is the number of to-
Wep
Σ wi # i=n+1
Σ w()l
F(n),
kens shared by two sliding windows and w;, w/ r N(O, V00,0"0a) o2o2). @ and are the ReLU function and step function respectively.
i=++1
(52)
where n is the number of shared tokens between sliding windows 62 and Cj-
Remark. can be seen that the similarity between sliding windows influences the converged NTK. As the variances are constants, we can focus on We and w2, which can be viewed as similarity met- ries for sliding windows. The former does not take positional information into consideration, while the latter does. Particularly, we can have that we(ci, cj) 2 wy(ci, oj). When the two sliding win- dows share tokens in the right order, w2(c;, 0ÿ) be- comes large.
We prove term B in Equation 48 can be writ- ten as a function of (n) under the infinite-width condition.
Proof. With the infinite network width, we can have
1 d d d
lim d d" D.D. v- JUDDD..
lim
Σ
0 Diry
D Djrr
Proof of Lemma A.2 We give the proofs for cach part in the NTK shown in Equation 48. Let us prove that the ReLU output multiplication (term A in Equation 48) can be written as a function of F(r F(n) under the infinite-width condition.
B D jrr!
lim T (Wo) W e
d
= -
Proof. The ReLU output multiplication can be writ- ten
(53)
where r is the row number. Dar and Djrr refer the r-th elements in the diagonal positions of D; and Dj. respectively. Term B will obey
d d Σ (cnr))
d 20 (c))(c))
(49)
where refers to the r-th element and
T D.D. ve VeT (W T W' W'efth11
K
lim
Σ
Cir
W
kr
Vd
W"
eitb-11
d
k=1
d
k=l
(50)
2
Di D jrrle DireeTTk 1644h-1
(54)
Cjr
W Wo
br
Vd
e
We can compute the expectation D D jrr (outputs of the step function) similarly to that ReLU output multiplication.
1=1
"D(x) = if x > 0, 0 otherwise.

---

Let wy (0 Dirr (D equals 1 when the corresponding Car 0 (cjr > 0), 0 oth- erwise. Then, we can obtain
The expectation can be computed as
F(n)
E
w
+
w))
W + Σ wS)l
E Dirr Djrrl
1=++1
i=++1
K
7
K
=E[D(
W
Σ 1=
D
W
Σ
w
(( + Rm xn+1x)) + +
1-
p(w)) ... p(wk) (ikk)))(( )p(w'n wnn 1)... p((kk)))))) dwn dwx f...+.....du
(n),
(55)
where n is the number of shared tokens between sliding windows ej and Cj.
Elp(wi.... ... w.)))))).))) ),)
(57)
For the third part (term in Equation 48), when the network width approaches infinity, the expecta- tion will be
where
1/( (w) .......... (.nn.)..))-- = 100
(z + 2n 1 1 (p( 1)
... p(w ) .... dwk
V' lim D (W T DJ vele, eg
100 d
(56)
Zim Σ Wi
- E D Dimle,
K
Σ
It is obvious that this term will be positive if the two windows e; and ej do not share any tokens, otherwise.
In+1:K =
Wi
1=++1
2xt1:K
1=++1
(58)
Now, let look at the F and II functions, which have interesting properties regarding the similari- ties between sliding windows.
and 1_ refers to the integrals for all the variables involved.
The expectation for the multiplication of two sliding windows without sharing tokens can be written as
Proposition A.1. Both F and H functions in Lemma A.2 are monotonically increasing as the on on-negative integer n increases. Given two non- negative integers n and n', when n' > n, the two functions obey
F(O)= =
K wa
K
E W4 +
wa)
+ i=n
Σ
w!)]
F(n') 2 F(n), H(n') 2 (n).
= E (R
(1
wx,))
(59)
which can be described
F(n)
Remark. This indicates the more similar the two sliding windows are (i.e., the more tokens they share), the larger and H will be.
K
K
E
+
w!)l
1=
The core idea leveraged in the proofs based the inequality, Var(x)= = H(x²) E(x)² 2 where is random variable yiclding Gaussian distribution.
F(0)
K
K
E (
wO},
(60)
Monotonicity of F function
where wy and w/ are L.I. random variables and 2 1. This indicates that when two sliding win- dows share tokens, the expectation will be larger than that in the case where two sliding windows do not share any tokens.
Proof. First, let us consider the scenario n > O. which means the two sliding windows share tokens.

---

We can prove that one more shared tokens be- tween two sliding windows can result in an in- crease in expectation. Let us increase the num- ber of shared tokens by I between the two sliding windows, the expectation can be written as
Monotonicity of H function We can obtain the expectation of the other terms in the NTK similarly as the activation function & can be replaced with different activation functions with non-negative out- puts.
F(n-
+ +1)=
Proofs of Lemma 4.3 and Theorem 4.4 With Lemma A.2 and Proposition A.1, we can prove Lemma 4.3 by replacing the test input x' with slid-
K Σ
K
E
Wig +
wi)
W
wi)
ing window c. Similarly, we can prove Theorem 4.4 with Lemma 4.3.
i=++2
Φ M2i++1 + 2n+2:KK ( s1m+1 + 2.+2.K)
Let us focus on a single sliding window c. The converged NTK between a sliding window e (con- sisting of tokens 62,02,.... 02,,..,. ex) and instance x obeys
p(w1) ((w))...p((ee p(uk ple dwi... "wnl...duddee dwxdw .du''
pl((k)
+x (c, [((w)(a)))
[/ E(w),.... wn+1) 1 ) Sp )dwn+1
+ Σ ( (oncdcccc))))(ccc)
p(w1). p(w,))d) dwn.
(65)
(61)
where E(w1:..., Wn+1)
S d(zim+1 +
+ Σ L To H(w (W)): c)))(((c))))
In 2: (p(Wa42).... ... (1.42)..)(((Ap)) ((wp)((wwee )dwn+2 ... duk.
Note that the expectation Elo(E= wi +
EK wi) wi + EK' w/]l can be
If e does not share tokens with any of the slid- ing windows in x, the NTK Oco x) will reach its minimum, namely, (c, = E; (O).. Otherwise, Oxe (c, x) will be significantly large if bears similarity to the sliding windows of æ. Then the dynamics of the label score of e obey
written as
F(n)
K
n
B Bl
w
W
Σ
w
i=++1
++++
[ f(w1..... Fer...,,pp+ll )p( (n+ d
2
8
P m
Σ (R.
vgl-ys (1)-y()))000) (c, x).
p(w1). ... ((w))dde )ww dwn-
(62) The sliding windows can be interpreted as n- grams. Suppose an n-gram represented by a sliding window e bears similarity only to the sliding win- dows in positive (negative) instances. In that case, it will receive positively (negatively) large gains in one direction during training and will likely end up being significant.
As we have
E(w Wn+1 ) ((p )dwn+1
2
2 (63)
[ E(wi,..., wn11):..
00
Let examine what features will be learned for single token e. We define rositional-relevance label score as follows,
we can arrive at
F(n+1)=
x+1
K
n+1
K Σ i==+2
1 T of W..").. U
1
E
W +
Σ
wi)d
wi)
Wy +
wf)l
(66)
1=1
i=++2
d=1
F(n)
which reflects the label score of token in the k-th position of sliding window.
K
K
E
Wi +
w
+
w()l-
It should be noted that based our analysis, the kernel size K in Equation 10 does not affect the monotonicity of the F and H functions. Suppose the sliding window e shares tokens with instance æ;; the NTK will learn the features regardless of the kernel size K.
1=n+1
i=n+
(64)
We can prove recursively for the case F(n + I) F(n) where >> 1. Therefore, the expectation will be monotonically increasing with n.

---

A.3 SA Model
ents can be computed
Let us define an intermediate score 8j, correspond- ing to the label score in the j-th output, as follows,
0f 1
Σ (WW o (Wo)+)))
Ov dvd
Of 1
[(x]
Σ E(( (W 0
OW dvd
8j
W ej-
(67)
0f OW
dvd Σ W e(W"e (
The gradients be computed as
1
d Σ 1 (WW
((0)
0s;
Csj
ej
where @ refers to element-wise multiplication. (71) Given a bigram ca the NTK will be computed as
dv
Vd
((0)
(())
0s:
V W
aig ask)
OWe
d'd
j=1 k=
Cco x)) x)> Σ
eee T
+ Weeeet + P.ek P
((a
Σ J=1 Vd aijve
(WW") diag(W efn WWe,
d3
(68)
(w") W e.e (Wo) diag(v) ( W'e
where ((a) is the instance length and dzk- j = k, otherwise. Ps is the positional embedding at position k.
d3 eae, 1 (Wo) (iag(()) W W
We assume that the parameters are initialized as W ~ (0, o.), and vy (0,,0)) and the dis- tribution of the attention weights are independent of the parameters W and v. Let consider the case where the test input is simply token e. If the network width approaches infinity, T. W'e; + and NTK between the token e and the instance x will converge to Oo (c, which obeys
J c (WWe) (iiag))) (raago))WW's
T e ((e)))
d ex Y(ee),
(72)
where
X(es) = W (W" es) o
((0) (())
(73)
((e,) = e (WW*e;)).
2)
Blacg)(o + of)eTe ej,
When the network width approaches infinity, the NTK will converge to
(69)
where B((sy)) is the expectation of ouj when ele- ments of W obeys Wi o²)
x)= (caett) = E
A.4 MV Model
o eaestie
The label score of an instance is defined as
of e
(74)
1 dvd
+ ooo I E eae
s= v Σ M(e,)
(70)
of eae
Ei + 30.0 e eae
where M(e)) = diag( (Wee) e (diag converts vector into diagonal matrix.) ., (c) 1.
and j
which can capture the co-occurrence between bi- grams.
The proof sketch is given as follows: The gradi-

---

A.5 L-RNN Model
Lemma A.3. Assume initialize parameters
following Gaussian distributions, i.e., Wij (0, o W ()), W ~ (0, and Vj ~ (0,00 (0) When the network width ap- proaches infinity, the NTK ( x) converges to a deterministic one Ox (a', ,r), which obeys
We follow the work of Emami et al. (2021) and Gu et al. (2021) and focus on a linear RNN, whose hidden state is defined as follows,
1 w h + WW Vd
ht
(75)
d
min 11 -1
where / E Raxd and W E Rdxd and the initial hidden state zero vector. We can expand the hidden states across time steps and obtain
2x (x
Σ
h o.
2h T
keT E1
min(T
Σ
2k POO T
hh
Ei kET-k
Wh )r 2--W ew
h4
(76)
d j=1
min(T TT
2h 2k hh o ker-k
where 1h 0 I. Vd
min(T
The label score of an instance computed based on the final hidden state as
Σ
Ko 2k
2000 keT-k-
+
s = e hT
where Ao = o......, min 1, which indi- cates the distance from the last tokens.
1 T
Wh
Σ
WW
Proof. The multiplication between b and its transpose can be computed as
1=1
(77)
I
W
Wh
8j1
(W T Wh,
(79)
Vd
1=1
where each element wy in Wh follows wy N (0, az) Each element w' in the output of the multiplication will be
where 40 T Wh - and T is the vd
final time step. Note that T j means the distance between the current token and the last token in an instance and sj can be viewed as the label score for the token with a distance of T - from the last token. The gradients will be calculated as
lim w/y = lim Σ d-000
WrWWy = E[wn(wrj]
d=
oh if
00, 0,
otherwise
Ös
Wh
T- T-iWWee
(80)
Ov
dd
j=1
Hence,
0s
1
Σ
T
To
Wh
T
WW
lim d d
W
Wh
II
(81)
dWe
d
Vd
==1
where Rdxd an identity matrix.
0s
T
Wh
T-j
T
W e,lT
We can also obtain
OW
(78)
W k
0s
lim d=o00
-
- 0,
(82)
OWK
where & is an integer and the elements vectors va and Ug E Ed are Gaussian distributed with zero means. Based on this, we can arrive
TT
T Wh
T-j-k-IjT
Vd
Wh
Vd
"wWee,]
lim 3 100
d
1
T
Wh Va
k'
Wh
U3
(83)
Wh
where T 1. When T 1, ow5
does not exist.
2k lim d
0
T
Vd
0g
0.

---

where V & (both are integers). A similar conclu- sion can be obtained for the case k' < k. Given instances x' and z, whose label scores are s' and respectively, we can have
the hidden states into a L-dimension vector (L is the label space), and the sigmoid layer will be placed by a softmax layer.
The label score of an instance will be described as:
Os' Os
lim 100
s(1) fe(1; 0z),
(88)
Ov dv
T
W
where s(t) is vector with dimension of L. For brevity, we omit the denotation L.
lim ΣΣ il
e (Wo) T
( ) vd
d
We can get the probability distribution for all the labels:
Wh T WWW Vd
p = softmax (8),
(89)
(84)
where E R
Like what we have done previously, we can obtain that the elements in e, (Wr)* w and e are Gaussian distributed. Let vn Va (We)TW" and U3 = WW we
The cross cross-entropy loss will be used, and the loss can be computed as:
log (p(at) T (x)) Σ (I,Y)E
vd
(90)
will get
0s' Ōs
where p(() refers to the softmax output for instance x and y is the one-hot label vector for instance x.
lim
d000
Ov" dv
Σ
The derivative of L with respect to us is computed
hh T T e ej
(85) as follows,
-i-T-
Oc 0g(²)
min(T",T) Σ
(z) - (())
-p(a)).
(91)
2h T 0 w
Re
Given a test input x', the learning dynamics of its output from the model with the infinite-width network can be described as
where T" -i= T'
k
Similarly, we can obtain terms aw
0s oW
Σlv (1) (z0 (Z.))D
0d' 0s OW > and < WW
0x 0s OW
(92)
OW
where 4 (x E I.X. refers to the converged
Let h- = T j. we can re-write sj in Equation
NTK determined at initialization.
77 as
1 ' d
Wh
We can have a similar analysis on this dynamics (z) - p()) will only be positive in the dimen- as (g
WW
(86)
sion where y (x) = (z is the dimension index in y). Suppose @ (a', x) E RRLL works in way that increases the influence of the corresponding dimension z in s' when a' is associated with label i.e., the label score corresponding to such a label will receive positive gain and grow to be large.
which means the label score for token at position l from the last token. We thereby define an NTK (c,, k; .) to represent the interaction between to- ken at position A and instance x.
When the network width approaches infinity, the NTK (e,, k x)) converges to a deterministic one 3o (c, kt x), which obeys
B Influence of Activation Functions
Ozo k, = ppk)) er er-k;
(87)
As we previously mentioned, the ML. model can have an induced bias for the features learned be- tween positive tokens and negative tokens, while the self-attention model does not have this problem. Actually, this problem is caused by the activation function ReLU used in the MLP model. Apart from the ReLU function, we considered the identity map (Φb - I), and tanh activation functions.
where p(k) o2th 2220002 to 220000 t0 2.02+ 26 ko "oo00o and k is non negaavete integer. To
t0
make it consistent, we can also replace the instance length with 1(x)
A.6 Multi-class Classification
Compared to the binary architecture in the main pa- per, the last linear layer will be modified to project
Let us consider token e from the vocabulary and its NTK with the training instance x..

---

- In this case, the model is linear. When the network width approaches infinity, (xr', x) con- verges to deterministic NTK cc (a', x) during training which obeys
Dataset
Train
Dev
Test
Token Num
887,521
70,390
78.669
PTB
Vocab Size
10,000
Token Num
2,088,628
217,646
245,569
Wiki2
Vocab Size
33,278
He' (()
Shakespeare
Token Num
1,003,854
111,540
Vocab Size
65
(x) =
2. e
Table 4: Language modeling datasets statistics.
((0 [(m]
00
Dataset
MLP
CNN
SA
MV
L-RNN
valid
79.4
79.7
80.4
77.3
79.4
SST
81.5
81.0
80.7
77.8
80.7
((' [(m)
valid
79.6
78.7
79.9
80.5
2 70
SSTwsub
test
79.0
79.2
81.3
77.5
valid
91.5
91.8
91.7
IMDB
test 91.8
91.4
91.6
(e' (0)
valid
91.5
91.4
91.7
ΣΣ =
3 + +0.00 00 ej.
Agnews
test
91.3
91.1
91.5
(93)
Table 5: Accuracy (%) on the datasets. Adagrad opti-
mizers. 1_" refers to that the training unstable.
Given token the NTK Oo0 (c, obeys
((r)
Suppose e ee = 1, we can obtain
Oxc = (0.00 + 0.00 o
E E[tanh (Cer) tanh(e)))) ([tanh"(cer)))
(94)
1
lim d+oo d²
DD (Wo) W
which means the NTK is affected by the frequency that is seen in instance This is similar to the MLP model with the ReLU activation function but without the induced bias.
LF D2 Je
(97)
lim DW)WW) D De
d d2
D e,,
Φ = tanh Suppose Φ tanh, then we have that
where Djr is the r-th diagonal element in Dj- Given token e, the converged NTK computed as
1 lim (co))((e)
d=000
1 lim tanh ( tanh (bez) E[tanh((or) tanh(c)+)).
(95)
(())
d vo0 à
(c) = Σ
Estann Ettanh"(cer)]e (cer)]eTes ejt
(98)
[x] Σ
D jrrl 2 2 ej. j=1
where is the row index and E[tanh(cer) tanh(c is a constant regard- less r.
This indicates the MLP model with the tanh acti- vation does not have the induced bias.
It can be inferred that cer yields to Gaussian distribution with a zero mean. As tanh is an odd function, we will arrive at E[tann((Crr)] =.
C More Experimental Results
Suppose e eTe, e, 0. namely, * ej. tanh(cer) and tanh(c)r) are two independent random vari- ables and Eltanh(cca (Cer) (anh(c,r) 0. Similarly, when Cj. we can obtain
The statistics of language modeling datasets listed in Table 4. The performances are listed in Table 510
Extracted Tokens& Bigrams We automatically extracted tokens associated with specific labels as shown in Table 6.
lim 1
d d²" a0 DD, (W9) W DD
(96)
1
"The MV and L-RNN models are not stable during training on IMDB and Agnews.
d 300 N0 (DW(W)) Djve e; =
lim

---

Dataset
Extracted Token (bigram) Num
150
SST
412 (+) 265 (-) 73 (+) 47 (-)
SSTwsub
SSTwsub (bigram)
779 (+)/ 542 (-)
TMDB
414 (+)/363 (-)
Agnews
441 (Iv540 (11)/297 (1I1/334 (IV)
Table 6: Numbers of extracted tokens. "(+)" and "(-)" refer to tokens associated with the positive and nega- rive labels, respectively. "(I)", "(II", "(IIl)", and "(IV)" refer to Class 1-4 Agnews.
20 13
(a) PTB
15.0
Adjectives with polarity were extracted automat- ically from SSTwsub. We first extracted tokens that were seen more frequently in either positive or neg- ative instances than in the other instances, i.e., the frequency ratio either larger than 3 or less than 1/3, Then we used the textblob package 1I to find out adjectives from those extracted tokens. Examples are shown in Table 7.
(b) Wiki2
Figure 7: Distributions the label scores for majorly co- occurring token pairs (in blue) and rarely co-occurring token pairs.
Tokens
inventive, nice, authentic, sympathetic, lovable, grand, happy, enthusiastic, noble, detailed, exotic, remarkable, charismatic,... inexplicable, feeble, sloppy. disastrous, stupid. terrible, unhuppy. horrible, atrocious, idiotic. angry, uninspired, vicious, unfocused, artificial,..
relatively larger values in the positions correspond- ing to the majorly co-occurring tokens than in the positions corresponding to the rarely co-occurring tokens.
Table 7: Examples of the extracted positive adjectives ("+") and negative adjectives ("-") from the SSTwsub dataset.
For the experiment on nanoGPT, we followed the settings in the quick start part and trained the model on the char-level Shakespeare datasct. The max iteration number set as 2000 instead of the default 5000. We extracted the most 20 frequent chars and calculated their co-occurrences with each char in the vocabulary (65 chars in total). For each extracted char, we extracted the top 5 and bottom 5 co co-occurring chars, respectively from the vocab- ulary. Feeding each extracted char into the model, we could get the output vector, as well as the label scores for majorly co-occurring chars and rarely co-occurring chars.
C.1 Language Modeling
Language modeling can also be viewed as classifi- cation task, where the label space is the vocabulary size and each label is a token in the vocabulary. The model output before the softmax layer is a vec- tor with a dimension of the vocabulary size. Each dimension corresponds to a label, i.e., a token in the vocabulary. We extracted 20 most frequent words from the FTB 12 and Wiki2 datasets (Mer- ity et al., 2016), respectively. For each extracted token, we searched for the top 30 token-label pairs. In addition, searched for the bottom 30 token- label pairs for comparison. We trained two-layer Transformer language mo on PTB and Wiki2, with the embeddings size and hidden size of 200. Figure shows that the co-occurrence can be gen- erally captured by the Transformer model. Specif- ically, given an extracted token, its output (with a dimension of the vocabulary size) will likely have
C.2 Negation
Aside from the results on the self-attention model, we also conducted experiments to verify the ability to capture negation phenomena on the CNN and the L-RNN models. It can be seen from Figure 8 that both the CNN and Transformer models can capture the negation for the positive adjectives. But it seems they do not capture such phenomena on the negative adjectives. The CNN model handles the negation bigrams in a linear combination way, and the polarity of a negation bigram is the combination httn:://gnncbooooounnollggggennr htwp:://maaaaommmmmttuutttnaaaaa the polarity of the two tokens involved. As the
"I hhttps://textblob rew.oddddeeaaheeeeeeeeeeoeoo iotp:///./////////////// chttps://catalog.lee upw.. uppem..eddu/oocc/LLcc5/// c193. html

---

12.5
pON BO) regaliio
7.5
Label Scoee
2.5
-12
-5.0
-14
1Teg [c] maagaller
20
60
10
20
30
40
(a) CNN, pos adjectives
(b) CNN, neg adjectives
0.
-1
Label Score
froe
20
52
(c) TR, pos adjectives
(d) TR, neg adjectives
Figure 8: Label scores for the positive adjectives (pos adjectives) negative adjectives (neg adjectives) well
their negation expressions. "TR" refers to the Transformer model (one head, one layer)
pos tokio
neg token
Score
-2
1.2
100
200
300
400
500
Epoch
Figure 9: The label score differences between the sub- phrases and the phrases. SSTwsub. SA model.
Figure 10: Label scores for extracted positive tokens and negative tokens from SST. MI. with SiLU.
token appears more in negative instances, is assigned negative label scores. Therefore, adding the token not to positive adjective can weaken its positive polarity but for negative adjective, adding not can strengthen its negative polarity. This also implies the limitation of such models: they rely largely on taken-label features.
C.3 Feature Extraction
SiLU Figure 10 shows SiLU can also prevent an induced bias in the features captured.
Adam Optimizer We conduct experiments on the Adam optimizer and observed patterns (shown in Figure 11) similar to those from the Adagrad optimizer in the main paper.
We also extracted 65 phrases (less than 1l words, positive and 48 negative) starting with negation words not, never, and hardly from SSTwsub. We computed their label scores and the label scores of their sub-phrase constructed by removing the negation words. Figure 9 shows the differences between the label scores from the subphrases and the phrases are all positive, indicating the negation words play negative roles in linear combination and do not reverse the polarity of negative sub- phrases.
L-RNN Extracting sufficient tokens that appear in a specific position and specific category of in- stances in real-world datasets like SST is not easy, Instead, created synthetic dataset (1,000 posi- tive instances and 1,000 negative instances) based on three types of tokens. One type of token is seen in positive instances with fixed distance from the last tokens; another type of token is seen in negative instances with a fixed distance from the last tokens.

---

pas takon
15
ney taken
-5
-15
100
200
300
400
500
Epoch
(a) = 0.1
pos token
30
--.. neg toten
100
200
300
400
500
Epoch
(b) Go
0.001
Figure 11: Label scores for the extracted tokens from SST. Adam optimizer.
The other tokens are seen randomly in both positive and negative instances with random positions. In this experiment,
set the fixed distance from the last ones as 2. Adagrad optimizers were used. It can be seen from Figure 12 that when k = 2, the positive and negative tokens are assigned signifi- cant label scores, while when = and k = 4, the label scores are less significant, supporting our aforementioned analysis on the L-RNN model.
IMDB IMDB is a dataset with relatively long instances. Our findings can also observed on the IMDB dataset in Figure 13. Particularly, we could also observe a feature bias on the IMDB dataset (as shown in Figure 13c) when we used the MLP model with ReLU, supporting our analysis in the main paper again that ReLu may cause feature bias. However, did not observe obvious performance decline on the IMDB dataset.

---

70g Mnan
oog Laka
(a) k=0
(b) k-2
(c) k=4
Figure 12: Label scores for the positive tokens (pos token) and negative tokens (neg token) at different positions for the L-RNN model. refers to the distance from the last tokens. Synthetic dataset.
1.0
(a) MLP, Ce 0.1
(b) Self-attention, 0.1
(c) MLP. 0.001
Figure 13: Label scores for extracted tokens from IMDB,
64.


Reviews:
[{"review_id": "Sap1DnrI8E", "reviewer": "Reviewer_SPeX", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for your valuable comments and suggestions!\n\nWe appreciate the importance of differentiating between correlational and causal analyses. The primary intention of our paper was not to present a causal analysis. Nevertheless, our work offers some causal insights where we considered the influence of significant factors in the learning dynamics and developed experiments to verify the phenomena predicted. \nFor example, as shown in our manuscript, Equation 8 indicates the presence of a bias in the evolution of the token label score for the MLP model with ReLU activation. This bias can cause an imbalance between positive tokens and negative tokens. Furthermore, we have expanded on this discussion in section A.1.3 (in the appendix), emphasizing the role of variance σ_v in influencing this bias. As shown in Figures 2a and 2d, with the change of σ_v, the bias becomes pronounced, thereby diminishing the prominence of negative tokens. We believe these discussions indicate a form of causal analysis.\n\nIn addition, we would like to clarify that we did not seek to conduct correlational analysis in our experiments. The main intention is to validate our analysis in terms of the phenomena predicted (e.g., Figure 2) based on the terms in the learning dynamics, not investigating the statistical correlation between factors and output. \n\nWe acknowledge the feedback regarding the potential ambiguity in our presentation. We will revise our paper to clarify our primary intention and explain in sections where the causal insights are discussed. We will also adjust Figure 3 for better readability.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The authors provide theoretical models to analyze the learning mechanism in MLPs, CNNs, and self-attention layers.\n* Empirical results provide insightful results relating to activation layers and the sensitivity of models to initialization.\n* Future work in this direction, especially on larger neural networks, can benefit from the groundwork laid in this paper.", "reasons_to_reject": "* The empirical results mostly offer a correlational analysis and not a causal or mechanistic understanding of the underlying phenomenon.", "questions_to_authors": "", "paper_topic_and_main_contributions": "The goal of the paper is to understand the learning dynamics of a neural network model. Specifically, the authors propose a theoretical approach to analyze the feature extraction mechanism of neural networks of infinite width using Neural Tangent Kernels. They analyze multi-layer perception, CNNs, linear RNNs, self-attention mechanism, and matrix-vector model and come up with theoretical models to analyze the output of these models. They train models with various activation functions and initialization on SST, Agnews, and IMDB datasets. The findings of their experiments indicate that reducing the initial variance of the final layer's weight can result in a substantial feature bias. This bias makes positive tokens more significant than negative ones in both the MLP and CNN models. The experiments also demonstrate that alternatives to ReLU such as GeLU and SiLU are more robust to changes in initialization compared to ReLU.They demonstrate that co-occurring characters in a given context have a higher model output score compared to non-co-occurring characters. They hypothesize that this may be a part of the reason why large datasets, which provide a large amount of co-occurrence information, helps improve the performance of deep learning models.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "Figure 3 is too small. Consider increasing the size of the plot or at least increasing the font size of the x and y axis labels."}, "extra_text": {}}, {"review_id": "KnxrtYN9pH", "reviewer": "Reviewer_tfVt", "summary": "", "questions": "Having a bias is typically quite useful if the inductive bias is chosen well. Can your theory suggest what kind of biases might be useful? Can we start from engineering a term in the learning dynamics and then reverse-engineer an architecture that will induce that bias?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for your constructive feedback!\n\n1. Promise for future work:\nOur main intention was to provide a theoretical grounding that mirrors the observed learning dynamics across various neural architectures. In our paper, we indeed discussed the potential limitations of the target models in their feature extraction capabilities. As an illustration, our observations confirmed that the one-layer self-attention model did not effectively capture the negation phenomenon, especially concerning negative adjectives. Contrastingly, the matrix-vector model demonstrated superior efficacy in this domain, which suggests that multiplication-based models might inherently be better at n-gram feature extraction (as discussed in Lines 605-609). Given these insights, there's a significant prospect in developing matrix-vector multiplication models tailored for downstream tasks with a strong reliance on n-gram features.\n\n\n2. Discussion about biases:\nSo far as we know, an inductive bias pertains to the inherent assumptions within a model that shape its predictions for data that it has not encountered before. In our work, the inductive bias posits that both positive and negative tokens are anticipated to impartially contribute to the ultimate decision-making process. Therefore, we believe the biases fostering this objective are useful.\nIn the context of our work, the bias exhibited by the MLP with ReLU specifically refers to the second term in Equation 8. A crucial observation we made was that such biases could compromise the fairness of the knowledge that the model extracts. As corroborated by our results in Table 2, the influence of this bias could, in certain scenarios, detrimentally affect performance, underscoring the importance of appropriately chosen biases. \nWhile our current work sheds light on these dynamics, it may offer insights into the roles of biases induced by architecture during gradient descent and help us tune their influences based on their roles. For example, the bias exhibited by the MLP with ReLU may potentially hinder classification performance. We may consider reverse-engineering an architecture devoid of this activation function.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Well-written paper with clean motivation and intuitive explanation of the derivations.\n\n2. Theoretical study the trends of which reflect empirical studies. \n\n3. Could be a useful framework to study novel phenomena in training fundamental models.", "reasons_to_reject": "1. We do not see a phenomenon that has not been observed empirically to be predicted by the theory. Typically when introducing a theory we hope to find insight that could help us build better feature extractors. Could you comment on that or show some promise for future work? ", "questions_to_authors": "Having a bias is typically quite useful if the inductive bias is chosen well. Can your theory suggest what kind of biases might be useful? Can we start from engineering a term in the learning dynamics and then reverse-engineer an architecture that will induce that bias?", "paper_topic_and_main_contributions": "The problem that this paper aims to address is to obtain more understanding of what type of features are learned by neural networks of fundamental models on tokenized data. The paper studies the neural tangent kernel (NTK) at the infinite-width limit for fundamental models (MLP, CNN, self-attention, matrix-vector, and linear RNN models). \n\nThe contributions and findings are several. The learning dynamics of particular tokens' scores depend on a term that is characterized by the co-occurrence of the token with the training data, and another term that does not depend on co-occurrences, and thus represents a bias in the features. Different fundamental models have bias, but notably, some do not, such as self-attention. The choice of variance of the initialization of the final layer, and the choice of activation functions affect the types of features learned. The theory also suggests the reason why ReLU might be replaced with other activations in Transformers, and why increasing the dataset size matters (because of the increased diversity of co-occurrences).\n\nThe theory is in the infinite-width but in practice, for finite width, the trends observed in the theory hold.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "Figure 1 is not referenced early enough and it is not very intuitive/ descriptive. Please improve the caption of the figure."}, "extra_text": {}}, {"review_id": "aL22gPCWjy", "reviewer": "Reviewer_ak69", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": 5, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for your valuable comments and suggestions!\nDue to the strict page limit, we had to prioritize the main content in the body of the paper and present the detailed theories and additional experiments in the appendix. However, we recognize the importance of ensuring a smoother reading experience and will consider restructuring the content for better flow. \n\nThe infinite width condition serves as a mathematical tool to elucidate the learning dynamics more clearly. Prior work, such as that by Sanjeev Arora et al. (2019), delved into the nuances of the infinite width condition. Their findings suggest that when the network width grows sufficiently large, the Neural Tangent Kernel (NTK) can approximate the infinite width NTK effectively while ensuring the error remains within specific bounds. In our empirical studies, we observed that the outcomes from finite network widths mirrored the theoretical insights derived from the infinite width assumption closely. To provide greater clarity on this topic and address potential ambiguities, we will expand on the discussions on the influence of finite and infinite widths in our revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The theory is presented well and covers both binary classification and multi-class classification (appendix). \n- The experiment results delve deep into some interesting behavioral aspects of neural network modules setting the stage for future work. For example, the authors describe the decline in the performance of MLP with ReLU.\n- Learning dynamics of popular neural network modules such as MLP, CNN, and SA have been explained. This sets the stage for more complex modules.\n", "reasons_to_reject": "- The results, experiments, and some theories described in the main paper require the reader to refer to the appendix often. As an example, theory and experiments around MV, and L-RNN are described in the appendix rather than the main paper.\n- This paper is better suited for a longer journal paper rather than a full conference paper. That being said, the authors can figure out if they want to just present a subset of the work for this conference.\n- It is a little unclear to what degree the learning dynamics change without the infinite width condition.\n", "questions_to_authors": "", "paper_topic_and_main_contributions": "\nThe paper describes a novel theoretical perspective on feature extraction mechanisms for fundamental neural modules when using gradient descent in text classification tasks. The authors begin with a theory around the learning dynamics for binary text classifiers. Following, the authors validate the proposed work with experimentation focusing on a deeper understanding of the factors that affect feature extraction and limitations of the models.\n", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper is a solid work and very interesting. By making NTK, the work studies the Feature Extraction Mechanisms of existing models, such as self-attention, cnn, etc.
Decision: Accept-Main