ADAPTIVE LARGE LANGUAGE MODELS BY LAYERWISE ATTENTION SHORTCUTS

Transformer architectures have revolutionized modern AI advancements, powering almost all diverse fields with a unified approach [1]. Modern Transformer language models consist of a simple stack of Transformer decoder blocks iteratively in modalities such as text [2], raw audio waveforms[3], acoustic and music tokens [4, 5, 6], videos [7] to name a few. The flexibility of the architecture to handle any input has generated an enormous interest in the field, and the advent of multi-modal architectures proposed by Google with its Gemini family [8] or multi-modal models like [9] that can allow these models to hear, see and read. However, the information is processed sequentially in all these architectures, one block after another, to allow features to learn representations at multiple levels and depths. In this paper, we explore if we can make these architectures’ depth and context adaptive depending on the contents of the input token. We allow the final layer’s attention mechanism to bypass intermediate layers’ computation and directly contribute to the next token prediction. This would allow more straightforward tokens present in the input, which are easier to predict, to directly learn features in shallow layers to predict the outout. It can thus better utilize and reserve deeper complex self-attention blocks for tougher token predictions. Mainly two ideas inspired this paper – First, several papers utilize intermediate layers that capture information at various scales [10, 11, 12, 13], for solving downstream tasks in natural language such as sentiment analysis, word representations, etc. These intermediate representations are often helpful for capturing useful features of text at multiple scales. Further, such behavior is observed for modalities like speech, where [14] showed it that different layer representations were useful for transformer-based architecture for several downstream tasks related to speech recognition and spoken language understanding. Even for a seq-to-seq architecture, Whisper-AT [15] showed the importance of intermediate layers by using a frozen Whisper-encoder and a temporal transformer to learn feature mapping to classify acoustic sounds. With all of these, we want to state that there is enough evidence to show that intermediate Transformer layers capture meaningful information, capturing different aspects of input signals across modalities. In this paper, we answer –

1. Can we leverage intermediate layer embeddings for improving pre-training of a Large Langauge Model?

The second evidence we address in our architecture is that for certain kinds of inputs, the output of several Transformer layers can be approximated by a single MLP layer, as shown in a recent paper [16]. This is fascinating find as they showed that for specific input tokens, we can bypass mapping of several layers in depth and context to predict the same output as a simple function approximation. Our paper allows this behaviour to be learned during pre-training. This allows our model to avoid wasting complex attention layers when learning trivial feature maps. This is the basis of our paper –

2. Can we make LLM context and depth adaptive with input content to bypass several layers and tokens in context by learning simple feature maps for the final layer? This differs from Pathway’s architecture [17] proposed by Google (Palm) as they have separate streams of MLP and attention blocks, and the overall output is the addition of the output of two blocks. However, the attention mechanism is still confined to the current layer depth. It is not adaptive to allow for a direct bypass of several layers to the output in context and depth. The contributions of the paper are as follows:

• •

  To the best of our knowledge, we, for the first time, show the powerfulness of a decoder-only transformer architecture to attend to all intermediate layer features to improve pre-training performance.

• •

  We show for four different datasets that our method improves the pre-training performance and gives evidence via attention maps that model learns attention shortcuts by attending to the desired depth, depending on the complexity of the current token and its past context.

We utilize four publicly available commonly used datasets for showcasing the powerfulness of our method, with context length of 512. These are chosen from three modalities namely speech, symbolic music and natural language. For speech, we take LibriSpeech corpus [18], which consists of 960 hours of spoken speech for training. We convert it into discrete tokens using ENCODEC [19]. We use the entire training corpus i.e. train-clean-100/360, and train-other-500 as our training data. This gave us a total of 270 million tokens. We utilize the coarsest codebook as proposed in Vall-E [20] to report our results. The model is trained with the vocabulary size of 1024 tokens. For natural language we used two dataset, using two different tokenizer. The first is character level language modelling on text-8 which uses 26 characters and an extra space token as its vocabulary with all other text or symbols removed and replaced with space. This choice removes any biases that may occur due to tokenizers and the dataset itself is widely reported. It has 100M characters from wikipedia. We also use Wiki-103 dataset with GPT-2 tokenzier that has 50,257 tokens. For both of them, we randomly crop 1 million crops with 512 context length yielding a total of 512 million tokens. This is done to have variable context for the same token to make the data more robust. Finally, we use MAESTRO [21], which consists of 1000 piano pieces from classical music in MIDI format. We convert the MIDI track to discrete tokens using Google’s tokenizer yielding 388 sized vocabulary. We compare the performance of our architecture with and without the modifications proposed. The paper is written in an academic setting on shrunk down version of GPT architecture ubiquitously used for language modelling.

In all methods, we use Transformer decoder layers. The architecture for all modalities is the same except for attention mechanism of the final layer. We have already mentioned vocabulary sizes for different modalities in the dataset section. For all the experiments, as a baseline, we use a stack of 10 Transformer decoder layers with a context length of 512, a model dimension of 128, and eight attention heads. The size of the feed-forward dimension was chosen to be four times that of the model dimension. The models were trained for 12-13 epochs, thus seeing more than 10 billion tokens. All architectures were trained with Adam optimizer with starting learning rate 2e-4 and then divide by half mid-way. We take the same architecture with the following modifications for our proposed architecture. The output of each Transformer layer in the second, fourth, and eighth layers is passed to a feature learning module, a 2-layer MLP with 1024 neurons, followed by a dense layer of the size of the model dimension. This allows us to learn feature maps out of embeddings at intermediate layers from depths for all the tokens present in our context. A standard decoder only Transformer formulation is,

