ContraNorm: A Contrastive Learning Perspective on Oversmoothing and Beyond

In this part, we begin by highlighting the limitations of existing metrics in characterizing oversmoothing. These limitations motivate us to adopt the effective rank metric, which has been demonstrated to be effective in capturing the degree of dimensional collapse in contrastive learning.

Taking the oversmoothing problem of Transformers as an example without loss of generality, a prevailing metric is evaluating attention map similarity (Wang et al., 2022; Gong et al., 2021; Shi et al., 2022). The intuition is that as attention map similarity converges to one, feature similarity increases, which can result in a loss of representation expressiveness and decreased performance. However, by conducting experiments on transformer structured models like ViT and BERT, we find that high attention map similarity does not necessarily correspond to high feature similarity. As shown in Figure 2(a) and Figure 2(b), although attention map similarity is close to 0.8 or even higher, the feature similarity remains below 0.5 in most cases. It means the oversmoothing problem still occurs even with a low feature similarity. This finding suggests that similarity can not fully depict the quality of representations and the oversmoothing problem.

Intuitively, we can consider a case that representations in the latent embedding space do not shrink to one single point but span a low dimensional space. In such cases, feature similarity may be relatively low, but representations still lose expressive power. Such representation degeneration problem is known as dimensional collapse, and it is widely discussed in the literature of contrastive learning. In contrastive learning, common practice to describe dimensional collapse is the vanishing distribution of singular values (Gao et al., 2019; Ethayarajh, 2019; Jing et al., 2022). To investigate whether dimensional collapse occurs in Transformers, we draw the singular value distribution of features in the last block of 12-layer BERT. As shown in Figure 2(c), the insignificant (nearly zero) values dominate singular value distribution in the deep layer of vanilla BERT, indicating that representations reside on a low-dimensional manifold and dimensional collapse happens. To show the collapse tendency along layer index, we simplify the singular value distribution to a concise scalar effective rank (erank) (Roy & Vetterli, 2007), which covers the full singular value spectrum.

Based on erank, we revisit the oversmoothing issue of Transformers in Figure 2(d). We can see that the effective rank descends along with layer index, indicating an increasingly imbalanced singular value distribution in deeper layers. This finding not only verifies dimensional collapse does occur in Transformers, but also indicates the effectiveness of effective rank to detect this issue.

The core idea of contrastive learning is maximizing agreement between augmented views of the same example (i.e., positive pairs) and disagreement of views from different samples (i.e. negative pairs). A popular form of contrastive learning optimizes feature representations using a loss function with limited negative samples (Chen et al., 2020). Concretely, given a batch of randomly sampled examples, for each example we generate its augmented positive views and finally get a total of $N$ samples. Considering an encoder function $f$, the contrastive loss for the positive pair $(i,i^{+})$ is

where $\tau$ denotes the temperature. The loss can be decoupled into alignment loss and uniformity loss:

The alignment loss encourages feature representations for positive pairs to be similar, thus being invariant to unnecessary noises. However, training with only the alignment loss will result in a trivial solution where all representations are identical. In other words, complete collapse happens. While batch normalization (Ioffe & Szegedy, 2015) can help avoid this issue, it cannot fully prevent the problem of dimensional collapse, which still negatively impacts learning (Hua et al., 2021).

Thanks to the property of uniformity loss, dimensional collapse can be solved effectively. Reviewing the form of uniformity loss in Eq. (2), it maximizes average distances between all samples, resulting in embeddings that are roughly uniformly distributed in the latent space, and thus more information is preserved. The inclusion of the uniformity loss in the training process helps to alleviate dimensional collapse. Intuitively, it can also serve as a sharp tool for alleviating oversmoothing in models such as GNNs and Transformers.

An alternative approach to address oversmoothing is to directly incorporate the uniformity loss into the training objective. However, our experiments reveal that this method has limited effectiveness (see Appendix B.2 for more details). Instead, we propose a normalization layer that can be easily integrated into various models. Our approach utilizes the uniformity loss as an underlying energy function of the proposed layer, such that a descent step along the energy function corresponds to a forward pass of the layer. Alternatively, we can view the layer as the unfolded iterations of an optimization function. This perspective is adopted to elucidate GNNs (Yang et al., 2021; Zhu et al., 2021), Transformers (Yang et al., 2022), and classic MLPs (Xie et al., 2021).

