MetricBERT: Text Representation Learning via Self-Supervised Triplet Training

Learning textual similarity is a long-standing task with applications in information retrieval [1] , document clustering [2] , essay scoring [3] , and recommender systems [4, 5, 6, 7, 8, 9, 10]. In this work, we propose MetricBERT, a model for text similarity that employs a novel metric learning technique. MetricBERT is a self-supervised model that utilizes pseudo-labels of negative and positive text passages extracted from the dataset at hand. The pseudo-labels serve as a proxy to the ultimate (but absent) similarity labels. While MetricBERT is a general model for text similarity, we chose to focus on text similarity for recommender systems where we show that MetricBERT outperforms the current state-of-the-art by a sizeable margin.

MetricBERT jointly optimizes a masked-language model and a novel margin-based triplet (metric learning) loss. The metric learning strives to embed similar samples closer to each other, compared to dissimilar ones as explained in Sec. 3. Additionally, MetricBERT leverages a hard negative mining procedure [11] to accelerate convergence and improve performance.

In this work, we provide an extensive comparison between different models, focusing on downstream tasks of learning similarities for recommendations. In particular, we report the results of seven methods - LDA, TF-IDF, BERT [12], RoBERTa [13], RecoBERT [4], SBERT [14], and our proposed MetricBERT, evaluated on two datasets for wines and video games recommendations. Notably, we show that the choice of the triplet loss over a traditional contrastive loss is highly beneficial and that optimizing the angular distance, over the alternative of cosine similarity, is crucial for model convergence. Finally, as an additional contribution, we publish a new dataset for text-based similarity of video games along with a test set of similarity annotations crafted by a domain expert.

Recent Transformer-based representation learning methods in NLP rely on training large neural language models (LMs) on self-supervised objectives and massive corpora [12, 13]. These models consists of a two-phase training procedure: First the LM is trained on large unlabeled corpora. Then, the resulting pre-trained model is trained with supervision for a specific downstream task. These models show great promise in various linguistic tasks [12, 15, 16, 13, 17, 18], and usually produce better results than non-contextualized neural embedding methods [19, 20, 21, 22, 23, 24, 25, 26] .

Recent works suggested new approaches for improving the sentence embedding abilities of transformer-based representation learning models: For example, SBERT [14] proposed using distantly-supervised natural language inference (NLI) to create representations that significantly outperform other state-of-the-art sentence embedding methods. In the domain of text-based recommendations, RecoBERT [4] proposed a BERT-based model for text-based item recommendations, by continuing the training of the pre-trained BERT model on the catalog at hand. RecoBERT utilizes BERT’s original pre-training objective of masked language modeling, along with a novel title-description objective that leverages an additional self-supervision task to match between the embeddings of an item title and its description. In this work, we focus on a textual recommendation task and compare (among other models) to both [14] and [4].

Another relevant line of work investigated integrating metric learning techniques for LMs, mostly for the scenario where supervised data is abundant. SPECTER [27] was introduced for producing document-level representations for scientific documents. SPECTER embeds scientific documents by pre-training a Transformer language model using supervision obtained by citation graphs. Due to the lack of similarity labels (or citations) within our domain, we were unable to compare our method with SPECTER. Recently, in [28] the authors surveyed contrastive learning concepts for NLP applications and their relations to other fields, and concludes with an open call to bring more contrastive methods to NLP.

Let $D:=\left\{(a_{i},b_{i})^{N}_{i=1}\right\}$ be a dataset of $N$ items, where each item $i$ is a pair of two textual elements. Given a source item $s\in D$, the task is to rank all the other items in $D$ according to their semantic similarity to $s$.

MetricBERT training is initialized with a pre-trained BERT model $f$. The training proceeds using triplets of textual elements $(x_{a},x_{p},x_{n})$ where $x_{a}$, $x_{p}$, and $x_{n}$ are anchor, positive and negative samples, respectively. In our settings, each anchor and positive samples are textual elements associated with the same item $p\in D$, while the negative sample is a different, randomly sampled item $n\in D$ according to the procedure described in Sec.3.1.

The textual elements are separately tokenized and masked (similar to BERT pre-training), aggregated over the batch dimension, and propagated through $f$. The embeddings of each of the three elements are used to train MetricBERT to (1) rank the positive and negative samples, based on the anchor text, and (2) reconstruct the masked words in each element. This is obtained by minimizing a combination of a novel triplet-loss, and a standard masked language loss²²2similar to BERT denoted by $\pazocal{L}_{T}$ and $\pazocal{L}_{MLM}$, respectively. The triplet loss ($\pazocal{L}_{T}$) is defined by:

$$\pazocal{L}_{T}=L(a,p,n)=\max(0,m+d(a,p)-d(a,n)),$$

(1)

where $m$ is a pre-defined margin, $a,p,n$ are the embeddings:

$$a=f\left(x_{a}\right),\;\;p=f\left(x_{p}\right),\;\;n=f\left(x_{n}\right),$$

(2)

and $d$ is the angular distance, which can be expressed as:

$$d(u,v)=\arccos\left(C(u,v)\right)/\pi$$

(3)

where $u,v\in\mathbb{R}^{d}$ are d-dimensional vectors and $C$ is the cosine similarity:

$$C(u,v)=\frac{u\cdot v}{\|u\|\|v\|}$$

(4)

The choice of the angular distance over a standard cosine similarity stems from early empirical observations that optimizing a cosine similarity results in a sizeable degradation in performance, where the model “collapses” into a narrow manifold (i.e. the embedding of all items in the dataset retrieved a cosine similarity that approaches 1). This improved performance of the angular distance can be attributed to its enhanced sensitivity to micro-differences in the angle between vectors with similar directions.

The improved sensitivity of the angular distance is also apparent in the derivatives of the two metrics. While the derivative of the cosine distance is:

$$\frac{\partial}{\partial u}\operatorname{C}(u,v)=\frac{v}{\|u\|\|v\|}-\operatorname{C}(u,v)\cdot\frac{u}{\|u\|^{2}}$$

(5)

the angular distance derivative can be expressed as:

$$\frac{\partial}{\partial u}d(u,v)=-\frac{1}{\pi\sqrt{1-\left(\operatorname{C}(u,v)\right)^{2}}}\frac{\partial}{\partial u}\operatorname{C}(u,v)$$

(6)

for which, compared to the cosine derivative, the denominator dynamically scales the gradients by an inverse ratio of the cosine distance between the vectors. In other words, compared to the cosine similarity, the angular distance yields gradients of larger magnitude for vectors with similar directions. See Sec. 4.5 for more details.

Finally, the total loss of MetricBERT is:

$$\pazocal{L}_{total}=\pazocal{L}_{MLM}+\lambda\pazocal{L}_{T},$$

(7)

where $\lambda$ is a balancing hyper-parameter. The MetricBERT model is illustrated in Fig. 1.

We evaluate MetricBERT on the publicly available wines recommendations dataset³³3https://doi.org/10.5281/zenodo.3653403 as well as the video-games recommendations dataset⁴⁴4https://zenodo.org/record/6088355, which we release as an additional contribution, to further accelerate research in the field.

We introduce MetricBERT, a BERT-based model that jointly optimizes a triplet-loss and a masked language model. We focus on text-based recommendation tasks where we demonstrate the ability of MetricBERT to improve upon the state-of-the-art sometimes by a considerable margin. While this work is focused on a recommendation task, MetricBERT is a general model, that can be utilized for additional text similarity and information retrieval tasks. Finally, to accelerate research and as an additional contribution, we publish a dataset of video games descriptions along with a test set of similarity annotations crafted by a domain expert.