Playing with Keras: Automatic Title Generation with seq2seq

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). A basic seq2seq model is shown in the figure below.

Although the diagram has many lines and might look dizzying, the structure is actually quite simple. On the left is the encoder for the input. It is responsible for encoding the input (which may be of variable length) into a fixed-size vector. Many models can be chosen for this, such as RNN structures like GRU or LSTM, or CNN + Pooling, or Google’s pure Attention. This fixed-size vector theoretically contains all the information of the input sentence.

The decoder is responsible for decoding the vector we just encoded into our desired output. Unlike the encoder, we emphasize in the diagram that the decoder is "unidirectional and recursive" because the decoding process is performed step-by-step. The specific process is:

1. All output sequences start with a universal <start> token and end with an <end> token. These two tokens are also treated as words/characters;

2. Input <start> into the decoder to get a hidden layer vector. Mix this vector with the output of the encoder and feed it into a classifier. The classifier should output P;

3. Input P into the decoder to get a new hidden layer vector. Mix it again with the encoder output and feed it into the classifier. The classifier should output Q;

4. Repeat this recursion until the classifier outputs <end>.

This is the decoding process of a basic seq2seq model. During decoding, the result of each step is fed into the next step until the <end> position is reached.

Training Process

In fact, the figure above also illustrates the general training process of seq2seq. Since we have labeled data pairs during training, we know the input and output of the decoder at each step in advance. Therefore, the result is actually "input X and Y_{\text{[:-1]}}, predict Y_{\text{[1:]}}," which means training by shifting the target Y by one position. This training method is called Teacher-Forcing.

The decoder can also use structures like GRU, LSTM, or CNN. However, it is important to emphasize again that this "predicting the future" characteristic only exists during training and not during the prediction phase. Therefore, when the decoder executes each step, it cannot use inputs from future steps in advance. Consequently, if using an RNN structure, generally only unidirectional RNNs are used; if using CNN or pure Attention, the subsequent parts must be masked (for convolution, this means multiplying the kernel by a 0/1 matrix so it only reads the current and "left" positions; for Attention, it is similar, masking the query sequence).

Sensitive readers might notice that this training scheme is "local" and not truly end-to-end. For example, when we predict R, we assume Q is known (i.e., Q was successfully predicted in the previous step), but this cannot be directly guaranteed. If a prediction error occurs at an earlier step, it may cause a chain reaction, making subsequent training and prediction steps meaningless.

Some scholars have considered this issue. For instance, the paper "Sequence-to-Sequence Learning as Beam-Search Optimization" incorporates the entire decoding search process into the training process using pure gradient descent (without reinforcement learning). This is a very worthwhile approach. However, the computational cost of local training is lower, and we generally use local training to train seq2seq models.

Beam Search

We have mentioned the decoding process several times, but it is not yet complete. 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) Clearly, during decoding, we hope to find the \boldsymbol{Y} with the maximum probability. How do we do this?

If in the first step p(Y_1|\boldsymbol{X}), we directly choose the one with the highest probability (which we hope is the target P), then substitute it into the second step p(Y_2|\boldsymbol{X},Y_1), and again choose the highest probability Y_2, and so on, choosing the current maximum probability output at each step, this is called greedy search. It is the lowest-cost decoding scheme. However, note that this scheme may not yield the optimal result. If we choose a Y_1 that is not the most probable in the first step, substituting it into the second step might yield a very large conditional probability p(Y_2|\boldsymbol{X},Y_1), such that the product of the two exceeds the result of the bitwise maximum algorithm.

However, if we were to enumerate all paths to find the optimum, the computational load would be unacceptably large (this is not a Markov process, so dynamic programming cannot be used). Therefore, seq2seq uses a compromise method: beam search.

This algorithm is similar to dynamic programming, but even in problems where dynamic programming is applicable, it is simpler. Its core idea is: In each step of calculation, only keep the top_k best candidate results. For example, if we take top_k=3, in the first step, we only keep the three Y_1 values that maximize p(Y_1|\boldsymbol{X}). Then we substitute them into p(Y_2|\boldsymbol{X},Y_1) respectively and take the top three Y_2 for each. This gives us 3^2=9 combinations. We then calculate the total probability for each combination and again only keep the top three. This continues recursively until the first <end> appears. Obviously, it essentially belongs to the category of greedy search, but it preserves more possibilities during the greedy process. Ordinary greedy search is equivalent to top_k=1.

Attention

Attention is now basically a "standard" module for seq2seq models. Its idea is: at each decoding step, do not only combine the fixed-size vector encoded by the encoder (reading the whole text) but also look back at every original word (reading local details). The two work together to determine the output of the current step.

Regarding the specific implementation of Attention, I have previously written an introduction; please refer to "A Light Reading of ’Attention is All You Need’ (Introduction + Code)". Attention is generally divided into multiplicative and additive types. I introduced the multiplicative Attention introduced by Google. Readers can look up additive Attention themselves. As long as you grasp the three elements—query, key, and value—Attention is not difficult to understand.

Prior Knowledge