Note that the uniformity loss works by optimizing model parameters while what the normalization layer directly updates is representation itself. So we first transfer the uniformity loss, which serves as a training loss, to a kind of architecture loss. Consider a fully connected graph with restricted number of nodes, where node ${\bm{h}}_{i}$ is viewed as representation of a random sample ${\bm{x}}_{i}$. Reminiscent of $\mathcal{L}_{\rm{uniform}}$, we define $\hat{\mathcal{L}}_{\rm{uniform}}$ over all nodes in the graph as

This form of uniformity loss is defined directly on representations, and we later use it as the underlying energy function for representation update.

Till now, we are able to build a connection between layer design and the unfolded uniformity-minimizing iterations. Specifically, we take the derivative of $\hat{L}_{\rm{uniform}}$ on node representations:

Denote the feature matrix as ${\bm{H}}$ with the $i$-th row ${\bm{h}}_{i}^{\top}$, we can rewrite Eq. (4) into a matrix form:

where ${\bm{A}}=\exp({\bm{H}}{\bm{H}}^{\top}/\tau)$ and ${\bm{D}}=\deg({\bm{A}})$¹¹1$\deg({\bm{A}})$ is a diagonal matrix, whose $(i,i)$-th element equals to the sum of the $i$-th row of ${\bm{A}}$.. To reduce the uniformity loss $\hat{\mathcal{L}}_{\rm{uniform}}$, a natural way is to take a single step of gradient descent that updates ${\bm{H}}$ by

where ${\bm{H}}_{b}$ and ${\bm{H}}_{t}$ denote the representations before and after the update, respectively, and $s$ is the step size of the gradient descent. By taking this update after a certain representation layer, we can reduce the uniformity loss of the representations and thus help ease the dimensional collapse.

In Eq. (6), there exist two terms ${\bm{D}}^{-1}{\bm{A}}$ and ${\bm{A}}{\bm{D}}^{-1}$ multiplied with ${\bm{H}}_{b}$. Empirically, the two terms play a similar role in our method. Note that the first term is related to self-attention matrix in Transformers, so we only preserve it and discard the second one. Then Eq. (6) becomes

In fact, the operation corresponds to the stop-gradient technique, which is widely used in contrastive learning methods (He et al., 2020; Grill et al., 2020; Tao et al., 2022). By throwing away some terms in the gradient, stop-gradient makes the training process asymmetric and thus avoids representation collapse with less computational overhead, which is discussed in detail in Appendix B.1.

However, the layer induced by Eq. (7) still can not ensure uniformity on representations. Consider an extreme case where $\mathrm{softmax}({\bm{H}}_{b}{\bm{H}}_{b}^{\top})$ is equal to identity matrix ${\bm{I}}$. Eq. (7) becomes ${\bm{H}}_{t}={\bm{H}}_{b}-s/\tau\times{\bm{H}}_{b}=(1-s/\tau){\bm{H}}_{b}$, which just makes the scale of ${\bm{H}}$ smaller and does not help alleviate the complete collapse. To keep the representation stable, we explore two different approaches: feature norm regularization and layer normalization.

Proposition 1 gives a bound for the ratio of the variance after and before the update. It shows that the change of the variance is influenced by the symmetric matrix ${\bm{P}}=({\bm{I}}-{\bm{e}}{\bm{e}}^{\top})({\bm{I}}-\bar{{\bm{A}}})+({\bm{I}}-\bar{{\bm{A}}})^{\top}({\bm{I}}-{\bm{e}}{\bm{e}}^{\top})$. If ${\bm{P}}$ is a semi-positive definite matrix, we will get the result that $Var({\bm{H}}_{t})\geq Var({\bm{H}}_{b})$, which indicates that the representations become more uniform. In Appendix F, we will give out some sufficient conditions for that ${\bm{P}}$ is semi-positive definite.

II. Appending LayerNorm. Another option is to leverage Layer Normalization (LayerNorm) (Ba et al., 2016) to maintain the feature scale. The update form of LayerNorm is $\text{LayerNorm}({\bm{h}}_{i})=\gamma\cdot(({\bm{h}}_{i}-\text{mean}({\bm{h}}_{i}))/\sqrt{\text{Var}({\bm{h}}_{i})+\varepsilon})+\beta$, where $\gamma$ and $\beta$ are learnable parameters and $\varepsilon=10^{-5}$. The learnable parameters $\gamma$ and $\beta$ can rescale the representation ${\bm{h}}_{i}$ to help ease the problem. We append LayerNorm to the original update in Eq. (7) and obtain

where applying the layer normalization to a representation matrix ${\bm{H}}$ means applying the layer normalization to all its components ${\bm{h}}_{1},\dots,{\bm{h}}_{n}$.

