Improving sentence embeddings with BERT and Representation Learning | by Denis Antyukhov | Towards Data Science

In this experiment, we fine-tune a BERT model to improve its capability for encoding short texts. This yields more useful sentence embeddings for downstream NLP tasks.
While a vanilla BERT can be used for encoding sentences, the embeddings generated with it are not robust. As we can see below, the samples deemed similar by the model are often more lexically than semantically related. Small perturbations in input samples result in large changes of predicted similarity.
To improve, we use the Stanford Natural Language Inference dataset which contains sentence pairs manually labeled with entailment, contradiction, and neutral tags. For these sentences we will be learning such a representation, that the similarity between the entailing pairs is greater than the similarity between the contradicting ones.
To evaluate the quality of learned embeddings we measure Spearman rank correlation on STS and SICK-R datasets.
The plan for this experiment is:
preparing the SNLI and MNLI datasetsimplementing the data generatordefining the lossbuilding the modelpreparing the evaluation pipelinetraining the model
preparing the SNLI and MNLI datasets
implementing the data generator
defining the loss
building the model
preparing the evaluation pipeline
training the model
This guide contains code for building and training a sentence encoder on labeled data.
For a familiar reader, it should take around 90 minutes to finish this guide and train the sentence encoder. The code was tested with tensorflow==1.15.
The code for this experiment is available here. This time, most of the code from the previous experiment is reused. I recommend checking it out first.
The standalone version can be found in the repository.
We begin with downloading the SNLI, MNLI, STS and SICK datasets and the pre-trained English BERT model.
Define a loader function for SNLI jsonl format.
To handle the dataset more conveniently, we re-arrange it a little bit. For each unique anchor we create an ID and an entry containing the anchor, entailment and contradiction samples. The anchors lacking at least one sample of each class are filtered out.
A single entry looks as follows
{  'anchor': ["No, don't answer."],    'contradiction': ['Please respond.'],    'entailment': ["Don't respond. ", "Don't say a word. "]}
Finally, load the SNLI and MNLI datasets.
For training the model we will sample triplets, consisting of an anchor, a positive sample and a negative sample. To handle complex batch generation logic we use the following code:
The high-level logic is contained in the generate_batch method.
Batch anchor IDs are selected randomly from all available IDs.Anchor samples are retrieved from anchor samples of their IDs.Positive samples are retrieved from entailment samples of their IDs.Negative samples are retrieved from contradiction samples of their IDs.These may be considered hard negative samples, because they are often semantically similar to their anchors. To reduce overfitting we mix them with random negative samples retrieved from other, random ID.
Batch anchor IDs are selected randomly from all available IDs.
Anchor samples are retrieved from anchor samples of their IDs.
Positive samples are retrieved from entailment samples of their IDs.
Negative samples are retrieved from contradiction samples of their IDs.These may be considered hard negative samples, because they are often semantically similar to their anchors. To reduce overfitting we mix them with random negative samples retrieved from other, random ID.
We can frame the problem of learning a measure of sentence similarity as a ranking problem. Suppose we have a corpus of k paraphrased sentence pairs x and y and want to learn a function that estimates if y is a paraphrase of x or not. For some x we have a single positive sample y and k-1 negative samples y_k. This probability distribution can be written as:
The joint probability of P(x, y) is estimated using a scoring function, S:
During training it is infeasible to sum over all k-1 negative samples in the dataset. Instead, we approximate P(x) by sampling K responses from our corpus for each batch and using them as negative samples. We get:
We will be minimizing the negative log probability of our data. So, for a batch of K triplets for the loss we can write down:
Note: the expression above is known as Softmax-Loss.
Inner product is used as a similarity function S in this experiment. The code for computing the expression in the last brackets is below
First, we import the fine-tuning code from the previous experiment and build the BERT module.
The model has three inputs for the anchor, positive and negative samples. A BERT layer with a mean pooling operation is used as a shared text encoder. Text preprocessing is handled by the encoder layer. Softmax loss is computed on the encoded sentences.
For convenience, 3 models are created: enc_model for encoding sentences, sim_model for computing similarity between sentence pairs and trn_model for training. All models use shared weights.
Natural Language Encoders are usually evaluated by embedding labeled sentence pairs, measuring some kind of similarity between them, then computing the correlation of that similarity with human judgement.
We use STS 2012–2016 and SICK 2014 datasets for evaluating our model. For all sentence pairs in the test set we compute cosine similarity. We report Pearson rank correlation with human-annotated labels.
The Callback below handles the evaluation procedure and saves the provided savemodel to savepath, every time a new best result is achieved.
We train the model for 10 epochs, with 256 batches each. Each batch consists of 256 triplets. We perform an evaluation at the beginning of each epoch.
trn_model.fit_generator(tr_gen._generator, validation_data=ts_gen._generator, steps_per_epoch=256, validation_steps=32, epochs=5, callbacks=callbacks)*** New best: STS_spearman_r = 0.5426*** New best: STS_pearson_r = 0.5481*** New best: SICK_spearman_r = 0.5799*** New best: SICK_pearson_r = 0.6069Epoch 1/10255/256 [============================>.] - ETA: 1s - loss: 0.6858256/256 [==============================] - 535s 2s/step - loss: 0.6844 - val_loss: 0.4366*** New best: STS_spearman_r = 0.7186*** New best: STS_pearson_r = 0.7367*** New best: SICK_spearman_r = 0.7258*** New best: SICK_pearson_r = 0.8098Epoch 2/10255/256 [============================>.] - ETA: 1s - loss: 0.3950256/256 [==============================] - 524s 2s/step - loss: 0.3950 - val_loss: 0.3700*** New best: STS_spearman_r = 0.7337*** New best: STS_pearson_r = 0.7495*** New best: SICK_spearman_r = 0.7444*** New best: SICK_pearson_r = 0.8216...Epoch 9/10255/256 [============================>.] - ETA: 1s - loss: 0.2481256/256 [==============================] - 524s 2s/step - loss: 0.2481 - val_loss: 0.2631*** New best: STS_spearman_r = 0.7536*** New best: STS_pearson_r = 0.7638*** New best: SICK_spearman_r = 0.7623*** New best: SICK_pearson_r = 0.8316Epoch 10/10255/256 [============================>.] - ETA: 1s - loss: 0.2381256/256 [==============================] - 525s 2s/step - loss: 0.2383 - val_loss: 0.2492*** New best: STS_spearman_r = 0.7547*** New best: STS_pearson_r = 0.7648*** New best: SICK_spearman_r = 0.7628*** New best: SICK_pearson_r = 0.8325
For reference we can check the evaluation results from Sentence-BERT paper where the authors evaluated several pre-trained sentence embedding systems on STS and SICK tasks.
Our results with a vanilla mean-pooled BERT model are consistent with the published metrics, scoring 57.99 Spearman rank correlation score on SICK-R.After 10 epochs, the best Colab model achieved 76.94, which is on par with the best result of Universal Sentence Encoder with 76.69.
Since negative examples are shared between all samples in a batch, fine-tuning with larger batch_size tends to improve metrics up to a certain point. Unfreezing more encoder layers also helps, if your GPU can handle it.
Once the training is complete, we can compare the base and trained models by encoding test triplets and examining the predicted similarities side-by-side. Some examples:
Above we have proposed a method to improving sentence embeddings using labeled sentence pairs.
By explicitly training the model to encode the sentence pairs with respect to their semantic relations, we were able to learn more efficient and robust sentence representations.
Both automatic and manual evaluation demonstrate substantial improvement over the baseline sentence representation model.
Pre-training BERT from scratch with cloud TPUBuilding a Search Engine with BERT and TensorflowFine-tuning BERT with Keras and tf.ModuleImproving sentence embeddings with BERT and Representation Learning[you are here]
Pre-training BERT from scratch with cloud TPU
Building a Search Engine with BERT and Tensorflow
Fine-tuning BERT with Keras and tf.Module
Improving sentence embeddings with BERT and Representation Learning[you are here]