Unraveling Feature Extraction Mechanisms in Neural Networks
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 \textit{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 contributions offer insights into the roles and
capacities of fundamental components within large language models, thereby
aiding the broader understanding of these complex systems.
Xiaobing Sun, Jiaxi Li, Wei Lu
# Unraveling Feature Extraction Mechanisms in Neural Networks

###### 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 contributions offer insights into the roles and capacities of fundamental components within large language models, thereby aiding the broader understanding of these complex systems.

## 1 Introduction

Neural networks have become indispensable across a variety of natural language processing (NLP) tasks. There has been growing interest in understanding their successes and interpreting their characteristics. One line of works attempts to identify possible features captured by them for NLP tasks (Li et al., 2016; Linzen et al., 2016; Jacovi et al., 2018; Hewitt and Manning, 2019; Vulic et al., 2020). They mainly develop empirical methods to verify hypotheses regarding the semantic and syntactic features encoded in the output. Such works may result in interesting findings, but those models still remain _black-boxes_ to us. Another line seeks to reveal internal mechanisms of neural models using mathematical tools (Levy and Goldberg, 2014; Saxe et al., 2013; Arora et al., 2018; Bhojanapalli et al., 2020; Merrill et al., 2020; Dong et al., 2021; Tian et al., 2023), which can be more straightforward 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 modules perform better than others on specific tasks, while some exhibit similar behaviors. We may wonder what mechanisms are behind such differences and similarities between those modules. By acquiring deeper insights into the roles of those modules in a complex model with respect to feature extraction, we will be able to select or even design more suitable models for downstream tasks.

In this work, we propose a novel theoretical approach to understanding the mechanisms, through which fundamental models (often used as modules in complex models) acquire features during gradient descent in text classification tasks. The evolution of model output can be described as learning dynamics involving NTKs (Jacot et al., 2018; Arora et al., 2019), which are typically used to study various properties of neural networks, including 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., 2018), allowing us to intuitively interpret the learning dynamics and identify the relevant features cap

Figure 1: Example of co-occurrence features between tokens and labels (self-attention model).
tured by the model.

We applied our approach to several fundamental models, including a multi-layer perceptron (MLP), a convolutional neural network (CNN), a linear Recurrent Neural Network (L-RNN), a self-attention (SA) model (Vaswani et al., 2017), and a 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:

* We propose an approach to theoretically investigate feature extraction mechanisms for fundamental 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 findings and reveal their relevance to advanced architectures such as Transformers (Vaswani et al., 2017).

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

