On Separate Normalization in Self-supervised Transformers
Authors: Xiaohui Chen
Yinkai Wang
Yuanqi Du
Soha Hassoun
Liping Liu
Venue: NeurIPS 2023 poster (2023)
Keywords: Transformer
Self-supervised Learning
Normalization
TL;DR: A simple normalization technique for improving a transformer's performance.

Abstract: 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 [CLS] symbol and the tokens. We propose in this paper a simple modification that employs separate normalization layers for the tokens and the [CLS] symbol to better capture their distinct characteristics and enhance downstream task performance. Our method aims to alleviate the potential negative effects of using the same normalization statistics for both token types, which may not be optimally aligned with their individual roles. We empirically show that by utilizing a separate normalization layer, the [CLS] embeddings can better encode the global contextual information and are distributed more uniformly in its anisotropic space. When replacing the conventional normalization layer with the two separate layers, we observe an average  2.7% performance improvement over the image, natural language, and graph domains.

On Separate Normalization in Self-supervised Transformers
Xiaohui Chen Department of Computer Science
Yinkai Wang
Department of Computer Science
Tufts University Medford, MA 02155
Tufts University Medford, MA 02155
xiaohui, chen@tuf edu
yinkai. wang@tuf ts edu
Yuanqi Du
Soha Hassoun
Department of Computer Science
Department of Computer Science
Cornell University Ithaca, NY 14850 yd3928corne11. edu
Tufts University Medford, MA 02155 soha. www.ountuufss....mmmmm edu
Li-Ping Liu Department of Computer Science Tufts University Medford. MA 02155 liping liu@tufts edu
Abstract
Self-supervised training methods for transformers have demonstrated remarkable performance across various domains. Previous transformer- -osoomer--bsedd models, such as masked autoencoders (MAE), typically utilize single normalization layer for both the class token (CT.S) and the tokens. We propose in this paper 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 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.
are distributed more
Introduction
have revolutionized natural language processing 2018 Liu al. 2019 and demonstrated remarkable performances across tasks. The significance of transformer models lies in their ability to model context and capture complex linguistic patterns without being constrained by sequential nature of data. Beyond NLP transformers have further found their successes in areas such computer vision (CV) HHn et al. 2022), speech recognition (Karmakar et al., 202T], 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.
Transformer models (NLP) (Devlin et al.) wide range of NLP
Vaswani et al.
2017
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

