Disentangled Contrastive Learning for Learning Robust Textual Representations

There are multiple ways to align a textual representation. We utilize two transformers with a consistent momentum representation to explicitly guarantee feature alignment [8]. The two networks are defined by a set of weights $\theta$ and $\xi$. We use the exponential moving average of the parameters $\theta$ to get $\xi$. Formally, we have:

Given a sentence $\mathcal{X}$ and its augmentation $\mathcal{X}^{\prime}$ (e.g, via data augmentation) from the first original network, we may obtain output representations $q\triangleq f_{\theta}(X)$ and $p\triangleq f_{\theta}(X^{\prime})$. Note that previous works [3, 8] indicates that an projection $p$ in feature space improve the performance. We then leverage a projection function $g(p_{\theta})$ and $\ell_{2}$-normalize both $g(p_{\theta})$ and $q_{\xi}$ to $\bar{g}(p_{\theta})\triangleq g/\|g(p_{\theta})\|_{2}$ and $\bar{q}_{\xi}\triangleq q_{\xi}/\|q_{\xi}\|_{2}$, respectively. We leverage the mean squared loss as follows:

Additionally, we make the losses symmetrical $\mathcal{L}_{\mathrm{align}}$ by feeding $X$ to the augmented network and $X^{\prime}$, separately. We optimize $\mathcal{L}_{\mathrm{align}}+\widetilde{\mathcal{L}}_{\mathrm{align}}$ with respect to $\theta$ only, but not $\xi$, via the stop-gradient.

To ensure that feature vectors should have an approximately uniform distribution, we can directly optimize the Eq. 6. However, different from computer vision, in the original loss of BRET [5], we have already utilized the next sentence prediction loss. Such a contrastive object has explicitly made the sentence representation $f(.)$ scattered in the feature space; thus, the model may quickly collapse without learning. Inspired by [22], we argue that batch normalization can identify the common-mode between examples of a mini-batch and removes it using the other representations in the mini-batch as implicit negative examples. We can, therefore, view batch normalization as a novel method of implementing feature uniformity on embedded representations. Because vanilla batch normalization will lead to significant performance degradation when naively used in NLP, we leverage an enhanced power normalization [24] to guarantee feature uniformity. Specifically, we leverage the unit quadratic mean rather than the mean/variance of running statistics with an approximate backpropagation method to compute the corresponding gradient. Formally, we have the following:

Note that we compute the gradient of the loss regarding the quadratic mean of the batch. In other words, we utilize the running statistics to conduct backpropagation, thus, resulting in bounded gradients, which is necessary for convergence in NLP (see proofs in [24]).

We leverage synonyms from WordNet categories to conduct data augmentation for computation efficiency. We combine all the momentum representation consistency and power normalization results in a unified architecture with the mask language model object. We leverage the same architecture of the BERT-base [5]. We first pre-train the model in a large-scale corpus unsupervisedly (e.g., the same corpus and training steps with BERT) and then fine-tune the model using task datasets.

We conducted experiments on three benchmarks: GLUE, SQuAD, SNLI, and DialogRE.

GLUE [26] is an NLP benchmark aimed at evaluating the performance of downstream tasks of the pre-trained models. Notably, we leverage nine tasks in GLUE, including CoLA, RTE, MRPC, STS, SST, QNLI, QQP, and MNLI-m/mm. We follow the same setup as the original BERT for single sentence and sentence pair classification tasks. We leverage a multi-layer perception with a softmax layer to obtain the predictions.

SQuAD is a reading comprehension dataset constructed from Wikipedia articles. We report results on SQuAD 1.1. Here also, we follow the same setup as the original BERT model and predict an answer span—the start and end indices of the correct answer in the correct context.

SNLI is a collection of 570k human-written English sentence pairs that have been manually labeled for balanced classification with entailment, contradiction, and neutral labels, thereby supporting the task of natural language inference (NLI). We add a linear transformation and a softmax layer to predict the correct label of NLI.

DialogRE is a dialogue-based relation extraction dataset, which contains 1,788 dialogues from a famous American television situation comedy Friends.

