Bidirectional Awareness Induction in Autoregressive Seq2Seq Models

Our method, called Bidirectional Awareness Induction (BAI), is based on the concept of Pivots elements. We define Pivots as elements of the network that can be trained on tasks that are not necessarily correlated to the final objective function but can be beneficial to improve the quality of the result. In this work, we train pivots to reproduce the target output in Seq2Seq problems and they are selected such that the auto-regressive property of the network is preserved. In this way, we induce bidirectional awareness in auto-regressive models or relax the prediction dependency on the previously generated tokens. Ultimately, BAI intends to combine the best of two worlds without inheriting the obstacles of bidirectional models. Overall, BAI can be broken down into three steps: (i) Pivot Selection. Select a set of elements to be trained on the bidirectional task (e.g. encoder features). (ii) Length Equalization. Leverage decoder representations only to equalize the pivot elements to the one of the decoder. (iii) Decoder sequence reconstruction. Reconstruct the decoder sequence using exclusively the result of the previous step. Train the model to optimize the reconstruction discrepancy and the task-specific auto-regressive loss jointly. The encoder features typically represent the most straightforward selection of pivots since they can be leveraged in the learning without breaking the auto-regression condition. However, the concrete implementation of the strategy depends on the adopted architecture, but notably, BAI does not require architectural modifications.

We highlighted in the third step that the BAI loss function is intended to be jointly optimized with the Cross-Entropy training. In this way, the network can benefit from an increased bidirectional awareness, without harming the effectiveness of the traditional approach.

In our experiments, to showcase the flexibility of BAI we present in the following sections the application of our proposal to different architectures, such as the popular and established Transformer [29], the recent ExpansionNet v2 [8], and GPT-2 [21]. To study the robustness of the idea we performed experiments in multiple tasks such as Neural Machine Translation, Text Summarization, and Image Captioning.

In this Section, we propose a concrete implementation of the previous idea in the case of the Transformer [29]. The Transformer is a popular Encoder-Decoder architecture that succeeded in numerous NLP tasks. In essence, all layers are made of Self-Attention and FeedForward layers and their exact implementation is omitted since they do not impact the discussion.

We select the output of the last encoder layer $\overline{E}$=$\{e_{1},e_{2},\ldots,e_{N}\}$ to be the pivots, which consists of $N$ feature vectors of size $H$. During the second step, the length equalization step, we perform the dot product similarity between $\overline{E}$ and the target sequence embeddings $D\in\mathbb{R}^{M\times H}$ where $M$ is the target length. Then apply Softmax and multiply the result by $\overline{E}$:

Note that although it involves the embeddings $D$ the auto-regressive property is preserved. The final BAI loss is defined as the Mean Square Loss (MSE) between the results of Equation 1 and $D$:

where $r_{t}\in\mathbb{R}^{H\times 1}$ denotes the element $t$-th element of $R\in\mathbb{R}^{N\times H}$, and $d_{t}\in\mathbb{R}^{H\times 1}$ defines the $t$-th target embedding. During the training, Equation 2 is jointly trained with the standard Cross-Entropy (CE):

where $X$ and $Y$ denote the input and target sequences respectively. $Embed(\cdot)$ denote the embedding function, and $\lambda$ is a configurable hyper-parameter.

ExpansionNet v2 [8] is another Encoder-Decoder architecture developed for Image Captioning. The encoder implements the (Block) Static Expansion, which distributes the input content across a group of arbitrary numbers of vectors, denoted by $G=\{g_{1},g_{2}\ldots,g_{|G|}\}$, in the Forward Expansion. The result is called "expanded sequence" and is processed again in the Backward Expansion to retrieve the original sequence length. We refer to the original paper for the architectural details.