ContraNorm. We empirically compare the two proposed methods and find their performance comparable, while the second one performs slightly better. Therefore, we adopt the second update form and name it Contrastive Normalization (ContraNorm). The ContraNorm layer can be added after any representation layer to reduce the uniformity loss and help relieve the dimensional collapse. We discuss the best place to plug our ContraNorm in Appendix B.3.

ContraNorm-D: Scalable ContraNorm for Large-scale Graphs. From the update rule in Eq. (9), we can see that the computational complexity of the ContraNorm layer is ${\mathcal{O}}(n^{2}d)$, where $n$ is the number of tokens and $d$ is the feature dimension per token, which is in the same order of self-attention (Vaswani et al., 2017). Therefore, this operation is preferable when the token number $n$ is similar or less than the feature dimension $d$, as is usually the case in CV and NLP domains. However, for large-scale graphs that contain millions of nodes ($n\gg d$), the quadratic dependence on $n$ makes ContraNorm computationally prohibitive. To circumvent this issue, we propose Dual ContraNorm (ContraNorm-D), a scalable variant of ContraNorm whose computational complexity ${\mathcal{O}}(nd^{2})$ has a linear dependence on $n$, with the following dual update rule,

where we calculate the dimension-wise feature correlation matrix ${\bm{H}}_{b}^{\top}{\bm{H}}_{b}\in{\mathbb{R}}^{d\times d}$ and multiple it to the right of features ${\bm{H}}_{b}$ after softmax normalization. In Appendix A, we provide a more thorough explanation of it, showing how to derive it from the feature decorrelation normalization adopted in non-contrastive learning methods (like Barlow Twins (Zbontar et al., 2021) and VICReg (Bardes et al., 2022)). As revealed in recent work (Garrido et al., 2022), there is a dual relationship between contrastive and non-contrastive learning methods, and we can therefore regard ContraNorm-D as a dual version of ContraNorm that focuses on decreasing dimension-wise feature correlation.

In this part, we give a theoretical result of ContraNorm, and show that with a slightly different form of the uniformity loss, the ContraNorm update can help alleviate the dimensional collapse.

Proposition 2 gives a promise of alleviating dimensional collapse under a special update form. Denote $\hat{\mathcal{L}}_{\rm{dim}}=tr(({\bm{I}}-{\bm{H}}{\bm{H}}^{\top})^{2})/4$, then $\hat{\mathcal{L}}_{\rm{dim}}$ shares a similar form with Barlow Twins (Zbontar et al., 2021). This loss tries to equate the diagonal elements of the similarity matrix ${\bm{H}}{\bm{H}}^{\top}$ to 1 and equate the off-diagonal elements of the matrix to 0, which drives different features becoming orthogonal, thus helping the features become more uniform in the space. Note that $\frac{\partial\hat{\mathcal{L}}_{\rm{dim}}}{\partial{\bm{H}}}=({\bm{I}}-{\bm{H}}{\bm{H}}^{\top}){\bm{H}}$, therefore, the update above can be rewritten as ${\bm{H}}_{t}={\bm{H}}_{b}+s\frac{\partial\hat{\mathcal{L}}_{\rm{dim}}({\bm{H}}_{b})}{\partial{\bm{H}}_{b}}$, which implies that this update form is a ContraNorm with a different uniformity loss. Proposition 2 reveals that this update can increase the effective rank of the representation matrix, when $s$ satisfies $1+(1-\sigma_{\max})s>0$. Note that no matter whether $\sigma_{\max}>0$ or not, if $s$ is sufficiently close to $0$, the condition will be satisfied. Under this situation, the update will alleviate the dimensional collapse.

Setup. To corroborate the potential of ContraNorm, we integrate it into two transformer structured models: BERT and ALBERT, and evaluate it on GLUE datasets. GLUE includes three categories of tasks: (i) single-sentence tasks CoLA and SST-2; (ii) similarity and paraphrase tasks MRPC, QQP, and STS-B; (iii) inference tasks MNLI, QNLI, and RTE. For MNLI task, we experiment on both the matched (MNLI-m) and mismatched (MNLI-mm) versions. We default plug ContraNorm after the self-attention module and residual connection (more position choices are in Appendix B.3). We use a batch size of 32 and fine-tune for 5 epochs over the data for all GLUE tasks. For each task, we select the best scale factor $s$ in Eq. (6) among $(0.005,0.01,0.05,0.1,0.2)$. We use base models (BERT-base and ALBERT-base) of 12 stacked blocks with hyperparameters fixed for all tasks: number of hidden size 128, number of attention heads 12, maximum sequence length 384. We use Adam (Kingma & Ba, 2014) optimizer with the learning rate of $2e-5$.

