Getting Hands-On with Keras: Seq2Seq for Automatic Title Generation

玩转Keras之seq2seq自动生成标题 苏剑林 · 2018-09-01 · Machine Learning model NLP Keras text generation

For someone who claims to have been doing NLP for so long, I have to admit I'd never actually run the classic combination of NLP and deep learning — seq2seq. I got in the mood these past couple of days and decided to study and put seq2seq into practice, and of course, naturally, it ends up implemented in Keras.

Seq2seq can be applied to a huge range of tasks. Here I've picked a relatively simple one: generating a title (in Chinese) from the body of an article, which can also be understood as a form of automatic summarization. I chose this task mainly because "article-title" pairs as parallel corpora are relatively easy to find, so I could quickly run some experiments.

A Brief Introduction to Seq2Seq

Seq2seq refers to the general task of converting one sequence into another — machine translation, automatic summarization, and so on. The distinguishing feature of such tasks is that the input sequence and the output sequence are not aligned. If they were aligned, we'd call it sequence labeling instead, which is much simpler than seq2seq. So although sequence labeling can technically also be thought of as a sequence-to-sequence conversion, when we talk about seq2seq we generally don't include sequence labeling under that umbrella.

To implement seq2seq yourself, the key is to understand its principles and architecture. Once you've got that down, it's actually not complicated to implement in any framework. There used to be a third-party Keras seq2seq library, though the author has since abandoned it — probably because something this simple didn't seem worth maintaining as a separate library. Another reference worth reading is last year's official Keras blog post, A ten-minute introduction to sequence-to-sequence learning in Keras. more

Basic Structure

Suppose the original sentence is $X=(a,b,c,d,e,f)$ and the target output is $Y=(P,Q,R,S,T)$. Then a basic seq2seq setup looks like the diagram below.

Basic seq2seq architecture

Although the diagram has a lot of lines and can look dizzying at first, the structure is actually simple. On the left is the encoder for the input, which is responsible for encoding the (possibly variable-length) input into a fixed-size vector. There are many possible choices of model here — RNN structures like GRU or LSTM, CNN+pooling, or Google's pure-attention approach, among others. In theory, this fixed-size vector should contain all the information of the input sentence.

The decoder, meanwhile, is responsible for decoding the vector we just encoded into the output we want. Unlike the encoder, the diagram emphasizes that the decoder is "unidirectionally recursive," because the decoding process proceeds recursively, specifically as follows:

1. Every output sequence begins with a common <start> marker and ends with an <end> marker; these two markers are also treated as words/characters in their own right.
2. Feed <start> into the decoder to obtain a hidden-layer vector; mix this vector with the encoder's output and pass it into a classifier, which should output $P$.
3. Feed $P$ into the decoder to get a new hidden-layer vector, mix it again with the encoder's output, and pass it into the classifier, which should output $Q$.
4. Keep recursing this way until the classifier outputs <end>.

This is the decoding process of a basic seq2seq model: at each step, the result of that step's decoding is fed into the next step, until <end> is output.

The Training Process

As it happens, the diagram above also illustrates the general training process for seq2seq. Since during training we have labeled data pairs, we know in advance what the decoder's input and output should be at every step. So the whole thing effectively boils down to: "given $X$ and $Y_{\text{[:-1]}}$, predict $Y_{\text{[1:]}}$" — that is, we train by shifting the target $Y$ by one position. This training scheme is known as teacher forcing.