In the BAI’s Pivot selection step, instead of adopting directly the last encoder outputs, we adopt two intermediate results. Let $A^{g}_{l},B^{g}_{l}\in\mathbb{R}^{g\times H}$ for $g\in G$, the results of the Forward Expansion in the two operation paths, where $g\in G$ denotes the expanded sequence length and $l$ identifies the layer. In particular, we compute $\overline{A}^{g}=\sum_{l}^{L}{A}^{g}_{l}$ and $\overline{B}^{g}_{l}=\sum_{l}^{L}{B}^{g}_{l}$.

In the length equalization step we replicate a parameter-less version of the Backward Expansion:

where $\phi$ denotes a row-wise normalization function [8] and $[\bullet,\bullet,\ldots,\bullet]_{n}$ denotes the concatenation over the $n$-th axis.

The third step of $BAI$ is equivalent to the Equation 2 the case described in Section 3.2. We define the BAI loss as the MSE between the result of Equation 4 and $D$ and perform joint optimization with the Cross-Enropy.

Most popular large pre-trained models are based on the Transformer architecture, such as Flan-T5 [3], mBART [17], and GPT-2 [21]. For this reason, BAI can be applied in the same way as described in Section 3.2 with few exceptions.

In the case of GPT-2, the architecture differs compared to the vanilla Transformer because of the absence of the encoder. In this case, we select the Pivot elements to be the final state of the input sequence.

A total of four datasets are adopted in the experiments. IC is evaluated on the Microsoft’s COCO 2014 [14] which is split according to Karpathy [10]. Overall, it consists of 113K training images, 5K images for the validation, and an additional 5K for the test. Each image is paired with five ground-truth captions. They are pre-processed by lower casing and punctuation removal. Words that occur less than five times are discarded, resulting in a total of 10K unique tokens. NMT is tested on the IWSLT 2015 English-Vietnamese (En-Vi) corpus, which consists of 133K sentences for training and 1268 for evaluation. Sequences whose post-tokenization length is greater than 150 are discarded. In all training instances, the target and the source language vocabulary are shared. Each vocabulary is created using the BPE algorithm [24]. TS is evaluated on the novel TIFU dataset [11], which consists of threads from the "TIFU" subreddit equipped with the "TL;DR" summary. We follow the Pegasus Split [33], the training set consists of $\sim$33K pairs, and 5K are reserved for testing. Additionally, we adopt the popular DialogSum [2], made of 12,460 training and 1500 testing dialogue-summary pairs. Since pre-trained models are adopted for these tasks, no additional pre-processing is applied to the dataset besides the subword tokenization defined in each respective model. The reason we adopted two datasets is motivated by the fact that, in early experiments, some pre-trained baseline models produced only trivial answers in the case of the TIFU dataset. As a result, they were unsuitable for experimentation and comparison. In contrast, all selected models achieved a satisfactory level of performance in the case of DialogSum. In both datasets, training, and testing samples are filtered according to the maximum sequence length supported by the pre-trained model.

The Transformer follows the architecture of the Base Transformer reported in [29] consisting of $N=6$ encoder and decoder layers, dimension $H=512$, FeedForward size of $FF=2048$ and $8$ attention heads. ExpansionNet v2 follows the configuration of [8], it consists of the Swin-Transformer-Large backbone [19] and $N=3$ encoder and decoder layers. We adopt the Static expansion coefficients of $G=\{32,64,128,256,512\}$ in the encoder and a dynamic expansion coefficient of 16 in the decoder. The remaining hyper-parameters are the same as the Transformer. The Transformer will be adopted in the case of NMT and TS. ExpansionNet v2 will be deployed for IC. We refer to the original paper for the architecture details of Flan-T5-Small [3], mBART [17] and GPT-2 [21].

For all three problems, the standard Cross-Entropy (CE) loss is adopted. As described in Equation 3 we jointly train CE and BAI using a parameter $\lambda$ which weights the contribution of the BAI loss. For each task, we report the batch size, batching criteria, optimizer, learning rate strategy, and the hyper-parameters of the function $\Lambda:t\rightarrow[0,1]\subset\mathbb{R}$ which provides the BAI weight $\lambda=\Lambda(t)$ in the iteration $t$. We define $\Lambda$ as follows:

