SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiaohui Chen
Medford, MA 02155
xiaohui.chen@tufts.edu
Yinkai Wang
yinkai.wang@tufts.edu
Yuanqi Du
Ithaca, NY 14850
yd392@cornell.edu
Soha Hassoun
soha.hassoun@tufts.edu
Li-Ping Liu
liping.liu@tufts.edu
Self-supervised training methods for transformers have demonstrated remarkable performance across various domains. Previous transformer-based models, such as masked autoencoders (MAE), typically utilize a single normalization layer for both the class token [CLS] and the tokens. We propose in this paper a new yet simple normalization method that separately normalizes embedding vectors respectively corresponding to normal tokens and the [CLS] token, in order to better capture their distinct characteristics and enhance downstream task performance. Our empirical study shows that the [CLS] embeddings learned with our separate normalization layer better encode the global contextual information and are distributed more uniformly in its anisotropic space. When the conventional normalization layer is replaced with a separate normalization layer, we observe an average  $2.7\%$  performance improvement in learning tasks from the image, natural language, and graph domains.
Transformer models [Vaswani et al., 2017] have revolutionized natural language processing (NLP) [Devlin et al., 2018, Liu et al., 2019] and demonstrated remarkable performances across a wide range of NLP tasks. The significance of transformer models lies in their ability to model context and capture complex linguistic patterns without being constrained by the sequential nature of data. Beyond NLP transformers have further found their successes in areas such as computer vision (CV) [Han et al., 2022], speech recognition [Karmakar et al., 2021], and recommendation systems [Sun et al., 2019, Gu et al., 2020, Wu et al., 2020]. Their flexible architecture and ability to capture dependencies have made them adaptable to diverse data modalities in these domains.
Figure 1: Comparison of the shared normalization (ShareNorm, left) and the proposed separate normalization (SepNorm, right) configurations for token normalization. In the ShareNorm setup, both the [CLS] symbol and other tokens are normalized using a single-layer normalization. In the SepNorm setup, normalization is done separately: the [CLS] symbol is normalized through batch normalization, while other tokens are normalized via layer normalization.
Transformer architectures have been studied extensively from various perspectives such as attention mechanisms, positional encoding [Devlin et al., 2018], and normalization techniques. Specifically, layer normalization [Ba et al., 2016] and batch normalization [Ioffe and Szegedy, 2015] are employed to enhance stability and speed up convergence during training. The literature on transformers also explores parameter initialization [Xu et al., 2019a], optimization algorithms [Huang et al., 2020], regularization techniques [Steiner et al., 2021, Zhou et al., 2020], and improved architectures [Han et al., 2021]. This collective research has advanced transformer architectures and their applications in NLP, CV, and other learning domains.
The study of normalization in transformer architectures is motivated by several factors [Xiong et al., 2020, Shen et al., 2020, Nguyen and Salazar, 2019]. For example, Xiong et al. [2020] emphasize the importance of the warm-up of the learning rate and the position of layer normalization layers for the purpose of stable training and faster convergence. Shen et al. [2020] investigates the disadvantage of using batch normalization in transformers and proposes power normalization. While most previous works focus on how the normalization layer can be modified to stabilize the training process, it is less understood how the normalization affects the encoding abilities of these embeddings.
In self-supervised transformers, the [CLS] symbol is frequently used as a global representation for various downstream tasks [Devlin et al., 2018, He et al., 2022]. Often, the normalization applied to the [CLS] symbol is shared with the rest of the tokens in the sequence, which we term it as Shared Normalization (ShareNorm). Given that the [CLS] symbol plays a special role in representation learning, a natural question is whether we should treat it separately in the normalization operation. Driven by this question, our research first scrutinizes the behavior of the current shared normalization in transformers, particularly the properties of the [CLS] embedding and its influence on downstream task performance. Subsequently, we propose a replacement of ShareNorm with Separate Normalization (SepNorm), the latter of which employs distinct normalization operations for the [CLS] symbol and the token features, as depicted in Figure 1. Through extensive analysis, we demonstrate that [CLS] embeddings learned using ShareNorm have the issue of dimensional collapse, which cannot be rectified even by enforcing uniformity [Wang and Isola, 2020]. However, the straightforward substitution of SepNorm for ShareNorm substantially mitigates this issue. We empirically validate the effectiveness of SepNorm in tasks from the image, text, and graph domains, demonstrating the universal advantage of the proposed SepNorm.
2.1 Pretraining Transformers with the [CLS] symbol
Unsupervised pretraining of a transformer-based model [Vaswani et al., 2017] is widely investigated in many domains, including NLP, computer vision (CV), and graphs.
Pretraining BERT for NLP. In NLP, Devlin et al. [2018] first develop the BERT model by pretraining a transformer-based network by performing the following two tasks - masked language modeling and next sentence prediction. During pretraining, BERT takes a pair of sentences  $(\mathbf{x},\mathbf{y})$
which are represented as a special sequence
$$
\mathbf {s} = \left([ \mathrm {C L S} ], \mathbf {x}, [ \mathrm {S E P} ], \mathbf {y}\right). \tag {1}
$$
Here [SEP] is a special token that separates the two sentences. A fraction (e.g.,  $15\%$ ) of the tokens in  $\mathbf{x}$  and  $\mathbf{y}$  are randomly replaced by a special symbol [MASK]. The first task in BERT is to predict the original tokens replaced by [MASK] with cross-entropy loss. The second task is to predict whether  $\mathbf{y}$  is the next sentence following  $\mathbf{x}$ , and the decision is made by classifying the final embedding of the [CLS] symbol. After pretraining, the representation in the [CLS] is usually used for sentence-level downstream tasks such as sentiment analysis [Medhat et al., 2014].
Pretraining MAE for CV. The Vision Transformer (ViT) [Dosovitskiy et al., 2020] applies the transformers to computer vision tasks. In ViT, an image is usually voxelized into  $16 \times 16$  patches, which are then flattened into a sequence of 256 tokens and fed into the ViT. He et al. [2022] proposes a self-supervised training scheme, Masked Autoencoder (MAE), for the ViT architecture. A training image has  $75\%$  of its patches masked. The MAE feeds Tokens of unmasked patches as well as a [CLS] token into the encoder and gets the representations for these tokens. Then the decoder tries to reconstruct the original image by minimizing the mean square error (MSE). Only the encoder will be used for downstream tasks after pretraining. The [CLS] symbol is treated as the class token for linear probing and fine-tuning in the downstream tasks.
Pretraining Graphormer for molecule discovery. Graphormer [Ying et al., 2021] is a transformer-based model designed for graph representation learning tasks. It is used to predict the property of a graph rather than a node or edge. Specifically, Graphormer introduces a new symbol [VNode] as a node connecting to all original graph nodes. Then the vector learned for [VNode] represents the global information of the entire graph. The mechanism of [VNode] is similar to the [CLS] symbol in BERT and MAE.
In typical applications of transformers, the [CLS] symbol is not a natural data token. It summarizes other tokens to capture global information, which is especially useful in downstream tasks. For these reasons, we argue that it should be treated differently in normalization operations.
2.2 Normalization Layers in Transformers
Given that transformers are initially proposed for NLP tasks, layer normalization (LN) Ba et al. [2016] is typically the normalization method of choice [Xiong et al., 2020]. LN normalizes across feature dimensions and is independent of the sequence length and the batch size. For any features  $\mathbf{h} \in \mathbb{R}^d$ , the LN has the following computation:
$$
\operatorname {L N} (\mathbf {h}) = \gamma \odot \frac {\mathbf {h} - \mu}{\sigma} + \beta , \quad \mu = \frac {1}{d} \sum_ {i = 1} ^ {d} h _ {i}, \quad \sigma = \sqrt {\frac {1}{d} \sum_ {i = 1} ^ {d} \left(h _ {i} - \mu\right) ^ {2}}. \tag {2}
$$
Here  $h_i$  is the  $i$ -th dimension of  $\mathbf{h}$ ,  $\odot$  represents element-wise multiplication, and  $\gamma, \beta \in \mathbb{R}^d$  are scale and bias parameters, respectively. In a transformer, all tokens, including special tokens, such as [CLS] and [SEP], are all treated equally and share the same LNs.
Batch Normalization (BN) [Ioffe and Szegedy, 2015] works by normalizing the input data to have zero mean and unit variance along the batch dimension, followed by an affine transformation to scale the result using gamma and beta parameters. BN normalizes a given vector  $\mathbf{h}$  as:
$$
\operatorname {B N} (\mathbf {h}) = \gamma \odot \frac {\mathbf {h} - \boldsymbol {\mu} _ {B}}{\boldsymbol {\sigma} _ {B}} + \beta . \tag {3}
$$
Here  $\pmb{\mu}_B, \pmb{\sigma}_B^2 \in \mathbb{R}^d$  are the running statistics (mean and variance) maintained by the BN. The running mean and variance are updated during training after each batch. They are usually calculated as an exponential moving average of the batch mean and variance. BN is widely adopted in CV but leads to significant performance degradation when naively used in NLP.
2.3 Uniformity of the Learned Representations
The dimensional collapse in self-supervised representation learning is a common phenomenon where the embedding vectors only span a lower-dimensional subspace [Jing et al., 2021] of the entire vector
space. This means that the model fails to capture data patterns with full power and instead collapses to a simpler representation. Contrastive methods [Oord et al., 2018, Chen et al., 2020b] have been one of the standard approaches to address this problem. Specifically, Wang and Isola [2020] propose the uniformity metric (loss) to quantify the degree of dimensional collapse. Given a set of representation vectors  $\{\mathbf{h}_1,\dots ,\mathbf{h}_N\}$  from a dataset of size  $N$ , the uniformity metric  $\mathcal{L}_{\mathcal{U}}$  is computed as follows:
$$
\mathcal {L} _ {\mathcal {U}} = \log \frac {1}{N (N - 1) / 2} \sum_ {\substack {n = 1, \\ m = n + 1}} ^ {N, N} \exp^ {- 2 \left\| \frac {\mathrm {h} _ {n}}{\| \mathrm {h} _ {n} \|} - \frac {\mathrm {h} _ {m}}{\| \mathrm {h} _ {m} \|} \right\| ^ {2}}. \tag{4}
$$
If the distribution of the representation is perfectly uniform, then the numerical value of  $\mathcal{L}_{\mathcal{U}}$  will converge to -4 as the dimension of  $\mathbf{h}$  increases to infinity [Wang and Isola, 2020].
In self-supervised transformers, the uniformity of the representation is also taken into consideration by some works. For example, Gao et al. [2021] finetune the pretrained BERT model using the InfoNCE loss [Oord et al., 2018], and Zhang et al. [2022] jointly train the MAE loss along with uniformity loss.
3.1 Separate Normalization
We present SepNorm, a normalization scheme that separately normalizes embeddings of the [CLS] symbol and embeddings of other tokens. In this work, we focus on the exploration of combinations of BN and LN for the two separate normalization channels.
For instance, if we apply BN to the [CLS] symbol and LN to other tokens, the learnable parameters are structured as  $g_{1} = (\gamma_{1},\beta_{1})$  and  $g_{2} = (\gamma_{2},\beta_{2})$ . Let  $\mathbf{H}\in \mathbb{R}^{L\times d}$  represent the feature sequence, where  $L$  denotes the sequence length, and  $d$  is the feature dimension. Assume embedding  $\mathbf{H}_0$  in the first position corresponds to the [CLS] symbol. The normalization process is as follows:
$$
\mathbf {H} ^ {\prime} = \left(\operatorname {B N} \left(\mathbf {H} _ {0}; g _ {1}\right), \operatorname {L N} \left(\mathbf {H} _ {1}; g _ {2}\right), \dots , \operatorname {L N} \left(\mathbf {H} _ {L}; g _ {2}\right)\right), \tag {5}
$$
where  $\mathbf{H}'$  denotes the normalized features. We can also run separate normalization with one of the three other combinations:
$$
\mathbf {H} ^ {\prime} = \left(\mathrm {B N} \left(\mathbf {H} _ {0}; g _ {1}\right), \mathrm {B N} \left(\mathbf {H} _ {1}; g _ {2}\right), \dots , \mathrm {B N} \left(\mathbf {H} _ {L}; g _ {2}\right)\right),
$$
$$
\mathbf {H} ^ {\prime} = \left(\mathrm {L N} \left(\mathbf {H} _ {0}; g _ {1}\right), \mathrm {B N} \left(\mathbf {H} _ {1}; g _ {2}\right), \dots , \mathrm {B N} \left(\mathbf {H} _ {L}; g _ {2}\right)\right),
$$
$$
\mathbf {H} ^ {\prime} = \big (\mathrm {L N} (\mathbf {H} _ {0}; g _ {1}), \mathrm {L N} (\mathbf {H} _ {1}; g _ {2}), \ldots , \mathrm {L N} (\mathbf {H} _ {L}; g _ {2}) \big).
$$
Separate normalization allows the [CLS] features to be encoded distinctly from other tokens.
As a comparison, the [CLS] token's embedding and other tokens' embeddings interfere with each other in a shared normalization structure. With ShareNorm, the update directions of the LN parameters  $\{\gamma ,\beta \}$  are primarily driven by the embeddings of normal tokens. Below is the gradient calculation for these parameters,
$$
\frac {\delta \mathcal {L}}{\delta \gamma_ {i}} = \sum_ {l = 1} ^ {L} \frac {\delta \mathcal {L}}{\delta \tilde {\mathbf {H}} _ {l , i}} \tilde {\mathbf {H}} _ {l, i}, \quad \frac {\delta \mathcal {L}}{\delta \beta_ {i}} = \sum_ {l = 1} ^ {L} \frac {\delta \mathcal {L}}{\delta \tilde {\mathbf {H}} _ {l , i}}, \tag {6}
$$
$$
\text {w h e r e} \tilde {\mathbf {H}} _ {l, i} = \frac {\mathbf {H} _ {l , i} - \mu_ {l}}{\sigma_ {l}}, \mu_ {l} = \frac {1}{d} \sum_ {i = 1} ^ {d} \mathbf {H} _ {l, i}, \sigma_ {l} = \sqrt {\frac {1}{d} \sum_ {i = 1} ^ {d} \left(\mathbf {H} _ {l , i} - \mu_ {l}\right) ^ {2}}. \tag {7}
$$
We see the summation in the gradient calculation is dominated by normal tokens given that the number of normal tokens is typically a large number. Given the potentially diverse characteristics (i.e., mean and scale) of feature distributions, it might be challenging for normalization parameters to accommodate both token types simultaneously. Moreover, mapping two types of token features into the same sphere may also mix the signal of [CLS] tokens with other tokens. Figure 2(a, b) demonstrates this phenomenon in the scenario where both token types utilize a ShareNorm and how using SepNorm mitigates this effect.
(a) Shared Normalization (ShareNorm)
(b) Separate Normalization (SepNorm)
Figure 2: The effect of SepNorm on feature distributions. Each subplot shows the distributions of the first 50 feature dimensions: [CLS] features are in blue, and other tokens' features are in red. The [CLS] features of ShareNorm are more concentrated around the mean and the mean deviates more from the zero centers, while in SepNorm, the [CLS] distribution is more centered and flattened.
3.2 Encourage the Uniformity of the [CLS] Embeddings via a Constrastive Term
We further relate SepNorm with the uniformity of embeddings. Higher uniformity values indicate that embeddings can better exploit the space to store information. Contrastive methods often employ negative instances to encourage uniformity. In particular, we incorporate SepNorm into transformers trained with U-MAE [Zhang et al., 2022], which uses a constrastive term to promote uniformity of features.
The U-MAE explicitly adds a uniformity loss term  $\mathcal{L}_{\mathrm{unif}}$  to the training objective to encourage uniformity of [CLS] embeddings.
$$
\mathcal {L} _ {\mathrm {U} - \mathrm {M A E}} = \mathcal {L} _ {\mathrm {M A E}} + \lambda \mathcal {L} _ {\text {u n i f}}, \text {w i t h} \mathcal {L} _ {\text {u n i f}} = \mathbb {E} _ {i} \left[ \mathbb {E} _ {j} \left[ \mathbf {h} _ {\mathrm {C L S}, i} ^ {\top} \mathbf {h} _ {\mathrm {C L S}, j} \right] \right] \tag {8}
$$
Here  $\mathcal{L}_{\mathrm{MAE}}$  is the MAE training objective. The two indices  $i$  and  $j$  represent two sequences within the same batch. [CLS] embeddings  $\mathbf{h}_{\mathrm{CLS},i}$  and  $\mathbf{h}_{\mathrm{CLS},j}$ , which are respectively for the two sequences, are obtained from our SepNorm during the transformer calculation. By minimizing  $\mathcal{L}_{\mathrm{unif}}$ , [CLS] features tend to be different from each other.
We examine the effectiveness of the proposed SepNorm component in three domains: CV, NLP, and graphs. We then further investigate how the ShareNorm and SepNorm affect the uniformity of the [CLS] embeddings.
4.1 Computer Vision
Datasets. We investigate the model performance on the four image datasets: STL10 [Coates et al., 2011], FGVC Aircraft [Maji et al., 2013], Street View House Numbers (SVHN) [Netzer et al., 2011], and Oxford 102 Flowers [Nilsback and Zisserman, 2008]. All four datasets are for classification tasks. We follow the train/test split provided in the papers introducing the datasets. We report top-1 and top-5 accuracy for all datasets.
Vision transformers (ViT) and MAE. We choose Vision Transformer (ViT) [Dosovitskiy et al., 2020] as our feature extractor for all datasets. To pretrain the ViT, we adopt the MAE training scheme [He et al., 2022]. We follow MAE and use a  $75\%$  masking ratio on input image. During the downstream tasks, we use the embeddings of the [CLS] token to predict the class labels.
Experiment setup. We follow the setup in He et al. [2022] to pretrain and evaluate the ViT. For pertaining, we train the ViT for 4000 epochs. For linear probing, we freeze the encoder's weight and train the last layer on the specific datasets for 2000 epochs. We use a batch size of 512 for pretraining and a batch size of 128 for linear probing.
Figure 3: t-SNE visualization of representations learned from the STL-10 dataset.
STL10	Aircraft	SVHN	Flower
ACC@1	ACC@5	ACC@1	ACC@5	ACC@1	ACC@5	ACC@1	ACC@5
MAE	92.01	99.5	52.54	84.16	88.97	99.13	27.63	53.73
+ SepNorm	93.84	99.7	59.02	86.65	89.18	99.21	32.51	60.92
Table 1: Comparison of linear probing performance of ShareNorm and SepNorm across 4 image classification datasets when the  $\mathrm{ViT_{base}}$  is pretrained with MAE.
Experiment results. The results presented in Table 1 demonstrate the performances of our model and the baseline model. Our model consistently outperforms the baseline across multiple datasets, indicating its effectiveness in image classification tasks. In the STL-10 dataset, our approach achieves the top-1 accuracy of  $93.84\%$  and the top-5 accuracy of  $99.7\%$ , higher than the baseline's respective accuracies of  $92.01\%$  and  $99.5\%$ . Similar improvements are observed in the Aircraft, SVHN, and Flower datasets, where our model consistently outperforms the baseline in both top-1 and top-5 accuracies. These results demonstrate the effectiveness of SepNorm in enhancing image classification performance. We also visualize the embeddings of ShareNorm and SepNorm using t-SNE in Figure 3. Compared with ShareNorm, SepNorm provides embeddings that have better separation among different classes.
4.2 Natural Language Processing
Datasets. We evaluated our approach using the STS dataset, which comprises seven semantic textual similarity (STS) tasks. These tasks, including STS 2012-2016 [Agirre et al., 2012, 2013, 2014, 2015, 2016], STS Benchmark [Cer et al., 2017], and SICK-Relatedness [Marelli et al., 2014]. We also evaluate our method on multiple transfer tasks, including MR [Pang and Lee, 2005], CR [Hu and Liu, 2004], SUBJ [Pang and Lee, 2004], MPQA [Wiebe et al., 2005], SST-2 [Socher et al., 2013], TREC Voorhees and Tice [2000], and MRPC [Dolan and Brockett, 2005]. Following the evaluation settings of SimCSE [Gao et al., 2021], we use Spearman's correlation coefficient as the evaluation metric.
BERT and RoBERTa. We conduct our study with pretrained checkpoints of BERT (uncased) [Devlin et al., 2018] and RoBERTa (cased) [Liu et al., 2019], instead of training them from scratch. Using pretrained models is common in this research field [Gao et al., 2021] because the findings are compatible with the common practice of finetuning pretrained models in actual learning tasks. This strategy also saves significant training time and computational resources, allowing us to extend the study to more learning tasks.
Experiment setup. We follow the experiment setup in Gao et al. [2021] and further finetune the BERT and RoBERTa models on English Wikipedia. We evaluate the models using established STS tasks and employ standard evaluation metrics such as Spearman's correlation.
STS12	STS13	STS14	STS15	STS16	STS-B	SICK-R	Avg.
Unsupervised Training
BERTbase	ShareNorm	65.28	78.82	69.65	79.02	77.21	76.4	71.74	74.04
SepNorm	67.01	82.16	72.48	81.38	79.11	77.56	71.36	75.87
RoBERTa base	ShareNorm	68.25	81.24	72.78	81.38	80.31	79.83	68.16	76.00
SepNorm	66.63	82.40	74.47	82.39	80.44	81.14	69.44	76.70
Supervised Training
BERTbase	ShareNorm	77.72	81.07	78.97	85.15	82.00	82.36	79.74	81.00
SepNorm	75.32	84.41	79.94	84.91	80.87	83.63	79.61	81.23
RoBERTa base	ShareNorm	77.38	80.87	78.72	84.02	82.56	83.08	78.25	80.70
SepNorm	75.80	84.94	80.33	85.51	82.11	84.88	79.72	81.90
MR	CR	SUBJ	MPQA	SST2	TREC	MRPC	Avg.
Transfer Learning
BERTbase	ShareNorm	82.78	88.79	94.69	89.86	87.94	84.44	75.99	86.36
SepNorm	82.82	89.08	94.30	89.70	87.97	83.88	75.21	86.14
RoBERTa base	ShareNorm	84.45	91.50	93.94	89.45	90.96	86.80	76.13	87.61
SepNorm	85.11	91.56	94.30	89.43	91.66	90.96	75.58	88.37
Experiment results. The experiment results presented in Table 2 highlight the performance of our model compared to the SimCSE baseline on NLP tasks. With the SepNorm layer, BERT<sub>base</sub> and RoBERT<sub>base</sub> achieve overall higher average accuracy compared to ShareNorm's average accuracy. Only in the transfer learning tasks, SepNorm works slightly worse than ShareNorm in BERT<sub>base</sub>, but the difference is marginal.
4.3 Prediction of Molecule Properties
Datasets. We conducted experiments using the ZINC dataset [Irwin and Shoichet, 2005], which contains approximately 250,000 molecular graphs. The task is to predict the properties of molecules from their graphs. We use a subset of 12,000 molecular graphs, as recommended by the benchmarking methodology outlined in [Dwivedi et al., 2020], so that our results are comparable with other studies. Despite being smaller, the subset retained sufficient diversity and complexity for effective evaluation. We also the MolHIV dataset from the OGB [Hu et al., 2020a] collection, which is widely used for training and evaluating graph-based models in molecular property prediction tasks.
Graphormer. We use Graphormer [Ying et al., 2021] as the transformer backbone to construct the predicting model. To obtain graph-level information, Graphormer adds a special node [VNode] to the graph and connects it to all normal graph nodes. The embedding of [VNode] is a summary of the entire graph and will be used in downstream classification tasks. The special node [VNode] serves the same purpose as the [CLS] token in traditional Transformer models. Graphormer has used three encodings to enhance the transformer's learning ability: centrality encoding captures node importance, spatial encoding considers spatial relations, and edge encoding incorporates edge features.
Experiment setup. We strictly follow Graphormer [Ying et al., 2021] in terms of the model architecture, hyperparameters, and training strategies. We replaced the ShareNorm in Graphormer with SepNorm to investigate the effectiveness of the proposed component. We evaluate the pretrained model on a broad class of graph-level prediction tasks. We report the mean absolute error for
Table 2: Sentence embedding performance on STS tasks and transfer tasks.
Dataset	ZINC	ZINC (subset)	MolHIV
Metrics	Mean absolute error↓	AUC↑
Graphormer	0.069	0.164	73.36%
+ SepNorm	0.052	0.144	75.64%