Footnote 1: Our code is available at [https://github.com/richardsun-voyager/ufemnn](https://github.com/richardsun-voyager/ufemnn).

## 2 Related Work

Probing features for NLP modelsProbing linguistic features is an important topic for verifying the interpretability of neural NLP models. Li et al. (2016) employed a visualization approach to detect linguistic features such as negation captured by the hidden states of LSTMs. Linzen et al. (2016) examined the ability of LSTMs to capture syntactic knowledge using number agreement in English subject-verb dependencies. Jacovi et al. (2018) studied whether the CNN models could capture n-gram features. Vulic et al. (2020) presented a systematic analysis to 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 interpretable features for downstream tasks. Our work seeks to explain how neural NLP models capture useful features during training from a theoretical perspective.

Infinite-width Neural NetworksResearchers found that there could be interesting patterns when the neural network'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; Bietti and Mairal, 2019; Nguyen et al., 2021; Loo et al., 2022) proposed that as the network width approaches infinity, the dynamics can be characterized by the NTK, which converges to a kernel determined at initialization and remains constant. This conclusion holds for 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 architecture. Very limited studies have delved into the analysis of feature extraction in neural NLP models. We will investigate the internal mechanisms of neural NLP models under extreme conditions.

## 3 Analysis

We use learning dynamics to describe the updates of neural models during training with the aim of identifying potentially useful properties. For the ease of presentation and discussion, we focus on binary text classification2.

Footnote 2: Analysis for multi-class classification can be found in Appendix A.

Model DescriptionAssume we have a training dataset denoted by \(\mathcal{D}\), consisting of \(m\) labeled instances. Let \(\mathcal{X}\) and \(\mathcal{Y}\) represent all the sentences and labels in the training dataset, respectively. \(x\in\mathcal{X}\) is an instance consisting of a sequence of tokens, and \(y\in\mathcal{Y}\) is the corresponding label. The vocabulary size is \(|V|\). Consider a binary text classification model, where \(y\in\{-1,+1\}\). The model output, denoted as \(s(t)\in\mathbb{R}\) at time \(t\) is

\[s(t)=\mathbf{f}_{t}(x;\mathbf{\theta}_{t}), \tag{1}\]

where \(\mathbf{\theta}_{t}\) (a vector) is the concatenation of all the parameters, which are functions of time \(t\). We refer to the model output \(s(t)\) as the _label score_ at time \(t\). This score is used for classification decisions, _positive_ if \(s(t)>0\) and _negative_ otherwise.

Learning DynamicsThe evolution of a label score can be described by learning dynamics, which may indicate interesting properties. Let \(\mathbf{f}_{t}(\mathcal{X})\in\mathbb{R}^{m}\) represent the concatenation of all the outputs of training instances at time \(t\), and \(y\in\mathcal{Y}\)
is the desired label. Given a test input \(x^{\prime}\), the corresponding label score \(s^{\prime}(t)\) follows the dynamics

\[\begin{split}\dot{s}^{\prime}(t)&=\nabla_{\theta}f_{t }^{\top}(x^{\prime})\nabla_{\theta_{t}}\mathbf{f}_{t}(\mathcal{X})\nabla_{\mathbf{f}_{t }(\mathcal{X})}\mathcal{L}\\ &=\Theta_{t}(x^{\prime},\mathcal{X})\nabla_{\mathbf{f}_{t}(\mathcal{ X})}\mathcal{L},\end{split} \tag{2}\]

where \(\Theta_{t}(x^{\prime},\mathcal{X})\) is the NTK at time \(t\) and \(\mathcal{L}\) is the empirical loss defined as

\[\mathcal{L}=-\frac{1}{m}\sum_{(x,y)\in\mathcal{D}}\log g(ys). \tag{3}\]

where \(g\) is the _sigmoid_ function. For simplicity, we will omit the time stamp \(t\) in our subsequent notations. The dynamics \(\dot{s}^{\prime}\) will obey

\[\dot{s}^{\prime}=\frac{1}{m}\underset{(x,y)\in\mathcal{D}}{\sum}g(-ys^{(x)})y \Theta(x^{\prime},x), \tag{4}\]

where \(s^{(x)}\) is the label score for the training instance \(x\). Obtaining closed-form solutions for the differential equation in Equation 4 is a challenge. We thereby consider an extreme scenario with the infinite network width, suggested by Lee et al. (2018).

Infinite-WidthWhen the network width approaches 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,

\[\dot{s}^{\prime}=\frac{1}{m}\underset{(x,y)\in\mathcal{D}}{\sum}g(-ys^{(x)})y \Theta_{\infty}(x^{\prime},x), \tag{5}\]

where \(\Theta_{\infty}(x^{\prime},x)\) refers to the converged NTK determined at initialization. This convergence may allow us to simplify the representations of the learning dynamics and offer more intuitive insights to analyze its evolution over time.

There can be certain interesting properties (regarding the trend of the label scores) harnessed by the interaction \(y\Theta_{\infty}(x^{\prime},x)\), where \(y\) controls the direction and \(\Theta_{\infty}(x^{\prime},x)\) may indicate the relationship between \(x^{\prime}\) and \(x\). Certain hypotheses can be drawn from these properties. First, the converged NTK \(\Theta_{\infty}(x^{\prime},x)\) may intuitively represent the interaction between the test input \(x^{\prime}\) and the training instance \(x\). This could extend to the interaction between the basic units (tokens or n-grams) from \(x^{\prime}\) and \(x\), 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). Second, if \(\Theta_{\infty}(x^{\prime},x)\) depends on the similarity between \(x^{\prime}\) and \(x\), a more deterministic trend can be predicted for a test input \(x^{\prime}\) that closely resembles the training instances of a specific type. For example, suppose \(\Theta_{\infty}(x^{\prime},x)\) exhibits a significantly large gain when \(x^{\prime}\) is similar to \(x\) at a particular \(y\), and the dynamics will likely receive significant gains in a desired direction during training, thus enabling us to predict the trend of the label score.

