A Meta-Learning Perspective on Transformers for Causal Language Modeling (2024)

Abstract: The Transformer architecture has become prominent in developing large causal language models. However, mechanisms to explain its capabilities are not well understood. Focused on the training process, here we establish a meta-learning view of the Transformer architecture when trained for the causal language modeling task, by explicating an inner optimization process that may happen within the Transformer. Further, from within the inner optimization, we discover and theoretically analyze a special characteristic of the norms of learned token representations within Transformer-based causal language models. Our analysis is supported by experiments conducted on pre-trained large language models and real-world data.

## introduction the transformer architecture for neural networks is based on a self-attention mechanism @xcite and has been widely used in natural language processing, computer vision, and scientific discovery @xcite @xcite , among other topical areas. causal language modeling (clm) aims to predict the next element in a sequence in an autoregressive manner and is one of the most important applications of the transformer model. large language models (llms) based on the transformer architecture and clm have shown impressive capabilities in many fields that have been considered difficult challenges in artificial intelligence (ai) @xcite . however, the underlying mechanisms of transformer-based causal language models are still not well understood. meta-learning seeks to learn how to learn by leveraging common knowledge across various tasks, often through a bi-level optimization @xcite @xcite . within this framework, an inner process mirrors a typical model optimization process such as stochastic gradient descent (sgd) @xcite , whereas an outer process is employed to learn hyperparameters for the inner process. studying transformer models from a metalearning perspective is an emerging research direction. @xcite treat the transformer model as a meta-learner to learn parameters for another model but did not identify a possible gradientbased inner optimization process undergone by the forward pass of the transformer model, as we will show in the sequel. some recent works of von @xcite ; @xcite ; von @xcite ; @xcite have attempted to interpret in-context learning in transformer-based language models from a meta-learning perspective by identifying a bi-level optimization process within in-context learning. however, they rely either on simplifications of the original transformer model such as linear attention, or on special constructions of parameters. we make only limited assumptions to analyze a model that is as close to the transformer used in practice as possible. however, we still do make some assumptions including zero initialization and omission of the impacts of normalization components, which makes the subject of our analysis different from transformer models used in practice. more importantly, rather than focusing on incontext learning that is regarded as an emergent capability of transformer-based clm models, we explore whether we can discover a meta-learning process within the training of transformer-based clm models. we attempt to give a plausible bilevel optimization process for a transformer model trained for the clm task. additionally, to investigate the inner optimization process from the metalearning perspective, we visualize the evolution of the token representations within transformerbased clm models. we uncover a characteristic of the evolution of the norms of token representations learned in transformer-based clms, which may indicate a special optimization trajectory. we further theoretically analyze this characteristic using a simplified linear model. various linear simplifications have been used in theoretical analyses of transformer models @xcite . although the linear model is limited in being different from practical settings, it helps us initiate the analysis of this characteristic in an otherwise complicated transformer model. we summarize our contributions as follows: ## transformer layer may approximate optimization steps on clm inspired by a bi-level optimization process discovered for in-context learning, we ask whether a bi-level optimization process also exists in training a transformer for the clm task. let us first formulate a linear projection weight matrix in a transformer model by expanding it via its gradient descent learning dynamics. a weight matrix @xmath0, . . . , x t to optimize a loss function l via backpropagation can be represented as in @xcite : the assumption of zero initialization may be better approximated by larger models. it is a common practice to initialize weights with smaller values for larger dimensions to control the initial variance of the output, thereby fostering a stable learning process and mitigating problems like exploding gradients. numerous initialization methods commonly employed, such as xavier @xcite and kaiming @xcite , incorporate a scaling factor proportional to the inverse of the number of dimensions in either or both the input and output. a transformer model is built from several transformer layers. each layer consists of a multi-head self-attention (mhsa) sublayer and a feedforward network (ffn) sublayer. we investigate them in the following sections. ## a meta-learning view of transformer we generally view the meta-learning process as a "learning to learn" process based on a variety of tasks. the "learning to learn" process is usually formulated as a bi-level optimization process. we have discovered an optimization process inner to the transformer's training process within its forward process, which forms a bi-level optimization process. the language data utilized for the clm task typically spans a broad range of topics, inherently encompassing various tasks. it has long been recognized that a language model trained on such data can proficiently undertake multiple tasks. @xcite . from our perspective illustrated in figure 1 , a transformer layer takes the context tokens as the input data and uses them to form the weightings to approximate the gradients as discussed in section 2. this can be considered as a way of learning by approximating the learning signal, the gradients from close points, different from the classical stochastic gradient descent (sgd), where the gradients are computed from a loss. then, the layer can work as a meta-optimizer to learn to arrange the points and compute the weightings based on current data to reflect expected closeness for efficiently solving a certain task. particularly, these close points could potentially come from different tasks such that ongoing learning could leverage knowledge from other tasks. then, the approximate gradients are used to update the current token representation to be a better one for minimizing the clm loss. particularly, a transformer layer or a meta-optimizer is trained using the clm loss by the backpropagation from an outer optimization process. the token representations can be viewed as model parame-ters from an optimization perspective. through a series of optimization steps performed by layers, the clm loss is expected to decrease over layers. particularly, we view the final layer along with the linear language head as a predictor model that takes a token representation as its parameter vector. so far, we have established a meta-learning view of the transformer trained for the clm task. interestingly, unlike standard gradient descents commonly used in meta-learning, we found that in principle, the transformer could utilize a transformed gradient descent within its inner optimization process, which could be more efficient than a standard gradient descent @xcite . we will leave investigations of these differences to future work. ## optimization trajectory in section 2, we developed an understanding of the forward pass of a transformer-based clm model as an inner optimization process. one further question to ask is whether there are any characteristics of this optimization process. a token representation is continuously updated across layers, which has a similar role to the model parameters being updated by optimization steps. an optimization trajectory can be characterized by the evolution of a token representation. many widely used learning algorithms consist of two parts: optimization and regularization. investigation of a single aspect won't give us a complete picture of the learning algorithms. l2-norm regularization is commonly used with sgd as a part of an optimization process. since the token representation is the subject of optimization, this inspires us to study the property of the norm of the token representations like that of model parameters to give a more comprehensive view of the inner optimization process. from the visualizations and quantitative analysis of the intermediate token representations in section 5.3, surprisingly, we find the norm of a token's vector representation approximately follows an incremental trend across layers in sequence. notice that the loss value has a decreasing trend at the same time across layers in sequence based on section 5.2. this is interesting because it indicates the inner optimization process may follow a specific type of optimization trajectory in general instead of an arbitrary one. in this section, we perform some analysis regarding this characteristic. we propose the following conjecture. conjecture 4.1 a transformer model trained for the causal language modeling task learns a solution such that the norm of the current token vector representation is approximately non-decreasing with increasing layer number l in the transformer model (except for the last layer). we do not consider the last layer because it is a part of the predictor model from our meta-learning view discussed in section 3. this conjecture is supported by both qualitative and quantitative analysis performed in section 5.3. we can also study this characteristic from a vector transformation perspective. due to the complexity of the nonlinear transformation of a transformer layer, we initiate studies on this characteristic by using a simplified linear model. we begin by simplifying the softmaxbased attention mechanism as a linear attention mechanism as done in some prior works @xcite ) such that the mhsa module becomes: by following conventions from (3). we can now treat the mhsa module as a linear transformation using matrix w mhsa . similarly, we obtain a linear version of the ffn module by removing its nonlinear activation function: we develop a linear model corresponding to a transformer layer by omitting the layer normalization and combining everything and residual connections: proposition 1 let @xmath1, and @xmath2. we can have the gram matrix of w decomposed by eigendecomposition as proof 1 since w t w is a symmetric matrix and has non-negative real eigenvalues, we can decompose w t w by eigendecomposition: with proposition 1, we have another conjecture. conjecture 4.2 one of the intrinsic characteristics of w linear in the aforementioned linear model constructed from a transformer model trained for the clm task is to have large enough eigenvalues for its gram matrix such that its output has norm no less than that of its input. we perform experiments to verify conjecture 4.2 in section 5.3 with results shown in ## experiments in this section, we demonstrate experimental results to provide evidence for our theoretical analysis presented in sections 2 and 4. ## conclusion in this work, we show the plausibility of an inner optimization process in a transformer model trained for the clm task by both mathematical derivations and empirical evidence. we also provide a meta-learning view of the transformer model trained for the clm task, which may provide useful insights into the learning dynamics of transformer-based causal language models. in addition, from our experiments on the inner optimization process, we discover a special characteristic about norms of the token representations learned by the transformer-based clm model and investigate this characteristic by theoretical analysis and experiments on real-world large language models. our overarching objective in this study is to provide a new perspective for studying transformer-based clm models as well as some empirical evidence to encourage further research in this direction. the perspective may also give more insights into the design and training of more advanced transformer models by having a clearer view of their internal mechanism. (a) gpt-2 (b) llama-7b (c) llama-13b figure 5: visuzalizations of current token vector representations for different models on the 1bw dataset.