SharcNorm
SepNorm
LN
BN
LN
CLS
Tokens
CLS
Tokens
Figure 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 al. 2018]. J and normalization techniques. Specifically, layer normalization [Ba et al. 2016] and batch normalization [loffe and Szegedy [2015] employed to enhance stability and speed up convergence during training. The literature on transformers also explores parameter initialization [Xu ct al. [ optimization algorithms (Huang al. [0200] regularization techniques [Steiner et al. 2021 Zhou et al., (020), and improved architectures [Han 2021]. This collective research has advanced transformer architectures and their applications in al.
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., (2200) Nguyen and Salazar 2019]. For example, Xiong et al. [2020) emphasize 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. [hhen 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 superviedd transformers, the [CLS] symbol is frequently used as global representation for various downstream tasks (Devlin et al.] 2018 He et al.) [022]. 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 special role in representation learning, natural question is whether we should treat separately the normalization operation. Driven by this question, our research first scrutinizes the behavior of the current shared normalization transformers, particularly the properties of the [CLS] embedding and its influence on downstream task performance. Subsequently, Separate Normalization (SepNorm), the latter which employs distinct normalization operations for the [CLS) symbol and the token features, depicted in Figure 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.
propose replacement of ShareNorm with
2 Background
2.1 Pretraining Transformers with the [CLS] symbol
Unsupervised pretraining of a transformer- model (asswan 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 pair of sentences (x,y).

---

which are represented as special sequence
= ([CLS), [SEP].y).
Here [SEP] is special token that separates the sentences. fraction (e.g., 15%) of the tokens in x and are randomly replaced by a special symhol [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 y is the next sentence following x, and the decision is made by classifying the final embedding of the [CLS] symbol. After pretraining, the representation the [CLS] is usually used for sentence-level downstream tasks such as sentiment analysis Meehat et al. 2014
Pretraining MAE for CV. The Vision Transformer (ViT) (Dosoviskkiy et al. 2020 applies the transformers to computer vision tasks. In ViT, an image is usually voxelized into 16 x 16 patches, which are then flattened into 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. training image has 75% of its patches masked. The MAE feeds Tokens of unmasked patches well [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.
2021 is a transformer- based model designed for graph representation learning tasks. It is used to predict the property of a graph rather than node or edge. Specifically, Graphormer introduces a new symbol [VNode] as node connecting 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.
Pretraining Graphormer for molecule discovery. Graphormer Ying et al.
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 al., 2020). LN normalizes across feature dimensions and is independent of the sequence length and the batch size. For any features Rdd the LN has the following computation:
LN(h) O h + 1
(1+ 1)²
Here h; is the i-th dimension of h, represents slement-wiss multiplication, and 7, 3 Rd are scale and bias parameters, respectively. a transformer, all tokens, including special tokens, such as (CLS) and [SEP], are all treated equally and share the same LNs.
Batch Normalization (BN) loffe 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 h as:
BB
BN(h) =
3.
OB
Here (B 2 E Rdd are the running statistics (mean and variance) maintained by the BN. The running mean and variance are updated during training after each batch. They usually calculated as an exponential moving average of the batch mean and variance, BN is widely adopted in CV but leads significant performance degradation when naively used in NLP.
2.3 Uniformity of the Learned Representations
The dimensional collapse in self-supervised representation learning a common phenomenon where the embedding vectors only span 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 simpler representation. Contrastive methods [Oord al. 2018 Chen 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 (h1....,hu b from dataset of size N, the uniformity metric Lu is computed as follows:
N.N
Lu log
exp
hmm
(4)
( 1)/2
m=n+1
If the distribution of the representation is perfectly uniform, then the numerical value of Lu will converge to as the dimension h increases to infinity (Wang and Isola (020).
In self-supervised transformers, the uniformity of the representation is also taken into consideration by some works. For example, Gao et al. [202T] finetune the pretrained BERT model using the InfoNCE loss [Oord al. 2018 and Zhang al. [2022] jointly train the MAE loss along with uniformity loss.
Approach
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, focus the exploration of combinations of BN and LN for the two separate normalization channels.
For instance, if we apply BN to the (CT.S symbol and LN to other tokens, the learnable parameters are structured as 91 = (71, (1) and 92 (72,, 32). Let H E Lxd represent the feature sequence, where L denotes the sequence length, and d is the feature dimension. Assume embedding Ho in the first position corresponds to the [CLS) symbol. The normalization process is as follows:
H' = (BN(Ho::) 91), LN(H:; 92)..... LN(H,,
(5)
where H' denotes the normalized features. We can also run separate normalization with one of the three other combinations:
H' = (BN (Ho: 91). ( (N Hz:g)))
H' = (H 9, BN (::,)....,,NNN
H1 Hz:g2))
H' = (LN(Ho:g1) L,, LN)
Separate normalization allows the [CLS) features to encoded distinctly from other tokens.
As comparison, the (CLS) token's embedding and other tokens' embeddings interfere with each other in shared normalization structure. With ShareNorm, the update directions of the LN parameters {r, B} are primarily driven by the embeddings of normal tokens. Below is the gradient calculation for these parameters,
8.C 8%/
8L
8L
8C
Σ 1=1
H
Σ
(6)
8H
83,
HH
d
d H -m))
H
P4
1
where H
d
H
H1,,
=
(7)
1=1
1=1
We see the summation in the gradient calculation is dominated by normal tokens given that the number of normal tokens typically 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 b) demonstrates this phenomenon in the scenario where both token types utilize ShareNorm and how using SepNorm mitigates this effect.

---

(a) Shared Normalization (ShareNorm)
(b) Separate Normalization (SepNorm)
The effect of SepNorm on feature distributions. Each subplot shows the distributions of the first 50 feature dimensions: [CLS] features in blue, and other tokens' features 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 [CI.S] distribution is more centered and flattened.
Figure 2:
3.2 Encourage the Uniformity of the (CLS) Embeddings via Constrastive Term
We further relate SepNorm with the uniformity of embeddings. Higher uniformity values indicate that embeddings better exploit the space store information. Contrastive methods often employ negative instances to encourage uniformity. In particular, we incorporate SepNorm into transformers trained with U-MAE (Zhang al. 2022]. which uses constrastive term to promote uniformity of features.
The U-MAE explicitly adds uniformity loss term Lunit to the training objective encourage uniformity of [CT.] embeddings.
LU-MAE LMAE + ALuuitt with Luiif E, E, h CCLS, heLS
Here CMAE the MAE training objective. The two indices and j represent two sequences within the same batch. [CLS] embeddings hCLs, and hets.J. which are respectively for the two sequences, are obtained from our SepNorm during the transformer calculation. By minimizing Lunifs (CL.S) features tend to be different from each other.
4 Experiments
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 CCLS) embeddings.
4.1 Computer Vision
Datasets. We investigate the model performance the four image datasets: STL10 (Coates et al.
2011], FGVC Aircraft [Maji et al.] 2013)] Street View House Numbers (SVHN) [Netzer et al. [201] and Oxford 102 Flowers NNlssbck and Zisserman, 2008). All four datasets 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) DDsoviiskii 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 75% masking ratio on input image. During the downstream tasks, we use the embeddings of the [CLS] token to predict the class labels.
follow the setup He 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 batch size of 512 for pretraining and batch size of 128 for linear probing.
Experiment setup.
5

---

ShareNorm(LN)
SepNorm
Figure 3: t-SNE visualization of representations learned from the STL-10 dataset.
STL10
Aircraft ACC@
SVHN
Flower
ACC@1
ACc@5
ACC@5
ACC@1
ACC@5
ACC@
ACC@5
MAE SepNorm
52.54
84.16
88.97
99.13
27.63 32.51
53.73
93.84
99.7
59.02
86.65
89.18
99.21
60.92
Table 1: Comparison of linear probing performance of ShareNorm and SepNorm across image classification datasets when the base is pretrained with MAE.
Experiment results. The results presented in Table 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.8 3.84% and the top-5 accuracy of 99.7%. higher than the baseline's '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 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 Figure Compared with ShareNorm, SepNorm provides embeddings that have better separation among different classes.
4.2 Natural Language Processing
Datasets. We evaluated 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 al. [017]. and Relatedessss [Marelli al. 2014 We also evaluate our method on multiple transfer tasks, including MR (Pang and Lee 2005]. CR [uu and Liu 2004]. SUBJ (Pang and Lee 2004]. MPQA [Wibbe al. 2005]. SST-2 Socher al. [2013].
TREC Voorhees and Tice [2000]. and MRPC (Dolan and Brocketi) 2005]. Following the evaluation settings of SimCSE [Gao 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) (De vlin 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 compatible with the common practice of finetuning pretrained models actual learning tasks. This strategy also saves significant training time and computational resources, allowing extend the study to more learning tasks.
Experiment setup. We follow the experiment setup in Gao et al. 202 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 Spearman's 's correlation.

---

STS12
STS13
STS14
STSI5
STS16
STS-B
SICK-R
Avg.
Unsupervised Training
BERT
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
RoBERTapus
ShareNorm SepNorm
68.25
81.24
72.78
81.38
80.31
79.83
68.16
76.00
66.63
82.40
74.47
82.39
80.44
81.14
69.44
76.70
Supervised Training
BERT base
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
RoBERTataas Bipase
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
SUBJ
MPQA
SST2
TREC
MRPC
Avg.
Transfer .earming
BERT hose
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
RoBERTapue
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
Table 2: Sentence embedding performance STS tasks and transfer tasks.
Experiment results. The experiment results presented in Table Sable@hight highlight the performance of our model compared to the SimCSE baseline NLP tasks. With the SepNorm layer, BERT base and RoBERTa achieve overall higher average accuracy compared to ShareNorm's 's average accuracy. Only in the transfer learning tasks, SepNorm works slightly worse than ShareNorm BERT base. the difference is marginal.
4.3 Prediction of Molecule Properties
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 subset of 12,000 molecular graphs, as recommended by the benchmarking methodology outlined in [Dwivedi al. 2020), that results 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.
Datasets.
Graphormer. We Graphormer YYig al. 2021 as the transformer backbone construet the predicting model. To obtain graph-level information, Graphormer adds special node [VNode] the graph and connects it to all normal graph nodes. The embedding of [VNode] is 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 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 the ZINC and ZINC (subset) datasets and the area under the curve (AUC) for the MolHIV dataset.
Dataset
ZINC
ZINC (subset)
MolHIV
Metrics
Mean absolute error],
AUCt
Graphormer SepNorm 0.052
0.069
0.164 0.144
73.36%
75.64
Table 3: A comparison of ShareNorm and SepNorm in three tasks of graph property prediction.
7

---

Recostruction loss (MSE)
0.57
Log of singulur valie
EN
LN
BN
BN-BN BN-LN
0.55
RNHLN
LS-NN
5.5
LN+IN
Log of singulur
6.0
IN-EN
RNSUN
5.0
5.5
0.53
4.5
3.0
4.0
3.5
0.49
3.0
2k
Training epochs
The index of sigular value
The index of sigular value
(a) Pretraining loss
(b) [CLS] embeddings
(c) normal token embeddings
Figure 4: (a) Reconstruction loss of the MAE pertaining MAE with SepNorm achieves lower MSE loss compared to ShareNorm, demonstrating better ability encode global contextual information. (b) & (c) Comparison of the singular values of learned (CCLS) 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.
Experiment results. Table the performances of our model and the Graphormer baseline. For the ZINC datasets, Graphormer with SepNorm achieves 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 STLI0 dataset via four different losses:
MAE loss LMAR without any Lumif (CLS) and token embeddings. This setting is study with MAE training only.
MAE loss [MAE with Lunif on the [CLS] embeddings. We treat all [CLS] embeddings within the same batch (except itself) as negative instances.
MAE loss [MAE with Luiit on the token embeddings. We treat all token embeddings within the same batch same images (except itself) negative instances.
MAE loss LMAE with Lunif on both [CLS) and token embeddings.
We choose A {0. 0.01, 0.1, 13. Note that the second loss with 0.1 corresponds to the U- MAE (Zhang et al. 2022]. We also replace the normalization layer of the ViT in MAE with one of the following: (L.N, BN, BN+LN, BN+BN). The combination of different losses, different X's, and different normalization layers yields 40 specifications of the experiments.
We first report our results with MAE training only. The uniformiit of learned embeddings is first measured by singular values of the decomposition of an embedding matrix: we randomly choose 10k embeddings 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 the STL.10 and Aircraft datasets [Coates et al. 2011 Maji 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 that the uniformity of [CLS] embeddings 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.

---

94
STL10
Aircraft
BN BN
Norml.ayer
[CL.S]
token
[CLS)
token
LN
-2.5911
-3,4998
-0,0977
-3,3514
Accuracy
93
LN
BN
-2.9531
-3.6785
-15321
-3,3659
LN + LN
-1.1045
-3.3685
-1,6874
-3,3633
LN + BN
-1.3153
-3.3735
-1,9634
-3,5056
92
LN
LN+ BN
LN LN
BN + LN
-3,4426
-3,3797
-2.7705
-3.3068
BN BN
-3,4763
-3.4758
-2.9788
-3.7349
-3.5
-2.5
-1.5
-0.5
Uniformity
(a) Uniformity of different normalization settings.
(b) Uniformity vs. accuracy
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 STLI0 dataset.
We then report results from studies with U-MAE training. Figure shows the uniformity met- rics obtained using different X's. When using ShareNorm, the uniformity of the [CLS] embed- dings is no better than -3,088, and even the explicit uniformity loss does not help much. On the con- trary, embeddings learned from the proposed Sep- Norm 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.
BN
LN-LN
BN-LN
Uniformity
0.01
0.05
Figure 5: vs uniformity. A higher À empha- sizes more the uniformity of [CLS] embed- dings.
The results provide strong evidence that the uni- formity of the [CLS] embeddings is held down by ShareNorm even the minimization of an ex- plicit contrastive loss cannot increase it. We 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.
4ablort 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 (A o 0), SepNorms' advantages are further highlighted. For example, when À 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 SepNorm, the performance of ShareNorm is less satisfactory when using contrastive methods. We believe 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 (BN+LN) gaining accuracy with increasing À 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 much information possible. Our empirical study also shows that, when contrastive loss [Oord et al. 2018 is used to encourage the uniformity of [CLS] features in superviee transformers, the difference between BN and LN on [CLS] features is not significant anymore.
5 Related Works
The training of transformer architectures with self-supervised learning has seen significant advance- ments in both contrastive and non-contrastive training. Among self-supervised learning methods,

---

Normalization layer
A
Negative pairs
X 0.01
À 0.1
À
Best
token
93.65
94.15
93.94
94.15
SepNorm (BN+BN)
93.84
CLS
93.73
93.85
93.93
93.93
CLS + token
93.40
94.25
94.28
94.28
token
92.98
93.60
94.30
94.30
SepNorm (BN+LN)
92.80
CCSS
92.98
93.48
93.36
93.48
CLS) token
92.74
93.18
94,40
94,40
token
92.48
93.38
92.78
93.38
ShareNorm (BN)
92.84
[CLS]
93.10
93.33
92.93
93.33
CLS token
93.41
93.46
92.99
93.46
token
92.61
92.74
92.14
92.74
ShareNorm (LN)
92.01
[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 varied. We report downstream task accuracy for the STL10 dataset.
non-contrastive ones do not rely negative samples for learning, They have emerged 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 al. a 2020bb]
Contrastive methods on the contrary train networks using positive and negative samples that are con- structed without manual labeling. They have also been used to train transformer-basee architectures. Gao ct 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. Mcanwhile, the potential of contrastive methods in vision transformers has been demonstrated by CCrrnn et al. [022]] and Radford et al. [222]] These collective efforts underscore the versatility and efficacy of contrastive methods in self self- supe-sueeeiieed learning of transformers.
Normalization layers, including layer normalization and batch normalization, are essential to trans- former 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 transform- ers. Similarly, Xu et al. (2099b 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, [henn et al. [2020] introduces new normalization layer, Powernorm, which a variant of batch normalization. Nguyen and Salazar (019 focus on the normalization process in the self-atted aff-attention mechanism of transformers and propose methods to optimize the normalization of self attentinno All the efforts above underscore the critical role normalization layers transformer models.
6 Conclusion
In this work, we have introduced SepNorm to separate the normalization of (CT.S) embeddings from that of other tokens. Across three application domains (images, text, and graphs), SepNorm shows consistent performance improvement when it incorporated into transformer models. Our analysis shows that SepNorm promotes uniformity of [CLS] embeddings and thus enhances the transformers' ability to encode information. As valuable technique for improving the foundational transformer architecture, SepNorm has the potential benefit wide range of applications.
Acknowledgement
We thank all reviewers for their constructive feedback. This research supported by the NIGMS of the National Institutes of Health, Awards 335GM118821 the Army Research Office, MURI program, contract # W91 INF2210239, and NSF Award 1909536. Chen and Liu are also supported by the NSF CAREER Award # 2239869.
10

---

References
E. Agirre, D. Cer, M. Diab, and A. Gonzalez-Agirre Semeval-2012 task 6: A pilot on semantic textual similarity. In SEM 2012: The First Joint Conference on Lexical and Computational PemaaticcsVVllueee I: Proceedings of the main conference and the shared task, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation (SemEval 2012), pages 385-393, 2012.
E. Agirre, D. Cer, M. Diab, A. Gonzalez-Agirre, and W. Guo. * sem 2013 shared task: Semantic textual similarity. In Second conference on lexical and computational semantics (* SEM), volume 1: proceedings of the Main conference and the shared task: semantic textual similarity. pages 52-43,22 2013.
E. Agirre, C. Banea, C. Cardie, D. Cer, M. Diab, Gonzalez-Agiree W. Guo, R. Mihalcea, G. Rigau. and J. Wiebe. Semeval-2014 task 10: Multilingual semantic textual similarity. In Proceedings of the 8th international workshop on semantic evaluation (SemEval 2014), pages 81-91, 2014.
E. Agirre, C. Banea, C. Cardie, D. Cer, M. Diab, A. Gonzalez-Agiree, W. Guo, I. Lopez- Gazpi,, M. Maritxalar, R. Mihalcea, et al. Semeval-2015 task 2: Semantic textual similarity, english, spanish and pilot on interpretabilltyy. In Proceedings of the 9th international workshop on semantic evaluation (SemEval 2015), pages 252-263,, 2015.
E. Agirre, C. Banea, D. Cer, M. Diab. A. Gonzalez Agirre, R. Mihalcea, G. Rigau Claramunt, and J. Wiebe. Semeval-2011 task I: Semantic textual similarity, monolingual and cross-lingual evaluation. In SemEval-2016. 10th International Workshop Semantic Evaluation; 2016 Jun 16- 17; San Diego, CA. Stroudsburg ACL; 2016. p. 497-511. ACL (Association for Computational Linguistics), 2016.
J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv:1607.06455 2016.
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages (550-9660) 2021.
D. Cer, M. Diab, E. Agirre, I. Lopez-Gazpio, and L. Specia. Semeval-2017 task 1: Semantic textual comilaritymmuliilingn and cross-lingual focused evaluation. arXiv preprint arXiv:1708.0005: 2017.
M. Chen, A. Radford, R. Child, J. Wu, H. Jun, D. Luan, and Sutskever. Generative pretraining from pixels. In International conference on machine learning, pages 1691-1703, PMLR, 2020a.
T. Chen, Kornblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020b.
X. Chen, S. Xie, and K. He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9640-9649, 2021.
A. Coates, A. Ng, and H. Lee. An analysis of single-layer networks unsupervised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 215-223. JMLR Workshop and Conference Proceedings, 2011.
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv 2018
B. Dolan and C. Brockett. Automatically constructing corpus of sentential paraphrases. In Third International Workshop on Paraphrusing (/WP2005), 2005.
A. Dosovitskiy, L.. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint ar arXiv:2010.11929 2020.

---

V. P. Dwivedi. C. K. Joshi, T. Laurent, Y. Bengio, and X. Bresson, Benchmarking graph neural networks. 2020.
Gao, X. Yao, and D. Chen. Simese: Simple contrastive learning of sentence embeddings. arXiv preprint arXiv:2104.088// 2021.
Y. Gu, Z. Ding, S. Wang, L. Zou, Y. Liu, and D. Yin. Deep multifaceted transformers for multi- objective ranking large-scale e-commerce recommender systems. In Proceedings the the 29th ACM International Conference on Information & Knowledge Management, pages 2493-2500, 2020.
K. Han, Xiao, Wu, J. Guo, C. Xu, and Y. Wang. Transformer transformer. Advances in Neural Information Processing Systems, 1-:15988--9111 2021.
K. Han, Y. Wang. H. Chen, X. Chen, J. Guo, Z. Liu, Y. Tang, A. Xiao, C. Xu, Y. Xu, et al. A survey on vision transformer. IEEE transactions on pattern analysis and muchine intelligence, 45(1): 87-110, 2022.
K. He, X. Chen, Xie, Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. Proceedings of the IEEE/CVE Conference on Computer Vision and Puttern Recognition, pages 16000-16009, 2022.
M. Hu and B. Liu. Mining and summarizing customer reviews. In Proceedings of the tenth ACM SIGKDD international conference on Knowledge discovery and data mining, pages 168-177, 2004.
W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovee. Open graph benehmark: Datasets for machine learning on graphs. Advances neural information processing systems, 03:22118-22133 2020a.
Z. Hu, Y. Dong, K. Wang, and Y. Sun. Heterogeneous graph transformer. In Proceedings of the web conference 2020, pages 2704-2710, 2020b.
X. S. Huang, Perez, Ba, and M. Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learning, pages 4475-4483, PMLR, 2020.
S. loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pages 448-456. pmir, 2015.
J. J. Irwin and B. K. Shoichet. Zinc- a free database of commercially available compounds for virtual screening. Journal of chemical information and modeling, 157 2005.
L. Jing, P. Vincent, Y. LeCun, and Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:21 10.09348, 2021.
P. Karmakar, S. W. Teng, and G. Lu. Thank you for attention: a survey attention artificial neural networks for automatic speech recognition. arXiv preprint arXiw:2102.07259, 2021.
Y. Liu, M. Out, Goyal, Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L.. Zettlemoyer, and V. Stoy- anov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
Maji, E. Rahtu, J. Kannala, M. Blaschko, and Vedaldi. Fine-grained visual classification of aircraft. arXiv preprint arXiv: 306.5151,
M. Marelli, S. Menini, M. Baroni, L. Bentivogli, R. Bernardi, R. Zamparelli, et al. A sick cure for the evaluation of compositional distributional semantic models. In Lrec, pages 216-223. Reykjavik, 2014.
W. Medhat, A. Hassan, and H. Korashy. Sentiment analysis algorithms and applications: A survey, Ain Shams engineering journal, 104) 2014.
Y. Netzer, T. Wang, A. Coates, Bissacco, B. Wu, and A. Ng. Reading digits natural images with unsupervised feature learning. 2.11.
12

---

T. Q. Nguyen and J. Salazar. Transformers without tears: Improving the normalization of self- attention. arXiv preprint 1r1.v:1910005895, 2019.
M.-E. Nilsback and A. Zisserman. Automated flower classification over large number of classes, In 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing, pages 722-772-777 229. IEEE, 2008.
A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 2018.
B. Pang and L. Lee. A sentimental education: Sentiment analysis using subjectivity summarization based minimum cuts. arXiv preprint cx/440958,, 2004.
B. Pang and L. Lee. Seeing stars: respect to rating scales. arXiv preprint cs/0506075, 2005.
Exploiting class relationships for sentiment categorization with
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748-8763. PMLR, 2021.
Shen, Z. Yao, A. Gholami, M. Mahoney, and Keutzer. Powernorm: Rethinking batch normaliza- tion in transformers. In International Conference Machine Learning, pages 8741-8751. PMLR, 2020.
R. Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Y. Ng, and C. Potts. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 conference on empirical methods natural language processing, pages 1631-1642, 2013.
A. Steiner, A. Kolesnikov, X. Zhai, Wightman, J. Uszkoreit, and Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint 102.7000002000 2021.
Sun, Liu, J. Wu, C. Pei. X. Lin, W. Ou, and P. Jiang. Bert4rec: Sequential recommendation with bidirectional encoder representations from transformer. In Proceedings of the 28th ACM international conference on information and knowledge management, pages 1441-1450 2019.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
E. M. Voorhees and D. M. Tice. Building a question answering test collection. In Proceedings of the 23rd annual international ACM SIGIR conference on Research and development in information retrieval, pages 200-207, 2000.
T. Wang and P. Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pages 9929- 9939. PMLR, 2020.
Y. Wang, A. Mohamed, D. Le, C. Liu, A. Xiao, J. Mahadeokar, H. Huang, A. Tjandra, X. Zhang, F. Zhang, et al. Transformer-based acoustic modeling for hybrid speech recognition. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). pages 6874-6878. IEEE, 2020.
J. Wiebe, T. Wilson, and C. Cardie. Annotating expressions of opinions and emotions in language. Language resources and evaluation, 39:165-210, 2005.
L. Wu, S. Li, C.-J. Hsich, and Sharpnack. Sse-pt: Sequential recommendation via personalized transformer. In Proceedings of the 14th ACM Conference Recommender Systems, pages 328-337, 2020.
R. Xiong, Y. Yang, D. He, K. Zheng, S. Zheng, C. Xing, H. Zhang, Y. Lan, L. Wang, and T. Liu. On layer normalization in the transformer architecture. In International Conference Machine Learning, pages 10524-10533, PMLR, 2020.
H. Xu, Q Liu, J. van Genabith, D. Xiong, and Zhang. Lipschitz constrained parameter initialization for deep transformers. arXiv preprint
2019a.
13

---

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, 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:22/0.08344, 2022.
W. Zhou, T. Ge, K. Xu, F. Wei, and M. Zhou, Scheduled drophead: regularization method for transformer models. arXiv preprint arXiv:2004. 13344, 2020.
14


Reviews:
[{"review_id": "UHDmHKyToj", "reviewer": "Reviewer_voLG", "summary": "This paper analyzes the normalization layer that is applied to both the context tokens and the [cls] token in an input. It argues that traditional normalzation that is applied to both type of tokens, dubbed ShareNorm, is not effective since they have different roles. The paper proposes to use separate normalization layers for these tokens, SepNorm, and conducted a series of analysis. Echoing prior work, they showed that a measure of uniformity of the representations correlates with downstream classification performance, and further that ShareNorm find hard to tackle this uniformity deficiency even with uniformity fixing methods, while SepNorm improves uniformity and hence downstream task performances.", "questions": "n/a", "limitations": "no limitations were discussed.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well written and easy to understand. It starts with sufficient background and explains the analysis clear.\n2. The hypothesis intuitively makes sense to me, and the empirical analysis provides good support.\n3. Experiments are conducted on a number of tasks in vision, nlp and graphs.", "weaknesses": "I don't have major concerns on this work. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are sincerely grateful for your time and effort in reviewing our submission. We greatly appreacited your positive feedback on our proposed method, and we are glad to know that you found the paper well-written, easy to understand. Once again, thank you for your kind words and for accepting the paper."}, {"review_id": "IrAjZfuB9N", "reviewer": "Reviewer_QRJ8", "summary": "The paper proposes a new method called SepNorm, which utilizes separate normalization layers for the [CLS] token and the remaining tokens, replacing the conventional single normalization layer (ShareNorm) in transformers. Experiments demonstrate the importance of applying separate normalization to the [CLS] token and the remaining tokens, and show that SepNorm enhances the transformer's ability to encode the input.", "questions": "1. Can the proposed method be scaled to large language models?\n\n2. How is the speed of SepNorm in comparison to ShareNorm?\n", "limitations": "The paper does not discuss the limitations of the proposed method. It would be better to discuss the potential impact of the proposed normalization method on learned representations, training time in addition to the performance on the downstream tasks.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper propose to explore separate normalization for the [CLS] token in pretrained transformers. Experiments on three different domains  show the effectiveness of the proposed method.", "weaknesses": "The paper mentions that \"Our method aims to alleviate the potential negative effects of using the same normalization statistics for both token types, which may not be optimally aligned with their individual roles.\" It seems that this statement is not empirically validated in the paper. \n\nIn the ablation study, only experiments on a CV dataset have been conducted. It would be better if more experiments on NLP and Molecule Discovery datasets could be carried out.\n\nThe paper only selects one or two models for each domain in the experiments, such as evaluating the STS task in the NLP domain. However, this limited selection of tasks is not sufficient to demonstrate the generalizability of SepNorm. It would be more appropriate to include a wider range of experiments across multiple tasks to provide a comprehensive evaluation of the effectiveness and generalizability of the proposed SepNorm method.\n\nThe paper does not empirically compare the proposed method against previous works, such as Powernorm mentioned in the section of related work.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## [weakness]\n**Q1: the statment \"Our method aims to alleviate the potential negative effects of using the same normalization statistics for both token types, which may not be optimally aligned with their individual roles.\" is not empirically validated in the paper.**\n\nR1: The potential negative effect we refer to is **the relatively low uniformity score achieved by using ShareNorm**. Uniformity has been widely used to assess the degree of dimension collapse and is often considered an essential metric for learned representation quality [1,2,3]. We quantitatively showed that the uniformity for ShareNorm is consistently worse than SepNorm, regardless of the emphasis that is placed on the metric. We will modify our statement accordingly to avoid such confusion.\n\n**Q2: It would be better if more experiments on NLP and Molecule Discovery datasets could be carried out.**\n\nR2: Thank you for your comment. As you suggested, we expanded our experiments to cover a broader range of NLP datasets. Due to time constraints, we are not able to explore on Molecule tasks. We will consider incorporating more Molecule datasets in the future. Below, we brief how we further validate our proposed method in various NLP datasets/tasks.\n* Models: We use Bert and RoBERTa models.\n* Input variants:\nWe finetune BERT and RoBERTa on multiple datasets using three input settings, using either the [CLS] or [MASK] tokens for sentence-level tasks. For [MASK] predictions, an actual word is predicted, and its semantic meaning determines the label, e.g., \"terrible\" indicates a negative label.\n  1. Standard: Input is \"[CLS] <Sentence> [EOS]\", we predict labels via [CLS] embeddings.\n  2. Prompt-based: Input is \"[CLS] <Sentence>. This is [MASK] [EOS]\", we use [MASK] embeddings for predictions.\n  3. Prompt-based with demo: Input is \"[CLS] <Sentence1>. This is [MASK]. [SEP] <demo1>. This is <label1>. [SEP] <demo2>. This is <label2> [EOS]\", we use [MASK] embeddings for predictions. Demos are drawn from training data. \n* Datasets:\nWe use the following datasets: Stanford Sentiment Treebank 2&5 (SST-2 & SST-5), Movie Reviews(MR), Customer Reviews(CR), Multi-Perspective Question Answering(MPQA), Subjectivity Dataset(SUBJ), Corpus of Linguistic Acceptability(CoLA), Text REtrieval Conference(TREC), Stanford Natural Language Inference(SNLI), Question Natural Language Inference(QNLI), Microsoft Research Paraphrase Corpus(MRPC), Quora Question Pairs(QQP). \n* Results: We report the result in the uploaded pdf in the general response. As we can observe, applying SepNorm on [CLS] can help improve the model performance with different prompting methods. \n\nWe hope the additional experiment can address the reviewer’s concern about our method’s effectiveness. \n\n**Q3: Include a wider range of experiments across multiple tasks to comprehensively evaluate the effectiveness and generalizability of SepNorm.**\n\nR3: Thank you for the feedback. In response, we've extended our evaluation to multiple transfer tasks, including MR, CR, SUBJ, MPQA, SST-2, TREC, and MRPC. To do this, we freeze the model weights learned from SimCSE and only train the classifier on top of the model. We kindly refer the reviewer to the attached pdf for more details on the experiment results. \n\n**Q4: The paper does not empirically compare the proposed method against previous works, such as Powernorm...**\n\nR4: We’d like to point out that we do not propose a new normalization layer. Instead, we propose a normalization strategy that can be combined with different normalizations such as BatchNorm, LayerNorm, and PowerNorm. Since there are many combinations of different normalization layers in SepNorm, we simplify the setting and focus on the canonical case.\n## [questions]\n\n**Q1: Can the proposed method be scaled to large language models?**\n\nR1: Our method can be applied to large language models that utilize the [CLS] tokens. For example, CLIP aligns the [CLS] embeddings from image and text via contrastive learning. SepNorm can be incorporated into the text and vision transformers.\n\n**Q2: How is the speed of SepNorm in comparison to ShareNorm?**\n\nR2: With SepNorm, the model needs no more computation but only stores separate mean and variances with a little extra space. In practice, the computation cost of replacing ShareNorm with SepNorm is negligible. For example, for training an MAE on STL-10 dataset in the machine with one A100 GPU with 32 CPU cores, both normalization methods took 36 hours to finish.\n## [limitations]\n**Q1: It would be better to discuss the potential impact of the proposed normalization method on learned representations, training time in addition to the performance on the downstream tasks.**\n\nR1: Thank you for the feedback. \n1. Learned Representations: We quantitatively evaluated representation uniformity, noting that better uniformity indicates better data information preservation[1]. Our qualitative analysis in Figure 2(a,b) shows that SepNorm results in a more uniform [CLS] feature distribution around a mean of 0. Moreover, the slower singular value decay in Figure 2(d,c), suggests that SepNorm can better span the feature space, leading to better representation.\n2. Training time: The replacement of SepNorm does NOT incur extra computation cost. The extra computation here is to perform another normalization operation over the feature, which is affordable compared to the self-attention modules.\n3. Performance on the downstream tasks: We verify SepNorm on downstream tasks from three domains. For example, we pretrain the transformers with ShareNorm/SepNorm in the CV and graph domains and finetune them on the corresponding downstream tasks (results are reported in Table 1, 3 in Section 5). \n## [Reference]\n[1]Wang, et al. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. ICML 2020.\n\n[2]Bojanowski, et al. Unsupervised learning by predicting noise. ICML 2017.\n\n[3]Mettes, et al. Hyperspherical prototype networks. NIPS 2019.\n"}, {"review_id": "DM4TVt1cjJ", "reviewer": "Reviewer_vuG3", "summary": "The authors propose to use a different normaliser to the [CLS] token and the rest of the tokens for masked autoencoders. \nThey motivate this as an improvement on the standard normalisation and combine it with contrastive uniformity loss. \nThey observe some improvements in classification tasks .", "questions": "While the [CLS] token stands among others,  it looks like this can be an extreme of different normalisation depending on the frequency. \nHave authors compared how much of the improvement comes from the token frequency vs how much comes from the semantic full sentence encoding of the token ?", "limitations": "The authors have studied some ablation studies.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The pape is proposing a new normalisation of MAE using 2 separated normalisers one for [CLS] token and another for other tokens. \nThey provide evidence this helps the models to obtain better perfomance indepndently and in combination with uniformity loss.", "weaknesses": "The contribution while important doesn't seem relevant for a larger audience but they do per discussion below for several tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your effort in reviewer our submission. Below we answer the concerns and questions that you raised, please feel free to ask if  you have further questions.\n\n## [weakness]\n**Q1: The contribution while important doesn't seem relevant for a larger audience.**\n\nR1: Thank you for your feedback. We feel that our proposed method can actually benefit a broad class of audiences in the machine learning community. Our proposed plug-and-play component, SepNorm, has the potential to benefit diverse domains. Our experiment demonstrated that it can be easily applied to various transformer-based architectures (e.g., MAE, BERT, Graphormer) used in different domains (CV, NLP, Graph). By incorporating SepNorm, researchers, and practitioners can improve the performance of these architectures without extensive modifications. We will enhance the clarity and presentation of our findings in the revised version of the paper.\n\n## [questions]\n**Q1: While the [CLS] token stands among others, it looks like this can be an extreme of different normalisation depending on the frequency. Have authors compared how much of the improvement comes from the token frequency vs how much comes from the semantic full sentence encoding of the token?**\n\nR1: Regarding the \"frequency\" aspect, are you referring to the occurrence frequency of individual words within the corpus? To rephrase your question, are you asking whether it's necessary to normalize the most frequently appearing word separately or if the special token requires its own normalization? Our response is that the unique [CLS] token necessitates distinctive normalization for two primary reasons. First, the vector representation of the [CLS] token is directly employed in downstream tasks. Second, this vector effectively encodes the comprehensive sentence semantics, desiring a distinct treatment. We are open to providing further clarification if our interpretation of your question is not accurate.\n"}, {"review_id": "oJpAD2a6ys", "reviewer": "Reviewer_hvnd", "summary": "The paper introduces SepNorm, a normalization technique for Transformer models. SepNorm separates the normalization of the [CLS] token from the rest of the tokens in a sequence, a departure from the traditional ShareNorm method that normalizes all tokens together. ", "questions": "Would SepNorm be helpful for multilingual machine translation tasks?", "limitations": "No visualization of the learned representation from before and after applying SepNorm", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. By applying SepNorm, the embeddings of the [CLS] symbol effectively captures the characteristics of the entire graph, leading to improved results in downstream applications.\n2. SepNorm achieves better uniformity by encouraging uniformity for both the [CLS] and token embeddings when used in contrastive methods.\n3. SepNorm seems to do well in the evaluation.", "weaknesses": "No uncertainty/confidence/error bars on experimental results or significance testing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the suggestion on our paper, below we address your concern accordingly.\n\n## [weakness] \n**Q1: No uncertainty/confidence/error bars on experimental results or significance testing.**\n\nR1: As suggested, we re-run the experiments to obtain the standard derivations in the NLP tasks as it is faster to train compared to the tasks in CV and graph domains. Moreover, we've also included the experiment results from the supervised training setting. We kindly ask the reviewer to check the results in the attached PDF. \n\n## [Questions]\n**Q1: Would SepNorm be helpful for multilingual machine translation tasks?**\n\nR1: Thank you for the insightful question! SepNorm can be better used in discriminative models rather than generative models. BERT originally proposes [CLS] token and is usually used for non-autoregressive transformers (the encoder part in [1]). We didn’t further investigate its use case in decoder-based transformers. Since machine translation usually requires an encoder-decoder model to perform conversion between two languages, it may not be an ideal scenario for applying SepNorm. However, SepNorm can be used for multilingual-related tasks such as alignment [2]. We will further discuss the suitable scenarios for using SepNorm in the future revised manuscript. \n\n\n## [Limitations]\n**Q1: No visualization of the learned representation from before and after applying SepNorm.**\n\nR1: Thank you for your feedback. In response, we’ve visualized the learned representation of the STL-10 dataset for both ShareNorm and SepNorm. We use t-SNE to reduce the dimension from 784 to 2 before visualization. Class label of each data point are highlighted with different colors. We kindly refer the viewer to the attached PDF for the visualization result.\n\n## [References]\n\n[1] Vaswani, Ashish, et al. Attention is all you need. NIPS 2017\n\n[2] Cao, Steven, et al. Multilingual alignment of contextual word representations. ICLR 2020"}, {"review_id": "qc3gqI33m8", "reviewer": "Reviewer_4DDo", "summary": "The paper proposes a simple modification to Transformer-based self-supervised learning by having distinct normalization (parameters) for the CLS token, which captures the global information for use in downstream tasks, and the remaining tokens. The motivation is that this allows the tokens to avoid dimension collapse as indicated by a lower uniformity score, which in turn correlates with higher downstream performance. The paper verifies in a case study that this is empirically the case.\nThe proposed method, SepNorm, is evaluated in a controlled setting against the standard approach, ShareNorm, on various datasets from different domains. SepNorm consistently achieves improvements. An additional ablation study confirms that this is also the case when uniformity is encouraged explicitly through an additional loss term.", "questions": "* In section 4, alternative normalization methods are discussed. In how far do these address the dimensional collapse problem? How would your results change with different normalizations, i.e., in how far is your method orthogonal to these?\n* Table 2: It is unclear what you mean by BERT-base and RoBERTa-base. Do these refer to the original models, or are these finetuned in the same way as the \"+SepNorm\" models? If it's the former, how do we know that the improvement isn't merely a result of continued finetuning?", "limitations": "The paper does not address limitations. I think it would be useful to discuss whether the results would extend to other types of normalization (see question above), and what potential avenues for future work are.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* The paper proposes an extremely simple, yet consistently effective solution to mitigate the problem of dimensional collapse and thereby improve performance.\n* The paper is concerned with general self-supervised Transformer architectures, which are relevant to the majority of the NeurIPS audience.\n* The experiments are very clear and well executed, the controlled experimental setup leaves little room for doubt regarding the method's merits.\n* The paper is generally self-contained, all necessary pieces can be understood by a general ML-savvy audience.\n* The idea to use separate sets of learnable parameters across different components of the model is of course not new per se, but the application context and the resulting insights constitute sufficient novelty.", "weaknesses": "The only main weakness is that the paper could be structured better. Sections 3.1. and 3.2 mix the description of the proposed method with experimental results that are supposed to motivate the method. However, those results already rely on the existence of the proposed method, and do thus not suit as motivation. Rather, they constitute useful analyses of the proposed method, which could just as well be discussed in the \"Experiments\" section. This has the added benefit that the datasets used for generating results in Figures 3 and 4 have already been introduced (which only happens in 5.1). In fact, section 3.1. and 3.2. never explicitly mention that these are computer vision datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer's positive feedback on our paper, as well as the suggestions that can improve our submission's quality. Below we try our best to address those concerns and questions from the reviewer.\n\n## [Weakness]\n**Q1: The only main weakness is that the paper could be structured better. Sections 3.1. and 3.2 mix the description of the proposed method with experimental results that are supposed to motivate the method. However, those results already rely on the existence of the proposed method, and do thus not suit as motivation. Rather, they constitute useful analyses of the proposed method, which could just as well be discussed in the \"Experiments\" section. This has the added benefit that the datasets used for generating results in Figures 3 and 4 have already been introduced (which only happens in 5.1). In fact, section 3.1. and 3.2. never explicitly mention that these are computer vision datasets.**\n\nR1:Thank you for your feedback! We will restructure the methodology section as suggested to enhance the clarity and coherence of our paper. We believe your suggestion will improve the readability of our work. \n\n\n## [Questions]\n**Q1: In section 4, alternative normalization methods are discussed. In how far do these address the dimensional collapse problem? How would your results change with different normalizations, i.e., in how far is your method orthogonal to these?**\n\nR1: Thank you for your insightful question. Our method is entirely distinct from previous normalization approaches in terms of both its motivation and methodology. Furthermore, our method can be seamlessly integrated with any normalization approach. As far as we know, the alternative normalization methods discussed in Section 4 do not specifically aim to tackle the dimensional collapse problem; instead, they primarily focus on stabilizing the training process of transformers.\n\nTo better understand how our method can complement alternative normalization techniques, consider the following insights:\n\n1. **For non-contrastive self-supervised training**, certain normalization methods like BatchNorm and PowerNorm can also promote uniformity. If PowerNorm demonstrates better training process stabilization than BatchNorm, we recommend using PowerNorm.\n2. **For contrastive self-supervised training**, it’s empirically shown that LayerNorm may perform better than BatchNorm [1] as it does not utilize batch statistics. This observation may also extend to PowerNorm.\n\nIt is worth noting that we have not discussed other normalization methods such as GroupNorm and PairNorm since they are not explicitly designed for transformers.\n\nWe hope that this response addresses the reviewer's concerns. Please feel free to inquire about any other questions you may have.\n\n\n**Q2: Table 2: It is unclear what you mean by BERT-base and RoBERTa-base. Do these refer to the original models, or are these finetuned in the same way as the \"+SepNorm\" models? If it's the former, how do we know that the improvement isn't merely a result of continued finetuning?**\n\nR2: We apologize for the confusion, “+SepNorm\" denotes BERT-base with SepNorm, and BERT-base denotes BERT-base with ShareNorm. Both models are finetuned in the same way to make a fair comparison. We have corrected the table in the attached PDF, all the results will be included in the revised script.\n\n## [Limitations]\n**Q1: The paper does not address limitations. I think it would be useful to discuss whether the results would extend to other types of normalization (see question above), and what potential avenues for future work are.**\n\nR1: We thank the reviewer for the insightful questions. A common limitation of such innovations of network layers is the increased complexities of model tuning: there are more hyperparameters to set. However, we hope the extensive benefit shown in this paper will help the community to settle for a default choice of SepNorm. We will add a paragraph discussing the limitation of the proposed method. \n\nIn terms of extending to other types of normalization, we kindly refer the reviewer to [Questions]-Q1 for the discussion of pairing with other types of normalization. \n\nRegarding potential future works, we provide two possible avenues in the following:\n\n1. Relationship between [CLS] and non-[CLS] tokens, and potential architecture improvement:\nA promising area for future work is exploring the relationship between the [CLS] token and non-[CLS] tokens. Our ablation study indicated that achieving better uniformity in the non-[CLS] tokens could positively impact downstream performance. Additionally, it remains unknown whether dimensional collapse issues exist in token-level tasks such as generation, object detection, etc. If there are any, we suggest investigating whether improvements to transformer designs can address these issues and further enhance performance.\n\n2. Extension to decoder-only large language models:\nWhile our current work focuses on encoder-only transformers, we recognize the growing popularity of large language models utilizing decoder-only transformers. To gain a comprehensive understanding of representation distributions and the role of uniformity in generative tasks, it is essential to extend our analysis to these decoder-only architectures. Exploring how uniformity influences the performance of generative tasks in decoder-only transformers may also provide valuable insights into their behavior.\n\nWe will revise our manuscript to incorporate those discussions suggested by the reviewer. Thank you for the valuable comments.\n"}]

Author rebuttal:
We thank all reviewers for their effort in reviewing our submission. And we appreciate all the positive and negative feedbacks on our manuscript. We've tried our best to address all the concerns and questions raised by the reivewers. The attached pdf includes the additional experiments and visualization asked by reviewer QRJ8 and hvnd. Thanks again for the valuable feedbacks.

Meta-review: The paper proposes a simple and straightforward modification to Transformer-based self-supervised learning by having distinct normalization (parameters) for the CLS token, which captures the global information for use in downstream tasks, and the remaining tokens. The proposed method, SepNorm, is evaluated in a controlled setting against the standard approach, ShareNorm, on various datasets from different domains.

After the response, the authors added more experiments on NLP tasks, which addressed the concerns of reviewers.
Decision: Accept (poster)