Lightweight Transformers for Conversational AI (2022)

Abstract: To understand how training on conversational language impacts performance of pre-trained models on downstream dialogue tasks, we build compact Transformer-based Language Models from scratch on several large corpora of conversational data. We compare the performance and characteristics of these models against BERT and other strong baselines on dialogue probing tasks. Commercial dialogue systems typically require a small footprint and fast execution time, but recent trends are in the other direction, with an ever-increasing number of parameters, resulting in difficulties in model deployment. We focus instead on training fast, lightweight models that excel at natural language understanding (NLU) and can replace existing lower-capacity conversational AI models with similar size and speed. In the process, we develop a simple but unique curriculum-based approach that moves from general-purpose to dialogue-targeted both in terms of data and objective. Our resultant models have around 1/3 the number of parameters of BERT-base and produce better representations for a wide array of intent detection datasets using linear and Mutual-Information probing techniques. Additionally, the models can be easily fine-tuned on a single consumer GPU card and deployed in near real-time production environments.

## introduction the development of the transformer @xcite ) -a multi-headed attention architecture with high capacity -caused a breakthrough in the pre-training of contextualized representations for text @xcite . this architecture can internalize large amounts of information from massive datasets, yielding powerful encoders that can be fine-tuned for various nlp tasks. the generative pre-training (gpt) model @xcite used a standard language modeling objective, learning to predict the next word in a sequence, and demonstrated the ability of transformers to learn long distance dependencies -a limitation of previous architectures. bert @xcite later introduced a masked language model (mlm) objective, where a portion of the text is masked out or perturbed, and the model learns to reconstruct those portions, yielding bi-directional representations. various datasets have been explored for pretraining including the toronto bookcorpus @xcite and wikipedia. more recently, much larger datasets have been used including common crawl datasets for roberta @xcite , t5 @xcite and others. larger datasets facilitate more parameters @xcite and with so much available data, many recent models range from tens to hundreds of billions of parameters. these large language models (llms) exhibit remarkable capabilities for many tasks, but they are massive, difficult to control, and expensive to deploy. 1 llms have yielded improvements over their predecessors, with little architectural modification, primarily by using larger datasets, more capacity, and training longer with more compute. the trend in the literature seems clear -use a denoising or language model objective and train on larger datasets with longer contexts and more capacity to improve nlp performance. the main approach seems to be "more is better" without much qualification to the type of textual content that is applied. there have been some attempts to qualify the corpora used for pre-training @xcite , but attempting to limit what goes into training becomes increasingly difficult as the datasets get larger. for this work, we attempt to target conversational ai, taking into account both model and dataset. it seems reasonable to assume that some carefully curated data, like wikipedia, should be broadly useful as it presents related concepts in close proximity with reliable structure. but for our purposes, it also seems desirable to have a large amount of data that is conversational, though its unclear how this data should be structured and how it should be balanced with non-conversational sources. knowing that task-oriented dialogue systems often have a need to understand domainspecific concepts and proper names, it also seems reasonable to assume that sources such as consumer reviews of products and services might be helpful. in this work, we attempt to build practical, compact models that excel for conversational ai, especially nlu. we set out to build on data sources that will be particularly useful for dialogue systems, and try to incorporate common-sense architectural modifications that should improve performance on dialogue. we focus our effort on mlm models as most nlu tasks benefit from bi-directionality @xcite . we also investigate a curriculum that teaches the model a grounded foundation first in generic language, followed by increasingly complex masking over conversational data. ## model description our model is an encoder-only transformer, similar to bert or roberta, trained primarily on fullcontext conversational input. unlike most previous mlms, our model is trained with relative attention @xcite . in this approach, relative positional representations are not conditioned on the global position of the token but instead use a local relative offset embedding at every layer as part of the self-attention computation, which we hypothesize makes them more suitable for dialogue applications where the global offset in a conversation is not meaningful. additionally, we perform layer normalization before each sub-layer and we also after the last transformer encoder layer @xcite . we also experiment with a curriculum for our mlm where the initial masking follows bert but, later in training, switches to masked turn modeling (mtm), where we mask entire turns, token-bytoken. we train eight-layer models with eight attention heads, and a hidden size of 512. we use byte pair encoding (bpe) @xcite with a vocab containing 30,000 lower-case tokens. we also include special tokens including "[cls]" and "[mask]" (borrowed from bert), "<@xmath0. we use fastbpe foot_1 to train, sampling 2 million posts from reddit as the bpe corpus. ## datasets in our investigation of sources of data for pretraining targeting conversational ai, we identified several potential source types. we bucket these into three basic categories: we hypothesized that, used together, each data source may provide complementary information that could improve model performance and robustness for a range of dialogue tasks. ## training details we train each model using mead-baseline @xcite ) on a single v3 tensor processing unit (tpu). 4 . to best utilize tpus, we use bucketing based on full conversation lengths, scaling the number of samples for each bucket length so that the number of tokens is constant per batch. we use adamw with a peak learning rate of 4e-4, a weight decay of 1e-3, and a linear warm-up of 10,000 steps followed by cosine decay over training to zero. for version 1, we use a maximum context window of length 256, training for 1 million steps with context windows between 64 and 256 tokens. for version 2, targeting the rf conversationonly corpus, we use a longer maximum context window of length 1024. for comparison purposes, in version 2, we train separate models using c4 only, c4 followed by rf, and rf only. from early experiments, we determined that the mtm objective was too difficult to learn from scratch, so for mtm models, we train the first 80% following the masking algorithm for bert, and we switch to mtm masking for the last 20%. ## experiments we use the senteval @xcite approach to perform linear probing on a several 4 each run takes 2-3 days, but curriculum branches are run from previous checkpoints, minimizing the total training time intent detection datasets, including few-shot scenarios. we also perform mutual information-based clustering probing experiments. ## deployment to support deployment into our intelligent virtual assistant (iva) environment, we compared finetuning of our version 1 models against our production models, which operate on asr n-best hypotheses and predict multiple outputs, representing intents and entities. for all fine-tuning, we trained on a single nvidia gtx1080ti and measured the joint accuracy on a real-world customer care application. we found the new models to be competitive, especially in few-shot environments. using distillation to a single model from an ensemfigure 2: oos dataset anmi ## conclusion using a combination of online user conversations and sensible design choices, we are able to provide models that are compact, efficient and perform well for conversational ai. we find that conversationalonly pre-training compares favorably to more traditional online data sources, but that combining the two in a curriculum can be advantageous in many cases. additionally we find that masking whole turns later in training is also particularly helpful for learning good dialogue representations. in future work, we will compare alternative architectures an a larger number of dialogue-oriented tasks. we are also interested in how model capacity affects downstream performance, particularly for few-shot learning. we also wish to explore the trade-offs between lm pre-training in the conversational domain versus equivalent adaptation using a contrastive loss function as used in @xcite , and to further isolate the impacts of individual data sources and lexical content.