We thereby propose the following approach to investigate a target model and verify our aforementioned hypotheses: 1) redefining the target model following the settings proposed by Jacot et al. (2018); Yang and Littwin (2021), which guarantees the convergence of NTKs; 2) obtaining the converged NTK \(\Theta_{\infty}(x^{\prime},x)\) and the learning dynamics under the infinite-width condition; 3) performing analysis on the learning dynamics of basic units and revealing possible features.

## 4 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.

NotationLet \(\mathbf{e}\in\mathbb{R}^{|V|}\) be the one-hot vector for token \(e\), \(l^{(x)}\) be the instance length, \(\mathbf{W}^{e}\in\mathbb{R}^{d_{in}\times|V|}\) be the weight of the embedding layer, and \(\mathbf{v}\in\mathbb{R}^{d_{out}}\) be the final layer weight. \(\mathbf{W}\in\mathbb{R}^{d_{out}\times d_{in}}\) is the weight of the hidden layer in the MLP model. \(\mathbf{W}^{e}_{k}\in\mathbb{R}^{d_{out}\times d_{in}}\) is the kernel weight corresponding to the \(k\)-th token in the sliding window in the CNN model. For simplicity, we let \(d_{out}=d_{in}=d\). Assume all the parameters are initialized with Gaussian distributions in our subsequent analysis, i.e., \(\mathbf{W}_{ij}\sim\mathcal{N}(0,\sigma_{w}^{2})\), \(\mathbf{W}^{e}_{ij}\sim\mathcal{N}(0,\sigma_{e}^{2})\), and \(\mathbf{v}_{j}\sim\mathcal{N}(0,\sigma_{v}^{2})\), and \(\mathbf{W}^{e}_{ij}\sim\mathcal{N}(0,\sigma_{w}^{2})\), for the sake of NTK convergence.

### Mlp

Following Wiegreffe and Pinter (2019), given instance \(x\), the output of MLP is defined as

\[s=\frac{\mathbf{v}^{\top}}{\sqrt{d}}\sum_{j=1}^{l^{(x)}}\mathbf{\phi}(\mathbf{W}\frac{1}{ \sqrt{d}}\mathbf{W}^{e}\mathbf{e}_{j}). \tag{6}\]

The label score \(s\) will be used for making classification decisions. \(\mathbf{\phi}\) is the element-wise _ReLU_ function. \(\mathbf{e}_{j}\) is the _one-hot_ vector for token \(e_{j}\). It is not straightforward to analyze \(s\) 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^{\prime}\) is simply
a token \(e\), we can get the corresponding NTK with the infinite network width.

**Lemma 4.1**.: When \(d\rightarrow\infty\), the NTK between the token \(e\) and instance \(x\) in the MLP model converges to

\[\Theta_{\infty}(e,x)=\rho\sum_{j=1}^{l(x)}\boldsymbol{e}^{\top}\boldsymbol{e}_ {j}+\sum_{j=1}^{l(x)}\mu, \tag{7}\]

where \(\rho=\frac{(\pi-1)\sigma_{e}^{2}\sigma_{w}^{2}}{2\pi}+\frac{\sigma_{e}^{2} \sigma_{w}^{2}+\sigma_{w}^{2}\sigma_{w}^{2}}{2}\) and \(\mu=\frac{\sigma_{e}^{2}\sigma_{w}^{2}}{2\pi}\).

