Unraveling Feature Extraction Mechanisms in Neural Networks (2023)

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.

## 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 @xcite @xcite @xcite . 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 @xcite @xcite @xcite @xcite , 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 @xcite , 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 @xcite , allowing us to intuitively interpret the learning dynamics and identify the relevant features cap-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 @xcite , and a matrixvector (mv) model @xcite 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: 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 stateof-the-art large language models and facilitate the development of enhanced model designs foot_0 . infinite-width neural networks researchers found that there could be interesting patterns when the neural network's width approaches infinity. @xcite linked infinitely wide deep networks to gaussian processes. a recent line of work @xcite @xcite 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 @xcite and rnns @xcite . later, @xcite 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. ## 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 classification foot_1 . model description assume we have a training dataset denoted by d, consisting of m labeled instances. let x and y represent all the sentences and labels in the training dataset, respectively. @xmath0, and @xmath1. the vocabulary size is |v |. consider a binary text classification model, where @xmath2, +1}. the model output, denoted as s(t) ∈ r at time t is learning dynamics the evolution of a label score can be described by learning dynamics, which may indicate interesting properties. let f t (x ) ∈ r m represent the concatenation of all the outputs of training instances at time t, and @xmath3. given a test input x ′ , the corresponding label score s ′ (t) follows the dynamics infinite-width when the network width approaches infinity, the ntk will converge and stay constant during training @xcite @xcite . therefore, the learning dynamics can be written as follows, there can be certain interesting properties (regarding the trend of the label scores) harnessed by the interaction yθ ∞ (x ′ , x), where y controls the direction and θ ∞ (x ′ , x) may indicate the relationship between x ′ and x. certain hypotheses can be drawn from these properties. first, the converged ntk θ ∞ (x ′ , x) may intuitively represent the interaction between the test input x ′ and the training instance x. this could extend to the interaction between the basic units (tokens or n-grams) from x ′ and x, as the semantic meaning of an instance can be deconstructed into the combination of the meanings of its basic units @xcite . second, if θ ∞ (x ′ , x) depends on the similarity between x ′ and x, a more deterministic trend can be predicted for a test input x ′ that closely resembles the training instances of a specific type. for example, suppose θ ∞ (x ′ , x) exhibits a significantly large gain when x ′ 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 @xcite ; @xcite , which guarantees the convergence of ntks; 2) obtaining the converged ntk θ ∞ (x ′ , x) and the learning dynamics under the infinite-width condition; 3) performing analysis on the learning dynamics of basic units and revealing possible features. ## 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 @xmath4, l (x) be the instance length, w @xmath5, and @xmath6. @xmath7. w c @xmath8. for simplicity, we let d @xmath9. assume all the parameters are initialized with gaussian distributions in our subsequent analysis, i.e., w ij ∼ n (0, σ 2 w ), w e ij ∼ n (0, σ 2 e ), and v j ∼ n (0, σ 2 v ), and w c ij ∼ n (0, σ 2 w ), for the sake of ntk convergence. ## experiments we conduct experiments to verify our aforementioned analysis in the following aspects: a) verify the features acquired by our models; b) explore factors that may affect feature extraction; c) examine the limitations of those models. ## 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.