where $T$ denotes the number of iterations in one epoch, $\eta$ denote the minimum weight, $\gamma$ determines the velocity of the weight increase and $\phi$ controls the base of the slope.

IC Training. RAdam [16] optimizer with $\beta_{1}=0.9,\beta_{2}=0.98$, batch size of 48, random batching criteria, an initial learning rate of 2e-4, warmed up for 10000 iterations, then annealed by 0.8 every 2 epochs for 20 epochs. $\Lambda$ is defined by $\eta$=1e-3, $\gamma$=0.5 and $\phi$=15 in the case of ExpansionNet v2. When reported, in the case of the Transformer is configured as $\eta$=1e-7, $\gamma$=2.0 and $\phi$=24.

NMT Training. Adam optimizer [12] with $\beta_{1}=0.9$ and $\beta_{2}=0.98$, batching criteria based on the source sentence length, token batch size of 4096 and the Noam learning rate as described in [29] with 4000 warm-up steps, trained for 300 epochs. The BAI weight function $\Lambda$ is characterized by $\eta$=1e-12, $\gamma$=0.1 and $\phi$=250 for models trained from scratch. Pre-trained models are fine-tuned for 3 epochs and $\phi$ is set to two.

TS Training. The same optimizer and learning rate for NMT is adopted. Trained for 2 epochs in TIFU and 10 epochs in DialogSum. A batch size of 2048 and no warming steps. $\Lambda$ is defined by $\eta$=1e-4, $\gamma$=0.5 and $\phi$=0.5 for Flan-T5-Small. Whereas $\eta$=1e-12, $\gamma$=0.1 and $\phi$=12 are the configurations for GPT-2.

While the general learning rate of the training follows the convention of the original papers, the definition of the $\Lambda$ function was chosen according to empirical rules formulated during the experimental stage. More details can be found in Section 5.3.

In IC, the model is evaluated on the standard evaluation metrics of CIDEr-D , METEOR , SPICE , ROUGE , BLEU . In TS we adopt the BLEU and ROUGE scores. For NMT we use the BLEU score only. Beam Search is adopted during inference, with a beam width of 3 in case of IC, and 4 in the case of NMT and TS. In the case of pre-trained models, greedy decoding is performed.

Tables 1, 2, and 3 report the performance improvements generated by our proposed method against the baselines, in IC, NMT, and TS.

Tables 1 and 2 showcase the performance improvements across standard Image Captioning metrics and Text Summarization. It can be observed that, both Transformer and ExpansionNet v2 benefit from our method, with an increase of 1.2 CIDEr-D in the first, and 2.4 CIDEr-D in the latter. Whereas, Flan-T5-Small and GPT-2 reported an increase of 0.88 and 0.35 ROUGE compared to the respective baselines. These results indicate the robustness and flexibility of our approach to different architectures, in particular, the improvement can be seen regardless of the model and pivot selection. In Table 1 a difference in the improvements of 1.2 CIDEr-D can be observed between the two architectures. This suggests that the method is sensible to the architecture or a proper selection of pivots. However, this characteristic can be regarded as a strength since BAI proved to be effective in the case of the most popular Seq2Seq model, the Transformer, but new architectures and strategies can be developed to benefit more from our approach.

In Table 3 we report the difference in BLEU score for the NMT task. It can be observed that augmenting the Cross-Entropy with BAI is beneficial in both situations, when the model is trained from scratch, represented by the base Transformer, and when it is done during the fine-tuning of a multi-lingual large pre-trained model mBART, when the standard Cross-Entropy is augmented with BAI. This aspect is also supported by Table 2, which showcase the performance improvements also in the case of the pre-trained models Flan-T5-Small and GPT-2.

In Table 3 a significant difference in the improvements can be noted in the pre-trained model compared to the one trained from scratch. We hypothesize this is caused by the fact that the BAI training appears to be the most effective when the weight term is low at the early stages of the training but high at the final epochs. This situation is perfectly recreated by the pre-training and fine-tuning of mBART, where the BAI’s contribution is virtually zero during the pre-training, and it is introduced only when the model reaches the performance plateau. More details can be found in Section 5.3.

