Seq2Seq with Bidirectional Decoding

seq2seq之双向解码 苏剑林 · 2019-08-09 · Machine Learning model NLP Keras text generation

In the earlier post Playing with Keras: seq2seq for Automatic Title Generation we already covered the basics of seq2seq, along with a reference Keras implementation.

This post pushes seq2seq one step further by introducing a bidirectional decoding mechanism, which can improve the quality of generated text to some extent (especially for longer sequences). The bidirectional decoding scheme described here is based on Synchronous Bidirectional Neural Machine Translation, and I've once again implemented it in Keras.

Background

Anyone who has studied seq2seq knows that the typical decoding process generates text from left to right, character by character (or token by token): first the encoder output is used to generate the first token; then the encoder output plus the already-generated first token are used to generate the second token; then the encoder output plus the first two tokens are used to generate the third token; and so on. In short, this amounts to modeling the following probability factorization

\begin{equation}p(Y|X)=p(y_1|X)p(y_2|X,y_1)p(y_3|X,y_1,y_2)\cdots\label{eq:p}\end{equation}more

Of course, one could also generate from right to left — generating the last token first, then the second-to-last, then the third-to-last, and so on. The problem is that no matter which direction you decode in, there is a directional bias. For instance, if decoding left to right, the accuracy of the earlier tokens will generally be higher than that of the later ones, and vice versa. Synchronous Bidirectional Neural Machine Translation gives the following statistics from a machine translation task:

$$\begin{array}{|c|c|c|} \hline \text{Model} & \text{The first 4 tokens} & \text{The last 4 tokens}\\ \hline \text{L2R} & 40.21\% & 35.10\%\\ \hline \text{R2L} & 35.67\% & 39.47\%\\ \hline \end{array}$$

$\text{L2R}$ and $\text{R2L}$ refer respectively to left-to-right and right-to-left decoding. From the table we can see that with left-to-right decoding, the accuracy of the first four tokens is around 40%, but the accuracy of the last four tokens is only about 35%; the reverse pattern holds for right-to-left decoding. This reflects the asymmetry inherent in decoding.

To eliminate this asymmetry, Synchronous Bidirectional Neural Machine Translation proposes a bidirectional decoding mechanism that maintains decoders running in both directions and further aligns the generation process using attention.

Bidirectional Decoding

Although this post is based on Synchronous Bidirectional Neural Machine Translation, I haven't read the original paper in full detail — I only skimmed it following my own intuition, roughly grasped the underlying principle, and then implemented the model myself, so my implementation is not guaranteed to match the original exactly. Also, this isn't the first paper to do bidirectional decoding for generation, but it's the first one I came across, so I only implemented this one without comparing it against other related work.

Basic Idea

Since it's called bidirectional "decoding," the changes only concern the decoder, not the encoder, so the description below focuses on the decoder part. Also note that bidirectional decoding is just a strategy, and what follows is just one reference implementation, not a standard or unique one — much like "seq2seq" itself is just a general term for sequence-to-sequence generative models, with plenty of room for variation in how exactly the encoder and decoder are designed.

First, here's a simple animated diagram illustrating the design and interaction process of the bidirectional decoding mechanism:

[

Your browser does not support video
](/usr/uploads/2019/08/3347752814.mp4)

Illustration of the bidirectional decoding mechanism for seq2seq

As shown in the figure, bidirectional decoding can basically be thought of as two decoding modules running in opposite directions coexisting. For convenience, let's call the top one the L2R module and the bottom one the R2L module. At the start, both are fed a start token (S in the figure above), after which the L2R module is responsible for predicting the first token while the R2L module predicts the last token. Next, the first token (along with historical information) is fed into the L2R module to predict the second token — but in addition to the L2R module's own encoding, the encoding already produced by the R2L module is also used. Conversely, the last token (along with historical information) is fed into the R2L module, plus the encoding information already produced by the L2R module, to predict the second-to-last token. This continues until an end token appears (E in the figure above).

Mathematical Description

In other words, when each module predicts each token, it uses not only the information encoded internally by that module, but also the sequence of information already encoded by the other module — and this "use" is implemented via attention. To put it in formulas: suppose in the current state, the L2R module is about to predict the $n$-th token, and the R2L module is about to predict the $n$-th token from the end. Suppose that after several layers of encoding, the resulting L2R vector sequence (corresponding to the second row from the top-left in the figure) is:

\begin{equation}H^{(l2r)}=\left[h_1^{(l2r)},h_2^{(l2r)},\dots,h_n^{(l2r)}\right]\end{equation}

