JANUS : Joint Autoregressive and Non-autoregressive Training with Auxiliary Loss for Sequence Generation (2022)

Abstract: Transformer-based autoregressive and non-autoregressive models have played an essential role in sequence generation tasks. The autoregressive model can obtain excellent performance, while the non-autoregressive model brings fast decoding speed for inference. In this paper, we propose JANUS , a J oint A utoregressive and N on-autoregressive training method using a U xiliary los S to enhance the model performance in both AR and NAR manner simultaneously and effectively alleviate the problem of distribution discrepancy.Further, we pre-train BART with JANUS on a large corpus with minimal cost (16 GPU days) and make the BART-JANUS capable of non-autoregressive generation, demonstrating that our approach can transfer the AR knowledge to NAR. Empirically, we show our approach and BART-JANUS can achieve significant improvement on multiple generation tasks, including machine translation and GLGE benchmarks. Our code is available at Github.

## introduction the transformer-based autoregressive @xcite (ar) generation model has achieved high-quality results in various natural language generation tasks. meanwhile, the nonautoregressive (nar) @xcite @xcite generation methods show great potential to reduce the inference latency by introducing parallel decoding. especially, iterative generative paradigms like cmlm @xcite can dynamically adapt the trade-off between performance and latency. recent works @xcite @xcite have successfully combined these two mechanisms by joint training. however, these approaches only consider the relevance of model parameters, ignoring the correlations between the two manners, which require figure 1 : the example of janus. we exhibit some cases that show the distribution discrepancy due to the difference in context with orange background. efforts for improvement. therefore, in this paper, we attempt to leverage the merits of both ar and nar mechanisms while avoiding their weaknesses to improve their performance. the principal difference between ar and nar is that they use distinct attention mechanisms. in particular, the ar uses unidirectional attention @xcite to simplify the sentence probability distributions by introducing markov hypothesis, which outputs the next token only depending on the previous context. this pattern makes the output distribution of each token accurate and unambiguous but also lacks diversity in inference. the nar introduces bi-directional attention @xcite that can capture bidirectional context by considering the whole sentence information. this pattern needs to predict multiple tokens simultaneously, which causes the token distribution to be ambiguous, stemming from the multi-modality problem @xcite . inspired by knowledge distillation @xcite and deep mutual learning @xcite , we try to regularize the model predictions by minimizing the distribution distance between the output generated by the two manners. however, the conditional output probability of each target position is inconsistent between ar and nar. we present an example in figure 1 to il-lustrate such a discrepancy, where the ar output probability p ar (y 3 |x, y 1 , y 2 ) for token y 3 depends on the context information {x, y 1 , y 2 }, and the nar output p nar (y 3 |x, y 1 , m 2 , y 4 , m 5 , y 6 ) relies on more context {x, y 1 , m 2 , y 4 , m 5 , y 6 }. the p ar (y 3 ) lack the token y 4 , y 6 as condition compare to p nar (y 3 ), in contrast, the p nar (y 3 ) misses the context y 2 . such distinct token distributions bring difficulty in directly regularizing the prediction distances d(ar|nar). to tackle this issue, we introduce an auxiliary distribution p aux to bridge the discrepancy between p ar and p nar with the help of minimizing the distribution distance d(ar|aux) and d(nar|aux) rather than the direct d(ar|nar), as shown in figure 1 . this approach benefits ar and nar to learn from each other with the following advantages: (1) the p aux possesses rich context information compared with ar and nar. for example, comparing to p ar and p nar , p aux (y 3 |x, y 1 , y 2 , y 4 , m 5 , y 6 ) contains all token y 2 , y 4 , y 6 as condition. (2) there is an autoregressive dependency between the predicted tokens, which guarantees the accuracy of output distribution and without ambiguity for each token. (3) the random mask mechanism applied to this distribution, similar to nar, enables the model to learn various token distributions. for effective implementation, we draw on the experience of two-stream self-attention @xcite and position compensation @xcite . we build the p aux by altering the attention matrix of nar, which merely modifies the training procedure without affecting the inference. we first launched experiments on multiple nmt datasets to verify whether janus can help improve performance in two manners. then, we explored our method on existing autoregressive pretraining models, in which we pre-train our bart-janus initialized by bart-base on a large corpus and fine-tuned to adapt various downstream tasks. here we use the glge datasets as our benchmark to validate the model performance. meanwhile, we also conducted comparative ablation studies to illustrate the effectiveness of our proposed method. experimental results show that our method can achieve similar results to the state-of-the-art nar model without distillation data. it simultaneously improves the ar model performance by more than 1.5 bleu scores on average. furthermore, our model exceeds the non-autoregressive pretraining model bang on the same glge tasks. perhaps surprisingly, it can achieve comparable performance with the ar manner at least two times speedup based on the iterative inference mechanism. ## related work how to take advantage of both ar and nar paradigms for sequence generation has been discussed heatedly. @xcite and @xcite introduced knowledge distillation @xcite into the training procedure of nar, which can reduce the data complexity and allow the model to learn the variations from data. some approaches @xcite @xcite try to transfer ar knowledge to nar by applying some specific learning strategy, which can gradually guide the model training. @xcite proposed that the nar model can effectively improve the ar performance by utilizing the nar output with slight speed degradation. several other works @xcite have also proposed to improve the nar performance by combining multiple decoding paradigms into a unified model. the most related work is bang @xcite , which proposed a novel cross-stream visible nstream self-attention structure to bridge the gap between ar and nar generation. specifically, it supports different decoding manners using ar, semi-nar, and nar loss functions as training objectives. unlike it, our approach not only focuses on joint training and parameter sharing but also considers the inner relationship of output distributions. compared with previous work, we can integrate the merits of two different ways while circumventing their drawbacks by introducing an auxiliary loss. some works @xcite @xcite have explored that distribution regularization can make the model more robust. nevertheless, the purpose of our work is more inclined to constrain two distinct training optimization objectives rather than the same. besides, the pretraining models @xcite @xcite can achieve robust generalization but suffer from huge model size. the model with massive parameters is prone to lead the unacceptable inference latency for generation tasks. our approach attempts to leverage an existing autoregressive model @xcite to make it obtain the abilities of non-autoregressive rather than pre-train a nar model from scratch. ## approach in this section, we elaborate on the details of our janus method, which supports ar and cmlmbased nar generation and mutually enhances their performance. as shown in figure 2 , we introduce the two-stream self-attention mechanism to unify these two manners into our single model. ## conclusion in this work, we propose a new training strategy named janus, which supports the ar, nar, and iterative refinement generation mechanisms. meanwhile, we tackle the problem of distribution discrepancy between the ar and nar by introducing an auxiliary distribution. experiments show that janus can significantly improve the ar and nar model performance. further, we pre-train bart-janus and achieve comparable performance with the nar pertained model bang. we are also interested in designing effective finetuning strategies to apply janus, which leaves it as future work.