While early works on bidirectionality focused on recurrent architectures, in the last years most techniques were developed for the Transformer in different applications. In this regard, we compare our proposal to NAT [6] (NMT), and CBTIC [39] (for IC) as two representatives of the algorithmic-based work on bidirectional and architecture-based approach for bidirectionality respectively. We train our models using the configurations reported in Section 4.2.

For completeness, we first showcase the ineffectiveness of inducing both Left-to-Right (L2R) and Right-to-Left (R2L) knowledge in the model in simple ways, i.e. without architectural or algorithmic incentives. For example, we use the ExpansionNet architecture for the image captioning task and adopt the training configuration described in Section 4.2. In Table 4 we observe that combining L2R and R2L data with simple strategies does not lead to improvements and sometimes can even harm the performances. In contrast, BAI improves the output quality regardless of the L2R or R2L¹¹1Training on R2L data leads to artefacts described in [9] regardless of BAI. approach training with an increase of 3.2 and 1.8 CIDEr-D in the two cases.

In Table 5, we compare our method against the CBTIC architecture. In particular, we re-create the CBTIC environment by training captioning models on Faster-RCNN features as described in [1]. Both the architectures of Transformer and ExpansionNet v2, shortened respectively as "Transf." and "ExpV2", significantly benefit from BAI, confirming its robustness to different visual features. In particular, in the "ExpV2" case, BAI can increase the initial score to 3.4 CIDEr compared to the 2.7 increase observed in the CBTIC architecture. Whereas, in the case of Transformer, BAI led to an increase of 17.8. However, the magnitude of such an improvement in the latter case is an outlier that should not be attributed to BAI alone, but rather an unfavourable experimental training setup for the baseline, since it was designed for ExpansionNet v2.

Compared to CBTIC, BAI improves performance without requiring architectural changes, which is reflected in the constant amount of FLOPs with or without BAI. This aspect can be particularly appreciated in the case of pre-trained models, where architectural modifications can be very time-consuming since they would require the re-training of the foundation model.

In Table 6, we showcase the comparison between our proposal and NAT, in the case of NMT. We re-create the experimental setup of the latter by training from scratch the Transformer architecture presented in Section 4.2 on the IWSLT16 En-De dataset. Here, it can be observed that our method improves the BLEU score by 0.3 without changing the model inference speed whereas, NAT, in the vanilla formulation, focuses on reducing the inference time at the cost of slight degradation of performances.

In this Section, we showcase the impact of the Weight Term choice on the effectiveness of BAI. Although we focused on the case of Image Captioning as the application example for the analysis, similar behaviours were observed in the other selected tasks.

In Figure 2 we plot the weight term function defined in Section 4.2 compared to other hand-crafted functions. The alternative weight functions were designed to represent orthogonal criteria such as constant weight, increasing weight, and decreasing weight.

In Figure 2.a-c) it can be observed that BAI leads to different results according to the weight term function. For instance, a large term that decreases over time ($\Lambda$5) seems to introduce too much noise to the standard Cross-Entropy learning. However, the problem does not lie in the initial magnitude of the term, as can be seen in the contrary case, $\Lambda$4, where the term increases over the epochs and improves the final CIDEr-D score. In general, the magnitude alone does not suggest a predictable pattern, since the cases of $\Lambda$1, $\Lambda$2, and $\Lambda$3 produced mixed results. Overall, the effectiveness of BAI depends on the weight function and it can be beneficial or detrimental depending on the design of the latter. Fortunately, in the case of $\Lambda^{*}$, we observed that the practice of keeping the term small in the early stages and increasing its magnitude when the model reaches the plateau in the Cross-Entropy loss (Figure 2-b) empirically leads to the best result. This motivated the design of the weight term function for all three tasks in Section 5.1.