Having Your Cake and Eating It Too: SimBERT, a Model That Combines Retrieval and Generation
A while back, we released a model called SimBERT, which is built on top of Google's open-sourced BERT and fine-tuned using a task design inspired by Microsoft's UniLM, integrating both retrieval and generation into a single model. As a result, it is capable of both generating similar questions and retrieving similar sentences. At the time, however, we only released the model weights and an example script, without explaining the underlying principles or the training process. In this post, we'll fill in that gap.
Open-source repository: https://github.com/ZhuiyiTechnology/simbert
UniLM
UniLM is a Transformer model that combines both NLU and NLG capabilities. It was proposed by Microsoft in May of last year, and upgraded to v2 this February. We already gave a brief introduction to UniLM in an earlier post, From Language Models to Seq2Seq: Transformers Are Like Theater, It's All About the Mask, and it has since been integrated into bert4keras.
The core idea behind UniLM is to endow the model with Seq2Seq capability through a special attention mask. Suppose the input is "What do you want to eat" and the target sentence is "White cut chicken." UniLM concatenates the two sentences into one: [CLS] What do you want to eat [SEP] White cut chicken [SEP], and then applies the attention mask shown in the figure below:
UniLM's Maskmore
In other words, the tokens [CLS] What do you want to eat [SEP] attend to each other bidirectionally, while the tokens White cut chicken [SEP] use unidirectional attention, allowing the model to recursively predict the tokens White cut chicken [SEP]. This is what gives it text generation capability.
Illustration of UniLM as a Seq2Seq model. The input segment can attend bidirectionally within itself, while the output segment only attends unidirectionally.
Seq2Seq alone only demonstrates that UniLM has NLG capability, so why did we say earlier that it has both NLU and NLG capability at the same time? Because of UniLM's special attention mask, the 6 tokens [CLS] What do you want to eat [SEP] only attend to each other, and have nothing to do with White cut chicken [SEP]. This means that even though White cut chicken [SEP] is appended afterward, it has no effect on the encoded vectors of the first 6 tokens. To put it more plainly: the encoded vectors of the first 6 tokens are exactly equivalent to what they would be if only [CLS] What do you want to eat [SEP] were present. If the vector for [CLS] represents the sentence vector, then it is the sentence vector for "What do you want to eat," not the sentence vector after appending "White cut chicken."
Because of this property, UniLM also randomly inserts some [MASK] tokens on the input side during training, so that the input portion can be used for an MLM task while the output portion is used for a Seq2Seq task. MLM strengthens NLU ability, while Seq2Seq strengthens NLG ability — killing two birds with one stone.
SimBERT
Once you understand UniLM, it's not hard to understand how SimBERT is trained. SimBERT is trained in a supervised manner, using a self-collected corpus of similar sentence pairs. The Seq2Seq component is constructed via a task of generating one sentence from another similar sentence. As mentioned earlier, the [CLS] vector effectively represents the sentence vector of the input, so it can simultaneously be used to train a retrieval task, as shown in the figure below:
Illustration of how SimBERT is trained
Suppose SENT_a and SENT_b form a pair of similar sentences. Then within the same batch, both [CLS] SENT_a [SEP] SENT_b [SEP] and [CLS] SENT_b [SEP] SENT_a [SEP] are included in training, forming a similar-sentence generation task — this is the Seq2Seq part.
On the other hand, we take the [CLS] vectors of the entire batch and form a sentence-vector matrix $\boldsymbol{V}\in\mathbb{R}^{b\times d}$ (where $b$ is the batch size and $d$ is the hidden size), normalize it along the $d$ dimension using $l_2$-normalization to obtain $\tilde{\boldsymbol{V}}$, and then compute pairwise inner products to get a similarity matrix $b\times b$ of shape $\tilde{\boldsymbol{V}}\tilde{\boldsymbol{V}}^{\top}$. We then multiply by a scale factor (we used 30), mask out the diagonal, and apply softmax to each row, training it as a classification task where the target label for each sample is its corresponding similar sentence (its own position along the diagonal having already been masked out). In short, all non-similar samples in the batch are treated as negative samples, and softmax is used to increase the similarity of similar samples while decreasing the similarity of the rest.
At the end of the day, the key point is that "the [CLS] vector effectively represents the sentence vector of the input," which is why it can be used for NLU-related tasks. The final loss is the sum of the Seq2Seq loss and the similar-sentence classification loss.
Other Details
Since the source code has already been released, readers can look into the training details themselves by reading it. The model is implemented with Keras + bert4keras, and the code is fairly clear, so most questions should be answerable just by reading through it.
Demo:
>>> gen_synonyms(u'微信和支付宝哪个好？')
[
u'微信和支付宝，哪个好?',
u'微信和支付宝哪个好',
u'支付宝和微信哪个好',
u'支付宝和微信哪个好啊',
u'微信和支付宝那个好用？',
u'微信和支付宝哪个好用',
u'支付宝和微信那个更好',
u'支付宝和微信哪个好用',
u'微信和支付宝用起来哪个好？',
u'微信和支付宝选哪个好',
u'微信好还是支付宝比较用',
u'微信与支付宝哪个',
u'支付宝和微信哪个好用一点？',
u'支付宝好还是微信',
u'微信支付宝究竟哪个好',
u'支付宝和微信哪个实用性更好',
u'好，支付宝和微信哪个更安全？',
u'微信支付宝哪个好用？有什么区别',
u'微信和支付宝有什么区别？谁比较好用',
u'支付宝和微信哪个好玩'
]
>>> most_similar(u'怎么开初婚未育证明', 20)
[
(u'开初婚未育证明怎么弄？', 0.9728098),
(u'初婚未育情况证明怎么开？', 0.9612292),
(u'到哪里开初婚未育证明？', 0.94987774),
(u'初婚未育证明在哪里开？', 0.9476072),
(u'男方也要开初婚证明吗?', 0.7712214),
(u'初婚证明除了村里开，单位可以开吗？', 0.63224965),
(u'生孩子怎么发', 0.40672967),
(u'是需要您到当地公安局开具变更证明的', 0.39978087),
(u'淘宝开店认证未通过怎么办', 0.39477515),
(u'您好，是需要当地公安局开具的变更证明的', 0.39288986),
(u'没有工作证明，怎么办信用卡', 0.37745982),
(u'未成年小孩还没办身份证怎么买高铁车票', 0.36504325),
(u'烟草证不给办，应该怎么办呢？', 0.35596085),
(u'怎么生孩子', 0.3493368),
(u'怎么开福利彩票站', 0.34158638),
(u'沈阳烟草证怎么办？好办不？', 0.33718678),
(u'男性不孕不育有哪些特征', 0.33530876),
(u'结婚证丢了一本怎么办离婚', 0.33166665),
(u'怎样到地税局开发票？', 0.33079252),
(u'男性不孕不育检查要注意什么？', 0.3274408)
]
Many of you are probably curious about the training data. To address this once and for all: we're not able to make the training data public, nor can we share it privately, so please don't ask about the data. The data was collected by crawling similar questions recommended by Baidu Zhidao (Baidu's Q&A platform), followed by simple algorithmic filtering. If you already have a large collection of questions on hand, you can also use common retrieval algorithms to find similar sentence pairs to use as training data. In general, there are no particularly strict requirements on the training data — in principle, any data with some degree of similarity should work.
As for training hardware, the released model was trained on a single TITAN RTX (22GB of GPU memory, batch_size=128) for about 4 days. There's no hard requirement on GPU memory or training time either — it depends on your actual situation. If you don't have that much memory, just lower the batch size accordingly. And if your corpus isn't very large, you don't need to train for as long (roughly enough to go through the dataset a few full passes).
That's about all I can think of for now — feel free to leave comments if you have further questions.
Summary
This post explained the training principles behind the SimBERT model we released earlier, and we've also open-sourced the training code. SimBERT is trained based on the UniLM idea, and it possesses both retrieval and generation capabilities at once. Feel free to try it out!
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.