while the R2L vector sequence (corresponding to the second-to-last row at the bottom-left in the figure) is:

\begin{equation}H^{(r2l)}=\left[h_1^{(r2l)},h_2^{(r2l)},\dots,h_n^{(r2l)}\right]\end{equation}

If we were decoding in a single direction, we would use $h_n^{(l2r)}$ as the feature to predict the $n$-th token, or use $h_n^{(r2l)}$ as the feature to predict the $n$-th token from the end.

Under the bidirectional decoding mechanism, we take $h_n^{(l2r)}$ as the query, and $H^{(r2l)}$ as the key and value, to compute an attention output, which is then used as the feature to predict the $n$-th token — this way, when predicting the $n$-th token, the model can already "sense" what comes later in the sequence. Similarly, we take $h_n^{(r2l)}$ as the query and $H^{(l2r)}$ as the key and value to compute another attention output, which is used as the feature to predict the $n$-th token from the end — this way, when predicting the $n$-th token from the end, the model can already "sense" what came earlier in the sequence. In the diagram above, the interaction between the top two layers and the bottom two layers is exactly this attention mechanism. In the code below, we use the simplest form of multiplicative attention (see A Casual Read of "Attention is All You Need" (Introduction + Code)).

Model Implementation

That covers the basic principle and approach of bidirectional decoding. As you can sense, this makes the seq2seq decoder symmetric, which is quite an elegant property. Of course, to fully realize this model, a few more questions need to be addressed: 1. How do we train it? 2. How do we do inference with it?

Training Scheme

Just like ordinary seq2seq, the basic training scheme uses so-called teacher forcing: when predicting the $n$-th token in the L2R direction, we assume the preceding $n-1$ tokens are all known exactly; and when predicting the $n$-th token from the end in the R2L direction, we assume the last $n-1,n-2,\dots,1$ tokens are all known exactly. The final loss is the average of the token-wise cross-entropy in both directions.

However, this training scheme is really a matter of necessity rather than choice, and we'll later analyze the drawback of information leakage that comes with it.

Bidirectional Beam Search

Now let's discuss the inference process.

With conventional unidirectional seq2seq decoding, we use the beam search algorithm to produce a sequence with as high a probability as possible. Beam search decodes token by token, keeping only the top-k "candidate paths" with the highest probability at each step, until an end token appears.

With bidirectional decoding, things get a bit more complicated. We still follow the beam search idea, but now we cache the top-k results for both directions simultaneously — that is, both the L2R and R2L directions each maintain top-k candidate paths. Furthermore, since in bidirectional decoding the L2R decoder needs to reference the R2L decoder's existing decoding results, when predicting the next token we need to enumerate not only the top-k candidate tokens and the top-k L2R candidate paths, but also the top-k R2L candidate paths — giving a total of roughly top-k³ combinations to compute. After computing all these, we adopt the simplest approach: for each "token – L2R candidate path" pair, we average its score over the "R2L candidate path" dimension, reducing the number of scores back down to top-k², used as the score for each "token – L2R candidate path" pair; from these top-k² combinations we then select the top-k with the highest scores. The R2L side undergoes the same processing in reverse. Finally, once both the L2R and R2L directions have produced complete sentences, the one with the highest overall probability (score) is chosen.

We call this whole process "bidirectional beam search." If you're already familiar with unidirectional beam search — or have even implemented it yourself — the process above shouldn't be too hard to follow (looking at the code makes it even easier); it's essentially a natural extension of unidirectional beam search. Of course, if you're not familiar with beam search itself, the description above will probably seem quite confusing. So readers who want to fully understand the principle should start with ordinary unidirectional beam search, make sure they understand that first, then come back to the description above, and finally look at the reference code given below — that should make it much easier to grasp.

Reference Code

Below is my reference implementation of bidirectional decoding. Overall it's consistent with the earlier post Playing with Keras: seq2seq for Automatic Title Generation, except that the decoder side has been changed from unidirectional to bidirectional:

https://github.com/bojone/seq2seq/blob/master/seq2seq_bidecoder.py

Note:
The test environment is roughly the same as before: Python 2.7 + Keras 2.2.4 + TensorFlow 1.8. If you're using Python 3.x or some other environment, feel free to make the necessary changes yourself. If you can't make the changes yourself, please don't come asking me to — I really don't have the time or obligation to get every environment running for everyone. This post is meant to discuss seq2seq techniques, if that's alright with everyone.