Results. As shown in Table 1, ContraNorm substantially improves results on all datasets compared with vanilla BERT. Specifically, our ContraNorm improves the previous average performance from 82.59% to 83.39% for BERT backbone and from 83.74% to 84.67% for ALBERT backbone. We also submit our trained model to GLUE benchmark leaderboard and the results can be found in Appendix E.1. It is observed that BERT with ContraNorm also outperforms vanilla model across all datasets. To verify the de-oversmoothing effect of ContraNorm. We build models with/without ContraNorm on various layer depth settings. The performance comparison is shown in Figure 3(a). Constant stack of blocks causes obvious deterioration in vanilla BERT, while BERT with ContraNorm maintains competitive advantage. Moreover, for deep models, we also show the tendency of variance and effective rank in Figure 3(b) and Figure 3(c), which verifies the power of ContraNorm in alleviating complete collapse and dimensional collapse, respectively.

We also validate effectiveness of ContraNorm in computer vision field. We choose DeiT as backbone and models are trained from scratch. Experiments with different depth settings are evaluated on ImageNet100 and ImageNet1k datasets. Based on the Timm (Wightman, 2019) and DeiT repositories, we insert ContraNorm into base model intermediately following self-attention module.

Setup. We follow the training recipe of Touvron et al. (2021) with minimal hyperparameter changes. Specifically, we use AdamW (Loshchilov & Hutter, 2019) optimizer with cosine learning rate decay. We train each model for 150 epochs and the batch size is set to 1024. Augmentation techniques are used to boost the performance. For all experiments, the image size is set to be $224\times 224$. For Imagenet100, we set the scale factor $s$ to 0.2 for all layer depth settings. For Imagenet1k, we set $s$ to 0.2 for models with 12 and 16 blocks, and raise it to 0.3 for models with 24 blocks.

Results. In Table 2, DeiT with ContraNorm outperforms vanilla DeiT with all layer settings on both datasets. Typically, our method shows a gain of nearly 5% on test accuracy for ImageNet100. For ImageNet1k, we boost the performance of DeiT with 24 blocks from 77.69% to 78.67%.

Setup. We implement ContraNorm as a simple normalization layer after each graph convolution of GCN, and evaluate it on fully supervised graph node classification task. For datasets, we choose two popular citation networks Cora (McCallum et al., 2000) and Citeseer (Giles et al., 1998), and Wikipedia article networks Chameleon and Squirrel (Rozemberczki et al., 2021). More information of datasets is deferred to Appendix D. We compare ContraNorm against a popular normalization layer PairNorm (Zhao & Akoglu, 2020) designed for preventing oversmoothing. We also take LayerNorm as a baseline by setting the scale $s=0$. We follow data split setting in Kipf & Welling (2017) with train/validation/test splits of 60%, 20%, 20%, respectively. To keep fair in comparison, we fix the hidden dimension to 32, and dropout rate to 0.6 as reported in Zhao & Akoglu (2020). We choose the best of scale controller $s$ in range of $\{0.2,0.5,0.8,1.0\}$ for both PairNorm and ContraNorm. For PairNorm, we choose the best variant presented by Zhao & Akoglu (2020).

Results. As shown in Table 3, in shallow layers (e.g., two layer), the addition of ContraNorm reduces the accuracy of vanilla GCN by a small margin, while it helps prevent the performance from sharply deteriorating as the layer goes deeper. ContraNorm outperforms PairNorm and LayerNorm in the power of de-oversmoothing. Here, we show the staple in diluting oversmoothing is ContraNorm, and LayerNorm alone fails to prevent GCN from oversmoothing, even amplifying the negative effect on Cora with more than 16 layers.

Recalling Section 3.3, we improve ContraNorm with stop-gradient technique (SG) by masking the second term. For solving data representation instability, we apply layer normalization (LN) to the original version, while for the convenience of theoretical analysis, layer normalization is replaced by $L_{2}$ normalization ($L_{2}$N). Here, we investigate the effect of these tricks and the results are shown in Table 4. Compared with the variant with only LayerNorm, ContraNorm with both stop gradient and layer normalization presents better performance. As for the two normalization methods, they are almost comparable to our methods, which verifies the applicability of our theoretical analysis.

In Appendix E.2, we further conduct an ablation study regarding the gains of ContraNorm while using different values of scale factor $s$, and show that ContraNorm is robust in an appropriate range of $s$.