$$x^{l+1}=x^{l}+\text{MLP}(\text{LayerNorm}(x^{l}+\text{Attention}(\text{LayerNorm}(x^{l}))))$$

where $x^{l+1}$, is the output of the $l^{th}$ Transformer decoder layer. Our paper introduces an MLP layer where we learn features directly from $x^{l}$, and we call it $x_{feat}^{l}$ which is, $x_{feat}^{l}=\text{MLP}(x^{l})$. Now, we keep the Transformer decoder blocks similar to the baseline architecture except for the final layer decoder block where we allow features with Attention shortcuts: We allow the final prediction at every token to bypass intermediate computations both in depth and context as it deems fit. It allows to directly attend to the desired feature of interest depending on complexity of input tokens and its past context. Further, as shown in the Figure 1, we want the computations to be adaptive, i.e., can we bypass in-depth and context? This follows [16], which found that a set of layers and context can be bypassed with simple features for easier input tokens and context. We focused on addressing these two insights when designing our architecture. In the ablation studies discussed in Section 4.2, we experimented with i) using at a particular depth $l$, learned features with MLP layer $x_{feat}^{l}$ ii) Simply using the $x^{l}$ to learn attention shortcuts. How can we learn computation shortcuts? Since most of the Transformer based LLMs are simply a stack of Transformer Decoder layers, we allow the architecture to use the intermediate layers. In addition to learning computational attention shortcuts, the other added benefit is that, as we have seen in the introduction, the intermediate layer has such rich information about the input data, which might be more useful than just learning features in a sequential hierarchy. Further, by giving the network the flexibility to attend to the intermediate layer feature maps directly, we do not waste precious resources which are powerful attention blocks and layers just learning simple linear maps as seen in [16].

By allowing the final decoder layer to directly attend to the intermediate features of the tokens at various depths and contexts, as we can see in Figure 1, we address the points that we have described above. We replace the self-attention block of the last layer of the stack of the Transformer decoder with a cross-attention layer. The attention mask has to be causal. Hence we only allow for the final Transformer block to attend to all of the previous current layer inputs $x^{l}(t)$, ($l$ being 9 in our case), and all features $x_{feat}^{2}(k)$, $x_{feat}^{4}(k)$, $x_{feat}^{6}(k)$, $x_{feat}^{8}(k)$ with $k$ going from $1,2,3,t$, and $L$ being context length as 512. We introduce causal masks so that at any instance $t\leq L$, $x^{9}(t)$ can attend to all features derived from intermediate layers $x_{feat}^{l}(k)$ i.e. $x_{feat}^{2}(k)$, $x_{feat}^{4}(k)$, $x_{feat}^{6}(k)$, $x_{feat}^{8}(k)$, with $k$ going from $1,2,3,...t$. This formulation of cross attention is similar to perceiver AR [22], Figure 1, using it for LLM like setups. We learn the value and the key matrices are learned from the concatenated stack of $x_{feat}^{2}(k)$, $x_{feat}^{4}(k)$, $x_{feat}^{6}(k)$, $x_{feat}^{8}(k)$, and the value matrices are learned from $x^{l}(k)$, with proper masking to ensure we do not attend to the future tokens. Since there are now 4 times more tokens to attend to in the cross attention layer, we have twice the number of heads just in the final cross attention layer to account for that. Our model has the same structure in all the layers as a vanilla transformer architecture except for using cross-attention in the final layer before softmax classification instead of a self-attention layer. Thus, the final layer decoder output $x^{10}$ with stacked features $x_{feat}=[x_{feat}^{2},x_{feat}^{4},x_{feat}^{6},x_{feat}^{8}],$ is

	$\displaystyle x^{10}$	$\displaystyle=x^{9}+\text{MLP}(\text{LayerNorm}(x^{9}+x_{att}))$
	$\displaystyle x_{att}$	$\displaystyle=\text{CrossAttention}(\text{LayerNorm}(x^{9}),\text{LayerNorm}(x_{feat}))$

We wanted to showcase the power of our architecture across four different datasets and modalities. As described in the previous section, we use text-8, Wiki-103, LibriSpeech, and MAESTRO. We compare the performance of our architecture with that of the baseline only in terms of negative log-likelihood scores. This is in line with other papers in natural language processing [23, 24] that report the performance of architecture only in terms of advancement of how well it does in pretraining in likelihood scores. Neural architectures scale well with the size of the architecture [25]. In addition, keeping all the aspects of the training recipe the same, i.e., training corpus, optimizers, and post-training like DPO or RLHF, the reason for the improved performance of LLMs is their ability to predict the next token correctly. Large Language Model-based generative architectures utilize sampling from the distribution of the next token to solve generative tasks, e.g., music generation, text generation, etc. Large architecture performs better because fewer errors in predicting the next token are present due to the probability distribution closer to the desired distribution. Further, the importance of likelihood-based metrics can also judged by the fact that in latest advanced models, like Gemini, these scores were conspicuously absent and were left blank from the results and plots. Hence, with this background, we will begin discussing our results.

We have showcased a method that improves large language model pre-training on small academic-scale architectures by allowing the neural architecture to attend to the intermediate feature representations. We give evidence that the attention map attends to all of the intermediate representations depending on the contents of the input signal. This allows the architecture to develop attention shortcuts as it deems fit to directly attend to utilize shallow representations if necessary, thereby freeing the subsequent layers to carry out heavier computations. Finally, exploring several possible variants of our model will be interesting, for example, combining it with pathways and other kinds of efficient neural architecture.