Returning to the task of generating article titles with seq2seq, the model can be simplified, and some prior knowledge can be introduced. For example, 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 significantly reduces the number of model parameters.

Additionally, there is a very useful piece of prior knowledge: most words in the title appear in the article (Note: they just appear; they are not necessarily continuous, nor can we say the title is contained within the article, otherwise it would be a standard sequence labeling problem). Thus, we can use the set of words in the article as a prior distribution and add it to the classification model during decoding, making the model more inclined to select words already present in the article.

Specifically, at each prediction step, we obtain a total vector \boldsymbol{x} (as mentioned before, it should be the concatenation of the decoder’s current hidden layer vector, the encoder’s encoding vector, and the Attention encoding between the current decoder and encoder). This is then fed into a fully connected layer to finally obtain a vector \boldsymbol{y}=(y_1,y_2,\dots,y_{|V|}) of size |V|, where |V| is the vocabulary size. After \boldsymbol{y} passes through softmax, we get the original probability: p_i = \frac{e^{y_i}}{\sum\limits_i e^{y_i}} This is the original classification scheme. The scheme introducing the prior distribution is: for each article, we obtain 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. This 0/1 vector is passed through a scaling and shifting layer to get: \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|}) where \boldsymbol{s}, \boldsymbol{t} are training parameters. Then, this vector is averaged with the original \boldsymbol{y} before performing softmax: \boldsymbol{y}\leftarrow \frac{\boldsymbol{y}+\hat{\boldsymbol{y}}}{2},\quad p_i = \frac{e^{y_i}}{\sum\limits_i e^{y_i}} Experiments show that the introduction of this prior distribution helps speed up convergence and generates more stable, higher-quality titles.

Basic Implementation

Based on the description above, I collected a corpus of over 800,000 news articles to attempt training an automatic title generation model. For simplicity, I chose characters as the basic unit and introduced four additional markers representing mask, unk, start, and end. For the encoder, I used a two-layer bidirectional LSTM, and for the decoder, a two-layer unidirectional LSTM. Specific details can be found in 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 as one epoch, after training for 50 epochs (about an hour), the model generated titles that looked reasonably good:

Article Content: On August 28, network leaks claimed that user data from hotel chains under the Huazhu Group were suspected of being leaked. From the content released by the seller, the data includes guest information from more than 10 hotel brands such as Hanting, Joyue, Orange, and Ibis under Huazhu. The leaked information includes Huazhu official website registration data, identity information from hotel check-in registration, and hotel room records, guest names, mobile phone numbers, emails, ID numbers, login account passwords, etc. The seller is selling this package of about 500 million pieces of data. Threat Hunter, a third-party security platform, verified 30,000 pieces of data provided by the seller and believes the data’s authenticity is very high. That afternoon, Huazhu Group issued a statement saying it had quickly launched an internal verification and reported to the police immediately. That evening, Shanghai police news stated they had received a report from Huazhu Group and the police had intervened in the investigation.
Generated Title: "Hotel User Data Suspected of Leakage"

Article Content: Sina Sports News, Beijing time October 16, the NBA China Games Guangzhou station started as scheduled, and the Rockets won again, defeating the Nets 95-85. Yao Ming gradually found his rhythm, playing 18 minutes and 39 seconds, making 5 of 8 shots, scoring 10 points and 5 rebounds, and he also had 1 block. The Rockets successfully concluded their China trip with a record of two wins.
Generated Title: "Direct Hit: Rockets Win Both Games, Yao Ming 10 Points 5 Rebounds in Guangzhou Station"

Of course, these are just two good examples. There are many poor examples, and it is certainly not enough for direct engineering use; many "black tech" optimizations are still needed.

Masking

In seq2seq, doing masking well is very important. Masking means hiding information that should not be read or useless information, generally by multiplying it by a 0/1 vector. Keras’s built-in mask mechanism is quite unfriendly; some layers do not support masking, and the speed of ordinary LSTMs drops by nearly half when masking is enabled. Therefore, I now directly use 0 as the mask marker and write my own Lambda layer for transformation. This way, the speed is basically lossless, and it supports embedding into any layer. Specifics can be found in the code above.

Note that we usually do not distinguish between mask and unk (out-of-vocabulary words). However, if you adopt my scheme, it is better to distinguish them because although we don’t know the specific meaning of an OOV word, it is still a real word and at least serves as a placeholder, whereas mask is information we want to completely erase.

Decoder Side

Beam search decoding has been implemented in the code. Readers can test the impact of different top_k values on the decoding results.

What needs to be mentioned here is that the implementation of decoding in the reference code is somewhat lazy, which significantly slows down the decoding speed. Theoretically, after obtaining the output of the current moment, we only need to pass it to the next iteration of the LSTM to get the output of the next moment. However, this requires rewriting the decoder’s LSTM (i.e., distinguishing between the training phase and the testing phase, with both sharing weights), which is relatively complex and not friendly to beginners. Therefore, I used a very crude scheme: the entire model is rerun for every step of prediction. As a result, the code is minimal, but it gets slower as it progresses, changing the computational complexity from \mathcal{O}(n) to \mathcal{O}(n^2).