SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 3: A comparison of ShareNorm and SepNorm in three tasks of graph property prediction.
the ZINC and ZINC (subset) datasets and the area under the curve (AUC) for the MolHIV dataset.
(a) Pretraining loss
(b) [CLS] embeddings
Figure 4: (a) Reconstruction loss of the MAE pertaining – MAE with SepNorm achieves lower MSE loss compared to ShareNorm, demonstrating a better ability to encode global contextual information. (b) & (c) Comparison of the singular values of learned ([CLS] and normal token) features with ShareNorm and different configurations of SepNorm. [CLS] embeddings learned from SepNorm have larger singular values, which suggests that vectors are better used to encode information.
(c) normal token embeddings
Experiment results. Table 3 shows the performances of our model and the Graphormer baseline. For the ZINC datasets, Graphormer with SepNorm achieves a significantly lower mean absolute error compared to that with ShareNorm. On the MolHIV dataset, SepNorm also improves the AUC to  $75.64\%$ , compared with ShareNorm's AUC of  $73.36\%$ . These results are strong evidence that the embeddings of the [VNode] can better summarize the properties of the entire graph and thus give superior performance on downstream tasks.
4.4 Uniformity Analysis
In this section, we investigate how, under both non-contrastive and contrastive training methods, ShareNorm and SepNorm respectively affect the uniformity of learned embeddings and further classification performances.
Experiment setup. We pretrain MAE on the STL10 dataset via four different losses:
- MAE loss  $\mathcal{L}_{\mathrm{MAE}}$  without any  $\mathcal{L}_{\mathrm{unif}}$  on [CLS] and token embeddings. This setting is a study with MAE training only.
- MAE loss  $\mathcal{L}_{\mathrm{MAE}}$  with  $\mathcal{L}_{\mathrm{unif}}$  on the [CLS] embeddings. We treat all [CLS] embeddings within the same batch (except itself) as negative instances.
- MAE loss  $\mathcal{L}_{\mathrm{MAE}}$  with  $\mathcal{L}_{\mathrm{unif}}$  on the token embeddings. We treat all token embeddings within the same batch or same images (except itself) as negative instances.
- MAE loss  $\mathcal{L}_{\mathrm{MAE}}$  with  $\mathcal{L}_{\mathrm{unif}}$  on both [CLS] and token embeddings.
We choose  $\lambda = \{0, 0.01, 0.1, 1\}$ . Note that the second loss with  $\lambda = 0.1$  corresponds to the UMAE [Zhang et al., 2022]. We also replace the normalization layer of the ViT in MAE with one of the following: {LN, BN, BN+LN, BN+BN}. The combination of different losses, different  $\lambda$ 's, and different normalization layers yields 40 specifications of the experiments.
We first report our results with MAE training only. The uniformity of learned embeddings is first measured by singular values of the decomposition of an embedding matrix: we randomly choose 10k embeddings to form the matrix. We do this separately for [CLS] embeddings and normal token embeddings. Figrue 4 shows the results, which indicate that [CLS] features learned from SepNorm exhibit better representational power and thus can better encode the global information.
Then the uniformity is measured by the score in Eqn. 4. Table 6(a) shows the numerical value of the uniformity on the STL10 and Aircraft datasets [Coates et al., 2011, Maji et al., 2013]. Compared to ShareNorm, SepNorm significantly enhances the uniformity of [CLS] embeddings. Interestingly, the uniformity of normal tokens' embeddings remains comparable. We also empirically verify that better uniformity on the [CLS] embeddings results in better performance on the downstream task (Figure 6(b)). Another observation is that the uniformity of [CLS] embeddings is clearly improved when they are normalized by BN instead of LN. Our hypothesis is that BN tries to make each feature dimension useful by controlling its variance while LN may still neglect some feature dimensions.
NormLayer	STL10	Aircraft
[CLS]	token	[CLS]	token
LN	-2.5911	-3.4998	-0.9097	-3.3514
BN	-2.9531	-3.6785	-1.5321	-3.3659
LN + LN	-1.1045	-3.3685	-1.6874	-3.3633
LN + BN	-1.3153	-3.3735	-1.9634	-3.5056
BN + LN	-3.4426	-3.3797	-2.7705	-3.3068
BN + BN	-3.4763	-3.4758	-2.9788	-3.7349
(a) Uniformity of different normalization settings.
(b) Uniformity vs. accuracy
We then report results from studies with U-MAE training. Figure 5 shows the uniformity metrics obtained using different  $\lambda$ 's. When using ShareNorm, the uniformity of the [CLS] embeddings is no better than -3.088, and even the explicit uniformity loss does not help much. On the contrary, embeddings learned from the proposed SepNorm can easily achieve better uniformity scores. The study with the contrastive approach further verifies the advantage of SepNorm in terms of encouraging uniformity of [CLS] embeddings.
The results provide strong evidence that the uniformity of the [CLS] embeddings is held down by ShareNorm even the minimization of an explicit contrastive loss cannot increase it. We
Figure 6: Uniformity Analysis. (a) Under SepNorm, the uniformity of the [CLS] embeddings are better preserved on the STL10 and Aircraft datasets. (b) Uniformity is positively related to the downstream task performance – lower uniformity results in higher classification accuracy on the STL10 dataset.
Figure 5:  $\lambda$  vs uniformity. A higher  $\lambda$  emphasizes more on the uniformity of [CLS] embeddings.
hypothesize that all features after LN will distribute in the same sphere, and [CLS] embeddings are squeezed to a small area of the sphere surface because they need to be different from embeddings of normal tokens.
Table 4 reports the downstream performance (accuracy) on STL10 across 40 different settings. We summarize our observations: (1) In the non-contrastive method MAE, with proper configuration, the performance of SepNorm is superior to that of ShareNorm. (2) In contrastive methods  $(\lambda \neq 0)$ , SepNorms' advantages are further highlighted. For example, when  $\lambda = 1$ , the performance of SepNorm (BN+LN) is improved by  $1.6\%$  compared to the non-contrastive method. The performance gain in SepNorm (BN+BN) is less obvious as the double BNs already impose implicit uniformity loss on both [CLS] and token embeddings.
In contrast to SepNorm, the performance of ShareNorm is less satisfactory when using contrastive methods. We believe it is very challenging to encourage the two types of embeddings to be uniformly distributed in the same sphere and keep them separable at the same time. (3) The uniformity of the token embeddings is also vital for learning a good [CLS] representation, as evident by SepNorm  $(\mathrm{BN} + \mathrm{LN})$  gaining accuracy with increasing  $\lambda$  on the token embeddings. We hypothesize that by enforcing uniformity, the token embeddings are forced to contain less information about others, which encourages the [CLS] embedding to encode as much information as possible. Our empirical study also shows that, when contrastive loss [Oord et al., 2018] is used to encourage the uniformity of [CLS] features in self-supervised transformers, the difference between BN and LN on [CLS] features is not significant anymore.
The training of transformer architectures with self-supervised learning has seen significant advancements in both contrastive and non-contrastive training. Among self-supervised learning methods,
Normalization layer	λ = 0	Negative pairs	λ = 0.01	λ = 0.1	λ = 1	Best
SepNorm (BN+BN)	93.84	token	93.65	94.15	93.94	94.15
[CLS]	93.73	93.85	93.93	93.93
[CLS] + token	93.40	94.25	94.28	94.28
SepNorm (BN+LN)	92.80	token	92.98	93.60	94.30	94.30
[CLS]	92.98	93.48	93.36	93.48
[CLS] + token	92.74	93.18	94.40	94.40
ShareNorm (BN)	92.84	token	92.48	93.38	92.78	93.38
[CLS]	93.10	93.33	92.93	93.33
[CLS] + token	93.41	93.46	92.99	93.46
ShareNorm (LN)	92.01	token	92.61	92.74	92.14	92.74
[CLS]	92.28	92.75	92.36	92.75
[CLS] + token	92.74	92.38	92.74	92.74
Table 4: Ablation study of the effect of [CLS] and token uniformity on the downstream tasks with  $\lambda$  varied. We report downstream task accuracy for the STL10 dataset.
non-contrastive ones do not rely on negative samples for learning. They have emerged as a powerful approach for training transformer models and demonstrated remarkable successes in various tasks. BERT [Devlin et al., 2018] and RoBERTa [Liu et al., 2019] were proposed in the NLP domain. Additionally, there are some works focus on the specific task, such as speech recognition [Wang et al., 2020], image generation [Chen et al., 2020a], and heterogeneous graph generation [Hu et al., 2020b].
Contrastive methods on the contrary train networks using positive and negative samples that are constructed without manual labeling. They have also been used to train transformer-based architectures. Gao et al. [2021] and Zhang et al. [2022] make significant strides in natural language processing tasks, while Chen et al. [2021] provide valuable insights into the pre-training of transformers. Meanwhile, the potential of contrastive methods in vision transformers has been demonstrated by Caron et al. [2021] and Radford et al. [2021]. These collective efforts underscore the versatility and efficacy of contrastive methods in self-supervised learning of transformers.
Normalization layers, including layer normalization and batch normalization, are essential to transformer architectures because they help stabilize the training procedure and accelerate convergence. Xiong et al. [2020] delve into the role of layer normalization in the transformer architecture and provide insights about how the layer improves the training stability and the performance of transformers. Similarly, Xu et al. [2019b] explores the intricacies of layer normalization and offers potential enhancements to its effectiveness. To address the limitations of traditional batch normalization in a transformer architecture, Shen et al. [2020] introduces a new normalization layer, Powernorm, which is a variant of batch normalization. Nguyen and Salazar [2019] focus on the normalization process in the self-attention mechanism of transformers and propose methods to optimize the normalization of self-attention. All the efforts above underscore the critical role of normalization layers in transformer models.
In this work, we have introduced SepNorm to separate the normalization of [CLS] embeddings from that of other tokens. Across three application domains (images, text, and graphs), SepNorm shows consistent performance improvement when it is incorporated into transformer models. Our analysis shows that SepNorm promotes uniformity of [CLS] embeddings and thus enhances the transformers' ability to encode information. As a valuable technique for improving the foundational transformer architecture, SepNorm has the potential to benefit a wide range of applications.
We thank all reviewers for their constructive feedback. This research is supported by the NIGMS of the National Institutes of Health, Awards R35GM148219, the Army Research Office, MURI program, contract # W911NF2210239, and NSF Award 1909536. Chen and Liu are also supported by the NSF CAREER Award # 2239869.
J. Xu, X. Sun, Z. Zhang, G. Zhao, and J. Lin. Understanding and improving layer normalization. Advances in Neural Information Processing Systems, 32, 2019b.
C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34: 28877-28888, 2021.
Q. Zhang, Y. Wang, and Y. Wang. How mask matters: Towards theoretical understandings of masked autoencoders. arXiv preprint arXiv:2210.08344, 2022.
W. Zhou, T. Ge, K. Xu, F. Wei, and M. Zhou. Scheduled drophead: A regularization method for transformer models. arXiv preprint arXiv:2004.13342, 2020.