Note that, for two tokens \(e_{j}\) and \(e_{k}\), their one-hot vectors satisfy \(\boldsymbol{e}_{j}^{\top}\boldsymbol{e}_{k}=0\) if \(e_{j}\neq e_{k}\); \(\boldsymbol{e}_{j}^{\top}\boldsymbol{e}_{k}=1\) if \(e_{j}=e_{k}\). The dot-product \(\sum_{j=1}^{l(x)}\boldsymbol{e}^{\top}\boldsymbol{e}_{j}\) can be interpreted as the frequency of \(e\) appearing in instance \(x\).

**Theorem 4.2**.: The learning dynamics of token \(e\)'s label score obey

\[\dot{s}^{e} =\frac{\rho}{m}\sum_{(x,y)\in\mathcal{D}}g(-ys^{(x)})y\omega(e,x) \tag{8}\] \[+\frac{\mu}{m}\sum_{(x,y)\in\mathcal{D}}g(-ys^{(x)})yl^{(x)},\]

where \(\omega(e,x)=\sum_{j=1}^{l(x)}\boldsymbol{e}^{\top}\boldsymbol{e}_{j}\), which depends on the training data and will not change over time.

The _non-linearity_ of the sigmoid function \(g(-ys)\) makes it a challenge to obtain a _closed-form_ solution for the dynamics.

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 \(y\omega(e,x)\) in each training instance. For instance, consider a token that only appears in positive instances, i.e., \(\omega(e,x)>0\) when \(y=+1\); \(\omega(e,x)=0\) when \(y=-1\). In this case, the first term remains positive and incrementally contributes to the label score \(s^{e}\) throughout the training process. The opposite trend occurs for tokens solely appearing in negative instances. If the impact of the second term is minimal, the label scores of these two types of tokens will be significantly positive or negative after sufficient 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 \(\omega(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.

Theorem 4.2 may explain how the MLP model leverages the statistical co-occurrence features between \(e\) and \(y\) as shown in Figure 1, and integrate them in final classification decisions, i.e., tokens solely appearing in positive/negative instances will likely contribute in the direction of predicting a positive/negative label.

### Cnn

We consider the 1-dimensional CNN, with kernel size, stride size, and padding size set to \(K\), 1, and \(K-1\) respectively. For each sliding window \(c_{j}\) comprising \(K\) consecutive tokens, the corresponding feature \(\boldsymbol{c}_{j}\in\mathbb{R}^{d}\) can be represented as

\[\boldsymbol{c}_{j}=\sum_{k=1}^{K}\boldsymbol{W}_{k}^{c}\frac{1}{\sqrt{d}} \boldsymbol{W}^{e}\boldsymbol{e}_{j+k-1}, \tag{9}\]

where \(\boldsymbol{W}_{k}^{c}\) is the kernel weight corresponding to the \(k\)-th token in the sliding window.

The label score of an instance is computed as

\[s=\frac{\boldsymbol{v}^{\top}}{\sqrt{d}}\sum_{j=-(K-1)}^{l(x)}\boldsymbol{ \phi}(\boldsymbol{c}_{j}), \tag{10}\]

where \(-(K-1)\) means the position for the leftmost padding token. The first and last \(K-1\) padding tokens in an instance are represented by zero vectors. \(\boldsymbol{\phi}\) is the element-wise _ReLU_ function. For brevity, we will denote \(\sum_{j=-(K-1)}^{l(x)}\) by \(\sum_{j}\).

Let us focus on a single sliding window and study the learning dynamics of its label score.

**Lemma 4.3**.: Consider a sliding window \(c\) consisting of tokens \(e_{1},e_{2},\ldots,e_{K}\), when \(d\rightarrow\infty\) the NTK between \(c\) and instance \(x\) converges to

\[\Theta_{\infty}(c,x)=\sum_{j}F[\omega_{c}(c,c_{j})]+\]

\[\rho\sum_{k=1}^{K}\sum_{j}H[\omega_{c}(c,c_{j})]\boldsymbol{e}_{k}^{\top} \boldsymbol{e}_{j+k-1}, \tag{11}\]

where

\[\omega_{c}(c,c_{j})\!=\!\!\sum_{k^{\prime}=1}^{K}\!\boldsymbol{e}_{k^{\prime}}^ {\top}\sum_{k=1}^{K}\boldsymbol{e}_{j+k-1},\ \ \rho\!=\!\sigma_{v}^{2}(\sigma_{e}^{2}\!+\!\sigma_{w}^{2}).\]

\(\omega_{c}\) means the number of shared tokens between \(c\) and \(c_{j}\) regardless of positions. \(F\) and \(H^{3}\) are _monotonically-increasing_ and _non-negative_ functions depending on \(\sigma_{e}^{2}\sigma_{w}^{2}\).

The first term in \(\Theta_{\infty}(c,x)\) captures the token similarity between sliding windows \(c\) and \(c_{j}\) regardless of token positions. In the second term, \(\sum_{j}H[\omega_{c}(c,c_{j})]\boldsymbol{e}_{k}^{\top}\boldsymbol{e}_{j+k-1}\) can be viewed as the weighted frequency of token \(e_{k}\) in instance \(x\), and when \(\sigma_{v}\) is sufficiently large, the converged NTK is majorly influenced by the sum of the weighted frequencies of the tokens in \(c\) appearing in \(x\).
**Theorem 4.4**.: The dynamics of the label score of the test sliding window \(c\) obey

\[\dot{s}^{c} =\frac{\rho}{m}\sum_{k=1}^{K}\!\!\sum_{(x,y)\in\mathcal{D}}\!\!g(- ys^{(x)})y\omega(e_{k},x) \tag{12}\] \[+\frac{1}{m}\sum_{(x,y)\in\mathcal{D}}\!\!g(-ys^{(x)})y\sum_{j}F[ \omega_{c}(c,c_{j})],\]

where \(\omega(e_{k},x)=\sum_{j}\!H[\omega_{c}(c,c_{j})]\mathbf{e}_{k}^{\top}\mathbf{e}_{j+k-1}\).

Theorem 4.4 indicates that with a sufficiently large \(\sigma_{v}\), the learning dynamics for window \(c\) may mainly depend on the linear combination of the weighted learning dynamics of its constituent tokens. Similar analysis can be performed on the label score of the sliding window. This may not exactly encode n-grams, which are inherently sensitive to order and can extend beyond their constituent elements. Instead, for each window, it is more akin to the composition model based on vector addition as described in the work of Mitchell and Lapata (2009). The second term in Equation 12 may not be zero even if \(c\) shares no tokens with \(x\), suggesting there can be an induced feature bias similar to the one in the MLP model.

When \(c\) 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-occurrence features between tokens and labels. Importantly, a 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.

### Sa

We employ a fundamental self-attention module, analogous to the component found in Transformers. The representation of the \(i\)-th output in the instance will be computed as a weighted sum of token representations as follows,

\[\mathbf{h}_{i}=\sum_{j=1}^{l^{(x)}}\frac{\alpha_{ij}}{\sqrt{d}}\mathbf{W}^{e}\mathbf{e}_{j}, \tag{13}\]

where \(\alpha_{ij}\) is the weight produced by a _softmax_ function as follows,

\[\alpha_{ij}=\frac{\exp(a_{ij})}{\sum_{j^{\prime}=1}^{l^{(x)}}\exp\bigl{(}a_{ ij^{\prime}}\bigr{)}}. \tag{14}\]

We define the attention score \(a_{ij}\) from position \(i\) to \(j\) as

\[a_{ij}=\frac{(\mathbf{W}^{e}\mathbf{e}_{i}+P_{i})^{\top}(\mathbf{W}^{e}\mathbf{e}_{j}+P_{j})}{ d}, \tag{15}\]

where \(P_{i}\) (\(P_{j}\)) is the positional embedding at position \(i\) (\(j\)) and will be fixed during training. The instance label score will be computed as

\[s=\mathbf{v}^{\top}\sum_{i=1}^{l^{(x)}}\mathbf{h}_{i}=\sum_{i=1}^{l^{(x)}}\sum_{j=1}^ {l^{(x)}}\frac{\alpha_{ij}}{\sqrt{d}}\mathbf{v}^{\top}\mathbf{W}^{e}\mathbf{e}_{j}, \tag{16}\]

which can be viewed as the weighted sum of token-level label scores if we define such a score for each token \(e\) as \(s_{e}=\frac{1}{\sqrt{d}}\mathbf{v}^{\top}\mathbf{W}^{e}\mathbf{e}\). We consider the case where the test input is also simply a token \(e\).

**Lemma 4.5**.: When \(d\rightarrow\infty\), the NTK between the token \(e\) and the instance \(x\) will converge to \(\Theta_{\infty}(e,x)\), which obeys

\[\Theta_{\infty}(e,x)\!\approx\!(\sigma_{e}^{2}\!+\!\sigma_{v}^{2})\sum_{i=1}^{ l^{(x)}}\!\!\sum_{j=1}^{l^{(x)}}\!\!\mathbb{E}(\alpha_{ij})\mathbf{e}^{\top}\mathbf{e}_{j}, \tag{17}\]

where \(\mathbb{E}(\alpha_{ij})\) is the expectation of \(\alpha_{ij}\).

**Theorem 4.6**.: The learning dynamics of the label score of a token \(e\) obey

\[\dot{s}^{e}=\frac{\rho}{m}\sum_{(x,y)\in\mathcal{D}}\!\!g(-ys^{(x)})y\omega(e, x), \tag{18}\]

where \(\omega(e,x)=\sum_{i=1}^{l^{(x)}}\!\sum_{j=1}^{l^{(x)}}\!\mathbb{E}(\alpha_{ij}) \mathbf{e}^{\top}\mathbf{e}_{j}\) and \(\rho=\sigma_{e}^{2}\!+\!\sigma_{v}^{2}\).

Theorem 4.6 shows the learning dynamics of token e's label score also depends on the weighted sum of the frequencies of \(e\) appearing in \(x\). The learning dynamics of a single token's label score will likely resemble that in the MLP model, capturing 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.

### Mv

We consider the matrix-vector representation as applied in adjective-noun composition (Baroni and Zamparelli, 2010) and recursive neural networks (Socher et al., 2012). It models each word pair through matrix-vector multiplication. The label score of an instance is defined as

\[s=\mathbf{v}^{\top}\sum_{j}\frac{1}{d\sqrt{d}}\mathbf{M}(\mathbf{e}_{j})\mathbf{W}^{e}\mathbf{e}_{j +1}, \tag{19}\]

where \(\mathbf{M}(\mathbf{e}_{j})=\text{diag}(\mathbf{W}\mathbf{W}^{e}\mathbf{e}_{j})\) (diag converts a vector into a diagonal matrix.) and \(j=1,2,\ldots,l^{(x)}-1\).

**Lemma 4.7**.: Given a bigram consisting of two tokens \(e_{a}e_{b}\), with the infinite network width the NTK will converge to
\[\Theta_{\infty}(e_{a}e_{b},x)\!=\!(\sigma_{e}^{2}\!+\!3\sigma_{w}^{2})\sigma_{e}^{ 2}\sigma_{w}^{2}\sum_{j}\mathbf{e}_{j}^{\top}\mathbf{e}_{a}\mathbf{e}_{j+1}^{\top}\mathbf{e}_{b}. \tag{20}\]

It is worth highlighting that the interaction \(\mathbf{e}_{j}^{\top}\mathbf{e}_{a}\mathbf{e}_{j+1}^{\top}\mathbf{e}_{b}\) is different from the interaction resulting from the aforementioned models. When \(e_{a}\equiv e_{j}\) and \(e_{b}\equiv e_{j+1}\) (i.e., \(e_{a}e_{b}\equiv e_{j}e_{j+1}\)), the NTK will gain a relatively large value, implying the ability to capture co-occurrence knowledge between bigrams and labels.

**Theorem 4.8**.: The dynamics of the label score of the test bigram \(e_{a}e_{b}\) obey

\[\dot{s}^{ab}=\frac{\rho}{m}\sum_{(x,y)\in\mathcal{D}}\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!
### Feature Extraction

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 1(a), 1(b), and 1(c) that the label scores for tokens in the extracted co-occurring pairs evolve as expected over epochs for the MLP, CNN, and SA models. The label scores of the tokens co-occurring majorly 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 patterns can be observed on IMDB in Appendix C. We also extract bigrams co-occurring majorly with either the positive or negative label from SSTwub 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 3.

Our analysis on the binary classification tasks can be extended to the multi-class classification scenario on Agnews of four-class. The label scores for the tokens associated with a specific class would be assigned relatively large scores in the dimension corresponding to the class as shown in Figures 3(a), 3(b), 3(c), and 3(d). These observations support our analysis of the feature extraction mechanisms within our target models.

In addition, we extend our experiments to language modeling tasks, which can be viewed as a variant of multi-class classification tasks, with the label space equivalent to the vocabulary size. Interestingly, we observe similar token-label patterns on Transformer-based language models incorporating self-attention modules despite their complexity, in both word and character levels. Particularly, we

Figure 3: Distribution of the label scores for extracted bigrams from SSTwub. “p” refers to _positive_ and “n” refers to _negative_.

Figure 2: Evolution of the label scores for the extracted tokens from SST over epochs. “pos token ” and “neg token ” refer to “positive tokens” and “negative tokens” respectively.
find that nanoGPT, a light-weight implementation of GPT, can capture the co-occurrence features between context characters and target characters on the character-level Shakespeare dataset, and reflect them in the label scores as shown in Figure 5. Given a context character, the model's output is more likely to assign higher scores to target characters 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 information between tokens. Further details can be found in Appendix C.

Induced BiasOur approach also indicates that factors such as activation and initial weight variances could affect feature extraction. We downscale the variances for the final layer weight vectors at initialization and compare the learning curves of the extracted tokens' label scores from models with different activations. As can be seen from Figures 1(d) and 1(e), a smaller initialization of the final layer weight variance can lead to a large feature bias, rendering negative tokens less significant than positive ones in the MLP and CNN models. This may not be a desirable situation, as Table 3 suggests a performance decline for the MLP model with _ReLU_. Furthermore, we compare other activation functions such as \(\tanh\), _GeLU_, and _SiLU_, which are alternatives to _ReLU_7. Figures 1(g) and 1(h) show that these alternatives are more robust than _ReLU_ in the MLP model. This also suggests that while non-linear activations may not significantly alter the nature of learned features during training, they can affect the balance of the extracted features. Figure 1(f) 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 1(i), suggesting a possible reason why _ReLU_ was replaced in models such as BERT Devlin et al. (2019), GPT3 Brown et al. (2020), and LLaMA Touvron et al. (2023), despite its presence in the original Transformer architecture Vaswani et al. (2017).

Footnote 7: The discussion of \(\tanh\) and visualization of _SiLU_ can be found in Appendix B and Appendix C, respectively.

Models' LimitationsWe aim to examine whether the CNN and SA models have a limitation in encoding n-grams in situations beyond constituent tokens' semantic meanings. We choose negation phenomena as our testbed, where a negation token can (partially) reverse the meanings of both positive and negative phrases, a task that is challenging to achieve by linear combination. We run experiments on the SSTwsub dataset with labeled sub-phrases, which contains rich negation phenomena, i.e., phrases with their negation expressions 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 5(a) shows that the SA model can capture negation phenomena for positive adjectives

Figure 4: Label scores for extracted tokens from Agnews, a dataset with four classes. SA model. \(d=64\).

Figure 5: Distribution of the label scores for target characters majorly (blue) and rarely co-occurring with each extracted context character. nanoGPT. Shakespeare Dataset.

\begin{table}
\begin{tabular}{c c c c c c c c c c} \hline \hline \multirow{2}{*}{**Dataset**} & \multicolumn{3}{c}{**ReLU**} & \multicolumn{3}{c}{**tanh**} & \multicolumn{3}{c}{**GeLU**} & \multicolumn{3}{c}{**SiLU**} \\  & & I & II & I & II & I & II & I & II \\ \hline \multirow{3}{*}{SST} & valid & 78.4 & 68.0 & 77.2 & 77.3 & 78.0 & 78.3 & 77.8 & 77.8 \\  & test & 80.0 & 67.3 & 78.9 & 78.9 & 79.9 & 79.8 & 79.7 & 79.9 \\ \multirow{3}{*}{Agnews} & valid & 91.1 & 90.4 & 91.1 & 90.7 & 90.8 & 90.0 & 90.7 & 90.1 \\  & test & 91.0 & 90.2 & 90.6 & 90.1 & 90.6 & 89.6 & 90.6 & 89.6 \\ \multirow{3}{*}{IMDB} & valid & 89.5 & 89.6 & 89.8 & 89.7 & 91.5 & 91.7 & 91.5 & 91.8 \\  & test & 89.9 & 89.5 & 89.9 & 90.2 & 91.2 & 91.4 & 91.2 & 91.4 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Average accuracy (%) on SST with scaled variances (I: \(\sigma_{v}=0.1\) and II: \(\sigma_{v}=0.001\)) and different activation functions. 3 trials for each run. MLP model.
but does not perform well for negative adjectives as shown in Figure 5(b). Specifically, prepending a negation word to the negative adjectives does not alleviate their negativity as expected but leads to the contrary. Based on our analysis, the polarity of a negation expression relies largely on the linear combination of the tokens' polarity in the SA model. As both the negation word _not_8 and negative 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); Orvieto 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 Transformer model in Appendix C. Conversely, the MV model demonstrates the efficacy of capturing such negation for negative adjectives, as shown in Figure 5(c), demonstrating that the multiplication mechanism may play a more effective role in composing semantic meanings.

Footnote 8: The negation word _not_ appears more frequently in negative instances (2086 times compared to 813 in positive instances).

### Discussion

Our experimental results verify our theoretical analysis of the feature extraction mechanisms employed by fundamental models during their training process. These findings are consistent even with network widths as small as \(d=64\), a scenario in which the infinite-width hypothesis is not fully realized. This observed pattern underscores the robustness and generalizability of our model, a conclusion that aligns with the insights presented by Arora et al. (2019). They suggest that as network width expands, the NTK closely approximates the computation under infinite width conditions while keeping the error within established bounds. In our study, we noted that both CNN and Self-Attention models predominantly rely on the linear combination of token-label features. However, they exhibit limitations in effectively composing n-grams beyond tokens, a deficiency highlighted in negation cases. This observation points towards a potential necessity for alternative models that are adept at handling tasks involving complex n-gram features. This observation aligns with studies by Bhatamishra 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-vector multiplication, can better capture such negation evident in both analytical and observational perspectives. This model emerges as a promising alternative for tasks that hinge on 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 alternative activation functions to mitigate this bias, enhancing the model's performance and reliability in diverse applications.

## 6 Conclusions

We propose a 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 fundamental models for text classification and explain how these models acquire features during gradient descent. Meanwhile, our approach allows us to reveal significant factors for feature extraction. For example, 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 hypothesis, the patterns observed are remarkable with finite widths. Our future directions include analyzing more complex neural architectures.

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 to the negation operation.
### Limitations

Despite the findings on the aforementioned fundamental models, applying our approach to analyze complex models like Transformers, which incorporate numerous layers, non-linear activation functions, and normalizations, presents challenges due to the increased complexity. These factors contribute to more intricate learning dynamics, making it less straightforward to gain comprehensive insights into the model's behavior. We would like to investigate and formulate them in future directions.

## Acknowledgements

We would like to thank the anonymous reviewers, our meta-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 3 Programme (The Award No.: MOET320200004), and Ministry of Education, Singapore, under its Academic Research Fund (AcRF) Tier 2 Programme (MOE AcRF Tier 2 Award No. : MOE-T2EP20122-0011). Any opinions, findings and conclusions or recommendations expressed in this material are those of the authors and do not reflect the views of the Ministry of Education, Singapore.