To evaluate the robustness of our approach, we also conduct invariance testing with CheckList¹¹1https://github.com/marcotcr/checklist.git [21] and adversarial attacks²²2https://github.com/thunlp/OpenAttack. To generate label-preserving perturbations, we used WordNet categories (e.g., synonyms and antonyms). We selected context-appropriate synonyms as permutation candidates. To generate adversarial samples, we leverage a probability-weighted word saliency (PWWS) [19] method based on synonym replacement. We manually evaluate the quality of the generated instances. We also conduct experiments that apply data augmentation and adversarial training to the BERT model. We utilize PyTorch to implement our model. We use Adam optimizer with a cosine decay learning rate schedule. We set the initial learning rate as 1e-5. We use a batch size of 32 over eight Nvidia 1080Ti GPUs.

Main Results

From Table 1 and 2, we can observe the following: 1) Vanilla BERT achieves poor performance in the robust set on both GLUE and SQUAD, which indicates that the previous fine-tuning approach cannot obtain a robust textual representation. This will lead to performance decay with permutations.

2) With data augmentation, BERT can obtain improved performance in the robust set; however, a slight performance decay is observed in the original test set. We argue that data augmentation can obtain better performance by fitting to task-specific data distribution; there is no guarantee that more data will result in robust textual representations.

3) Our DCL approach achieves improved performance in both the original test set and robust set compared with vanilla BERT. Note that our DCL is an unsupervised approach, and we leverage the same training instances with BERT. The performance improvements indicate that our approach can obtain more robust textual representations that enhance the performance of the system.

Adversarial Attack Results

From Table 3, we can observe the following: 1) Vanilla BERT achieves a poor performance with adversarial attacks; BERT with adversarial training can obtain a good performance. However, we notice that there exists a performance decay for adversarial training in the original test set. Note that adversarial training methods would lead to standard performance degradation [29], i.e., the degradation of natural examples. 2) Our DCL approach achieves improved performance in the test set with and without an adversarial attack, which further demonstrates that our approach can obtain robust textual representations that are stable for different types of permutations.

Quantitative Analysis of Textual Representation

As we hypothesize that power normalization can implicitly contribute to feature uniformity, we conduct further experiments to analyze the effects of normalization [1]. Specifically, we random sample instances and leverage the cosine similarity of the original input projection vectors and the augmented projection vectors. We calculate the average cosine similarity between positive instances (in blue) and random instances (in red) with different strategies, including without normalization (No Norm), batch normalization (BN), and power normalization.

From Figure 3, we observe that with no normalization in $p$ or $q$, the feature space is aligned for both positive and negative instances, which shows that there exists a feature collapse for textual representation learning. Considering DCL training (i.e., with power normalization), we notice that the textual representations are relatively more similar between the positive instances (0.9842) than random (negative) ones (0.7904); thus, we can obtain different vectors.

Next, we give an intuitive explanation of preventing feature collapse for textual representation learning. Given an input instance without negative examples, the model may always output the projection vector $z$ with $[0,1,0,0,…]$. Thus, the model can achieve a perfect prediction through learning a simple identity function, which, in other words, collapse in the feature space. With normalization, the output vector $z$ cannot obtain such singular values. Since the outputs will be redistributed regarding the learned mean and standard deviation, we can implicitly learn robust representations.

Qualitative Analysis of Textual Representation

We randomly selected instances to visualize a sentence with T-SNE [15] to better understand the behaviors of textual representations. The different color refers to the different sentence pairs for both random permutation and adversarial attack settings. From Figure 2, it may be observed that our approach can obtain a relatively similar semantic representation with permutations in both invariant tests and adversarial attack settings. Note that we explicitly align the projection of the textual representation with a random permutation, thereby encouraging similar semantic instances to have relatively similar representations.

Robust Representation with Contrastive Learning. Conventional approaches usually try to leverage instance-level augmentation aimed at achieving good performance on a robust set. However, there is no guarantee that robust textual representations will be obtained. Intuitively, directly aligning the representation of input tokens with slight permutations may contribute to robust representations. However, without any negative constraints, the model will easily collapse with a sub-optimal solution. In this study, we observe that power normalization identifies this common mode between examples. In other words, it can remove those trivial samples by using the other representations in the batch as implicit negative instances. We can, therefore, view normalization as an implicitly contrastive learning method.

Limitations. This work is not without limitations. We only consider the synonym replacement as a data augmentation strategy due to the efficiency of processing a huge amount of data. Other strong data augmentation methods can also be leveraged. Another issue is representation alignment, as there are lots of augmentations. We cannot enumerate all positive pairs for alignments; thus, there is still some room for designing more efficient feature-aligning algorithms. Moreover, as we utilize the square root loss, which is absolutely a Euclidean distance. Recent approaches [16] indicates that Euclidean space may be sub-optimal for textual representations, and we leave this for future works.