In this implementation, I think it's worth explaining how the start and end tokens are handled. In the earlier unidirectional decoding example, I used 2 as the start token and 3 as the end token. With bidirectional decoding, a natural question arises: should the L2R and R2L directions use two separate sets of start/end tokens?

I don't think there's really a standard answer here — whether you share one set or maintain two separate sets of start/end tokens, the results are probably similar. As for the scheme I used in the reference code above, it's a bit unconventional, but I think it's fairly intuitive: I still use just one set of tokens, but in the L2R direction, 2 is the start token and 3 is the end token, while in the R2L direction, 3 is the start token and 2 is the end token.

Discussion and Analysis

Finally, let's think a bit more deeply about this bidirectional decoding scheme. Although making the decoding process symmetric is an elegant property, that doesn't mean it's entirely without issues — thinking about it more deeply helps us understand and use it better.

1. Why does it improve generation?

An interesting question is: bidirectional decoding does seem to improve generation quality at the beginning and end of a sentence, but could it simultaneously degrade quality in the middle?

In theory, this is certainly possible, but in practice it doesn't turn out to be too severe. On one hand, the seq2seq architecture's encoding and decoding capacity is strong enough that information isn't easily lost. On the other hand, when we ourselves evaluate the quality of a sentence, we tend to focus mainly on the beginning and end — if both ends are reasonable and the middle isn't too bad, we tend to judge the sentence as reasonable overall; conversely, if either the beginning or the end is off, we tend to judge the whole sentence as bad. So by improving the generation quality at the beginning and end, the overall perceived generation quality improves as well.

Improvement brought by bidirectional decoding over other unidirectional models, as reported in the original paper

2. It doesn't correspond cleanly to a probabilistic model

For unidirectional decoding, we have a clear probabilistic interpretation: we're estimating the conditional probability $p(Y|X)$ (that is, $\eqref{eq:p}$). But with bidirectional decoding, we find that we simply don't know how to map it onto a probabilistic model at all. In other words, it feels like we're computing a probability, and the results seem to work, yet we don't actually know what it is we're really computing, because the conditional dependency structure has been completely scrambled.

Of course, if it genuinely works in practice, a bit of a shortfall in theoretical elegance isn't a big deal. This point is really just about aesthetic preference in theory, and readers are welcome to have their own opinions on it.

3. Premature information leakage

By information leakage, I mean that the label which is meant to be the prediction target ends up being used as input, causing the training-phase loss to appear artificially low (or the accuracy to appear artificially high).

Because in bidirectional decoding, the L2R side needs to read the vector sequence already encoded by the R2L side, and during training, predicting the $n$-th token on the R2L side requires feeding in the preceding $n-1$ tokens, the further along the decoding goes, the more severe the information leakage becomes. As shown in the figure below:

Illustration of information leakage. During training, when the L2R side is predicting "you," it is in fact making use of the "you" label that was fed into the R2L side; conversely, when the R2L side predicts "north," the same problem occurs — it makes use of the L2R side's "north" label

One observable symptom of information leakage is: later in training, the sum of the L2R and R2L cross-entropies under bidirectional decoding ends up smaller than the cross-entropy of a single-direction model trained on its own — and this is not because bidirectional decoding brings some huge improvement in fit, but rather a manifestation of information leakage.

So if information really is being leaked during training, why does the model still turn out to be useful? I think the rough explanation lies in the table given at the very start of this post. Take the same example: when the L2R side predicts the last token "you," it makes use of all the information already known on the R2L side; and since the R2L side decodes right to left, based on the statistics in the table at the start of the post, we can reasonably expect that, for the R2L side, the prediction accuracy of the very last token (i.e., the first token it decodes) should be the highest. So, assuming R2L really can predict that last token with high accuracy, the "leaked" information stops being a leak in any meaningful sense — because leakage only matters if we've artificially fed in a label that differs from what the model would have predicted anyway; if the predicted result matches the label regardless, then it's no longer really a leak.

Of course, the original paper also offers a strategy to mitigate this leakage problem. Roughly speaking, the approach is: first train a model as described above; then, for each training example, use the model to generate its corresponding predicted output (a pseudo-label); then train the model again, this time feeding in the pseudo-label to predict the correct label. This keeps training and inference as consistent as possible.

Summary

This post introduced and implemented a bidirectional decoding mechanism for seq2seq, which symmetrizes the entire decoding process and thereby improves generation quality to some extent. I think this kind of improvement is worth exploring, especially for readers who appreciate elegance of form. That's why I wanted to write it up.

Beyond that, the post also analyzed some potential issues with this bidirectional decoding scheme and offered my own perspective on them. I welcome further discussion from readers~