Conformer LLMs - Convolution Augmented Large Language Models

Large Langauge Models have demonstrated significant capabilities across modalities, ushering in a new revolution and frantic interest in artificial intelligence. They have developed super-human abilities in dialogue systems, speech recognition, and image processing. They are built on a core building block of Transformer architectures [1], which have been used not only for LLMs but also for understanding audio [2], speech recognition [3], image understanding [4], text [5], protein sequences [6], robotics [7] to name a few. They operate on a causality assumption, and by using a simple proxy goal of prediction of a text token, they can capture and model the characteristics of the input context and summarize it. Recent advancements in large language models have also trickled down to real-world setups such as reinforcement learning [8], step-by-step reasoning [9], building dialogue agents [10], solving math problems [11] as well as coding [12]. The ability to do one-shot reasoning by choosing a prompt from another modality has given these architectures nothing short of magical properties, as shown in the paper “Socrates architectures” [13]. They also exhibit emergent properties that develop by scaling the number of parameters [14]. They have also been able to learn multi-modal representations by solving proxy tasks of next token prediction [15], which previously would have required a complex pipeline to combine two modalities, e.g., text and audio as shown in [16].

Just before the advent of Transformer architectures, we saw a brief shift from a traditional LSTM-based architecture [17] to modeling sequences to that of one that uses dilated convolution for, e.g., wavenet [18]. They were also explored for natural language processing applications in ByteNet, showing gains over traditional LSTM architectures [19]. This shift to dilated convolutions resulted in several increased gains for the task of time-series prediction in various fields apart from raw audio as described in [20] and natural language processing[5]. However, the main drawback of this shift was that the topology of how the convolutional filters operated is fixed. In contrast, an attention mechanism could allow depending on the dataset, to decide the topology automatically. However, with the advent of Transformers, slowly convolutional-based architectures were phased out, and end-to-end architectures were trained for causal setups such as language modeling and non-causal setups such as in audio [21] or vision[22]. Conformers [23], first proposed for ASR, combined convolutional filters with self-attention and feed-forward layers in Transformers. It gave significant gains for speech recognition by utilizing the best of both worlds. Convolutional augmented image transformers similar in non-causal setups have shown similar gains for understanding images and visual content [22]. Conformers had its origins from CLDNN architectures [24] by replacing LSTM layers with Transformers in the sense that convolutional layers are good feature extractors, LSTMs are good modeling sequences, and MLP layers transform the learned features to more separable space. In this work, we utilize conformers to train large language models in a causal setup. All of our convolutional filters are causal. This architecture thus allows the model to have local and global connections while learning kernels that can filter out/understand dependencies according to the task. One can achieve significant gains by designing hand-made filters in a non-causal setup which was explicitly done by [25]. Our proposed architecture outperforms a purely Transformer based architecture by adding small shallow context convolutional kernels that are causal and which incorporate the best of both worlds. We report results on two standard datasets: For language modeling over text characters and raw audio. For this paper, we do not achieve the state of the art as the number of parameters we use is far smaller than a GPT-2/3. However more importantly, we show experiments showing increasing our systems’ complexity and gains achieved with/without using causal convolutional on every intermediate embedding sandwiched between transformer layers. The contributions of the paper are as follows: i) We showcase gains in using a shallow layer of convolutional filters sandwiched in between two transformer decoder modules and achieve significant gain training of large language models; ii) We show that a Conformer LLM, i.e., a convolution augmented LLM scales well similar to a traditional large language model with embedding size, number of heads and scaling of a convolutional block, which augurs well with integrating these architectures with a decoder only blocks. The paper’s organization is as follows: Section 2 describes the dataset we used for our experiments and explains the conformer architecture. Our experimental setup follows this section and shows the gains we achieve using convolution-augmented language models.

To showcase the strength of our paper, we run experiments for two datasets. The goal in a specific dataset is simple: given the context of previous samples/tokens/characters, predict the next token from the context. For raw waveforms, we make use of the YouTubeMix dataset. It consists of piano sounds at 8KHz, 8-bit quantized signal totaling about 4 hours of audio files. This can be treated as a long-time series of discrete values ranging from 0-255. We use YouTube Mix dataset as reported previously in [26] and [27] with the same training/validation/test splits. In all these experiments, the context is fixed to 256 tokens/samples/events. For the text data, we utilize the text8 dataset [28] that consists of predicting 27 tokens, which are 26 lowercase alphabets and a single token depicting space.

This section describes various blocks of our proposed architecture; concisely, we use causal convolutional filters on the intermediate embeddings after every Transformer decoder layer. However, for the sake of completeness, and this paper, we will explain the entire architecture in detail.

A Transformer decoder architecture is typically used in langauge modeling tasks. It differs from Transformer Encoder in that we use a causal mask so that every layer of intermediate/final transformer embeddings is only a function of the past samples/tokens/embeddings.

We have showcased the powerfulness of a convolutional augmented Transformer for the case of language modeling. We see that by adding a small number of convolutional parameters or, in other words, augmenting the Transformers with convolutional layers, we achieve significant gain in performance. We apply a causal convolutional block to the intermediate embeddings of every Transformer layer that can learn the best of two worlds: Transformers understand dependencies over long time scales, and convolutional filters act on those embeddings to transform them to a more separable space. This, similar to CLDNN or Conformer architectures, brings together three fundamental blocks of neural network advancements: attention, fully connected architecture, and convolutional layers. Through ablation studies, we show how our architecture works for two modalities and achieves gains in natural language processing, raw audio. Our architecture scales with the number of attention heads, parameters of convolutional blocks, and size of the intermediate embeddings, as shown in the experiment. The improved performance of conformers augers well for achieving faster convergence or performance gains with a similar number of parameters as compared to not adding a few parameters of convolutional block. There are several exciting ways of further exploring how we can combine these blocks, and it is an interesting future direction to improve the performance.

This work was supported by the Stanford Institute of Human-Centered AI (HAI) through a Google Cloud computing grant. This paper originated while working on genomic sequences in Kundaje Lab at Stanford University. I want to thank Prof. Kundaje and his group for providing a creative environment that led to this work.