The decoder can likewise be built with GRU, LSTM, or CNN structures, but it's worth stressing again that this "knowing the future" property is only available during training — it doesn't exist at prediction time. So at each decoding step, the decoder must not be allowed to peek at inputs from later steps. If you're using an RNN, you generally only use a unidirectional RNN; if you're using a CNN or pure attention, you need to mask out the later portions (for convolution, this means multiplying the convolution kernel by a 0/1 matrix so that convolution can only read the current position and everything to its "left"; for attention it's similar, except you mask the query sequence instead).

Sharp-eyed readers may notice that this training scheme is "local" in some sense — it's not truly end-to-end. For instance, when predicting $R$, we're assuming $Q$ is already known, i.e., that $Q$ was successfully predicted at the previous step, but this can't actually be guaranteed. If some earlier step's prediction goes wrong, it can trigger a chain reaction that renders the training and prediction of all subsequent steps meaningless.
Some researchers have looked into this issue — for example, the paper Sequence-to-Sequence Learning as Beam-Search Optimization folds the entire decoding search process into training, using pure gradient descent (no reinforcement learning needed), which is a very worthwhile approach to learn from. That said, local training is much cheaper computationally, so in practice we generally just use local training for seq2seq.

Beam Search

We've already mentioned the decoding process several times above, but the picture isn't complete yet. In fact, for seq2seq we are modeling

$$p(\boldsymbol{Y}|\boldsymbol{X})=p(Y_1|\boldsymbol{X})p(Y_2|\boldsymbol{X},Y_1)p(Y_3|\boldsymbol{X},Y_1,Y_2)p(Y_4|\boldsymbol{X},Y_1,Y_2,Y_3)p(Y_5|\boldsymbol{X},Y_1,Y_2,Y_3,Y_4)\tag{1}$$

Clearly, during decoding we want to find the $\boldsymbol{Y}$ with the highest probability — but how do we do that?

If, at the first step $p(Y_1|\boldsymbol{X})$, we directly pick the one with the highest probability (hopefully the target $P$), then substitute it into the second step $p(Y_2|\boldsymbol{X},Y_1)$ and again pick the highest-probability $Y_2$, and so on — always choosing the highest-probability output at each step — this is called greedy search, the cheapest decoding scheme available. But note that the result obtained this way isn't necessarily optimal: suppose at the first step we chose $Y_1$, which doesn't have the maximum probability, but plugging it into the second step happens to yield a very large conditional probability $p(Y_2|\boldsymbol{X},Y_1)$, so that the product of the two ends up exceeding what you'd get by taking the maximum at each position individually.

However, if we really tried to enumerate every path to find the optimum, the computational cost would be unacceptably large (this isn't a Markov process, so dynamic programming doesn't apply either). So seq2seq uses a compromise: beam search.

This algorithm resembles dynamic programming, but even in cases where dynamic programming would apply, it's simpler still. The idea is: at each computation step, only keep the $top_k$ best candidate results so far. For example, take $top_k=3$: at the first step, we keep only the top 3 $Y_1$ that maximize $p(Y_1|\boldsymbol{X})$, then substitute each of them into $p(Y_2|\boldsymbol{X},Y_1)$ and again keep the top three $Y_2$ each time. This gives us $3^2=9$ combinations, and we compute the total probability of each combination, keeping only the top three again, recursing this way until the first <end> appears. Clearly, this is still fundamentally a form of greedy search — it's just that the greedy process retains more possibilities along the way. Ordinary greedy search is equivalent to setting $top_k=1$.

Improving Seq2Seq

The seq2seq model described above is the standard version, but it encodes the entire input into a single fixed-size vector and then decodes from that vector alone. This means the vector must, in theory, contain all the information of the original input, which places heavy demands on both the encoder and the decoder — especially for information-preserving tasks like machine translation. This kind of model is a bit like asking someone to "read the Chinese text once and then immediately write out the corresponding English translation," which requires a very strong memory and decoding ability. In reality, ordinary people don't have to do this — we go back and re-read and cross-check the original text repeatedly. This observation motivates the following two techniques.

Attention

Attention has by now essentially become a "standard component" of seq2seq models. The idea is: at each decoding step, we shouldn't just rely on the fixed-size vector encoded by the encoder (a "skim-read" of the whole text) — we should also go back and consult each individual word of the original input (a "close reading" of local details), and combine the two to determine the current step's output.

Seq2seq with attention

As for the specifics of attention, I've written about this before — please see A Brief Read of "Attention is All You Need" (Introduction + Code). Attention generally comes in multiplicative and additive flavors; what I described there is the multiplicative attention introduced by Google's system. Readers can look up additive attention on their own — as long as you grasp the three components of query, key, and value, attention isn't hard to understand in either form.

Prior Knowledge

Coming back to the task of generating article titles with seq2seq, the model can be simplified somewhat, and we can also bring in some prior knowledge. For instance, since both the input and output languages are Chinese, the embedding layers of the encoder and decoder can share parameters (i.e., use the same set of word vectors). This substantially reduces the number of model parameters.

There's also another very useful piece of prior knowledge: most of the words in a title have appeared somewhere in the article (note: merely "appeared" — not necessarily consecutively, and certainly the title isn't a substring contained in the article, otherwise this would just reduce to an ordinary sequence-labeling problem). Given this, we can use the set of words appearing in the article as a prior distribution and incorporate it into the classification model used during decoding, so that the model is biased toward selecting words that already appear in the article when generating output.

Specifically, at each prediction step, we obtain an aggregate vector $\boldsymbol{x}$ (as described earlier, this should be the concatenation of the decoder's current hidden vector, the encoder's encoded vector, and the current attention encoding between decoder and encoder). This is fed into a fully-connected layer, ultimately yielding a vector $\boldsymbol{y}=(y_1,y_2,\dots,y_{|V|})$ of size $|V|$, where $|V|$ is the vocabulary size. After a softmax, $\boldsymbol{y}$ gives the original probability

$$p_i = \frac{e^{y_i}}{\sum\limits_i e^{y_i}}\tag{2}$$

This is the plain classification scheme. To introduce the prior distribution, for each article we construct a 0/1 vector $\boldsymbol{\chi}=(\chi_1,\chi_2,\dots,\chi_{|V|})$ of size $|V|$, where $\chi_i=1$ means the word appeared in the article, and $\chi_i=0$ otherwise. We pass this 0/1 vector through a scale-and-shift layer to obtain:

$$\hat{\boldsymbol{y}}=\boldsymbol{s}\otimes \boldsymbol{\chi} + \boldsymbol{t}=(s_1\chi_1+t_1, s_2\chi_2+t_2, \dots, s_{|V|}\chi_{|V|}+t_{|V|})\tag{3}$$

where $\boldsymbol{s},\boldsymbol{t}$ are trainable parameters. We then average this vector with the original $\boldsymbol{y}$ before applying softmax:

$$\boldsymbol{y}\leftarrow \frac{\boldsymbol{y}+\hat{\boldsymbol{y}}}{2},\quad p_i = \frac{e^{y_i}}{\sum\limits_i e^{y_i}}\tag{4}$$

Experiments show that introducing this prior distribution helps speed up convergence and produces more stable, higher-quality titles.

Keras Implementation

Time for the fun open-source part~

Basic Implementation

Based on the description above, I collected a corpus of over 800,000 news articles to try to train an automatic title-generation model. For simplicity, I used characters as the basic unit and introduced four extra tokens representing mask, unknown, start, and end respectively. For the encoder I used a two-layer bidirectional LSTM, and for the decoder a two-layer unidirectional LSTM. For details, see the source code (Python 2.7 + Keras 2.2.4 + TensorFlow 1.8):

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

Using 64,000 articles per epoch, after training for 50 epochs (a bit over an hour), the model was already producing titles that look reasonably decent:

Article content: On August 28, an online leak claimed that user data from hotels under the Huazhu Group chain had been compromised. According to the content posted by the seller, the data covers guest information from over ten hotel brands under Huazhu, including Hanting, Ni Hao, Orange Hotel, and Ibis. The leaked information reportedly includes Huazhu's official website registration data, ID information recorded at hotel check-in, and room records, including guest names, phone numbers, emails, ID numbers, and login credentials. The seller was offering roughly 500 million records for sale as a package. The third-party security platform "Threat Hunter" verified 30,000 of the records provided by the seller and judged the data highly likely to be genuine. That afternoon, Huazhu Group issued a statement saying it had launched an internal investigation and reported the matter to the police immediately. That evening, Shanghai police confirmed that they had received Huazhu's report and had begun an investigation.
Generated title: "Hotel User Data Suspected of Being Leaked"
Article content: Sina Sports — Beijing time, October 16. At the NBA China Games Guangzhou stop, the Rockets won again, defeating the Nets 95-85. Yao Ming grew steadily into the game, playing 18 minutes 39 seconds, going 5-for-8 from the field, scoring 10 points and grabbing 5 rebounds, plus one block. The Rockets finished their China tour undefeated in both games.
Generated title: "Live Report: Rockets Win Both Games Again, Yao Ming Scores 10 Points and 5 Rebounds at Guangzhou Stop"

Of course these are just two of the better examples — there are also plenty of poor ones, and this certainly isn't good enough to deploy in production as-is; it would need a lot of further "dark magic" tuning.

Masking

Doing masking properly in seq2seq is extremely important. Masking refers to blocking out information that shouldn't be read, or information that's useless, typically by multiplying it away with a 0/1 vector. Keras's built-in masking mechanism is quite unfriendly — some layers don't support masking at all, and ordinary LSTM layers get nearly twice as slow once masking is turned on. So these days I just use 0 directly as the mask marker and write my own Lambda layer to handle the conversion. This way there's essentially no speed penalty, and it can be embedded into arbitrary layers — see the code above for details.

One thing to note: we usually don't distinguish between "mask" and "unknown word" tokens, but with this approach it's better to distinguish them, because even though we don't know the exact meaning of an out-of-vocabulary word, it's still a genuine word — it at least serves as a placeholder — whereas a mask is information we want to erase entirely.

The Decoding Side

The code already implements beam search decoding, and readers are welcome to test how different values of $top_k$ affect the decoding results.

One thing to point out: the reference implementation of decoding in the code is a bit lazy, which drags decoding speed down a lot. In theory, once we obtain the output at the current time step, we only need to feed it into the next iteration of the LSTM to get the output for the next time step. But this requires rewriting the decoder's LSTM (i.e., distinguishing between the training phase and the testing phase while sharing weights between them), which is relatively complex and not very beginner-friendly. So I used a much cruder approach instead: rerun the entire model from scratch at every prediction step. This keeps the amount of code to a minimum, but it gets slower and slower the further along you go — what used to be $\mathcal{O}(n)$ in computational cost becomes $\mathcal{O}(n^2)$.

Final Words

Once again I've gotten an example running in Keras — great, great, I shall continue to firmly hold high the banner of Keras~

The corpus for automatic title generation is relatively easy to find, and among seq2seq tasks it's on the easier side, making it a good exercise task. If you're looking to get into this field, hop right in.