Why Bert transformer uses CLS token for classification instead of average over all tokens?
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
The BERT (Bidirectional Encoder Representations from Transformers) model revolutionized natural language processing tasks by introducing a new way to pre-train language representations. Unlike previous models that predominantly use unidirectional—or in some cases, shallow bidirectional—approaches, BERT uses deep bidirectional representations, allowing it to set new state-of-the-art performance on many challenging natural language understanding tasks.
One of the novel aspects of BERT is its way of handling inputs for classification tasks using the [CLS]
token. Instead of averaging over all encoded tokens, BERT uses the representation of the [CLS]
token for classification. This design choice has technical motivations and implications, which this article aims to elucidate.
What is a Transformer and How Does BERT Work?
Before diving into the [CLS]
token, it is important to understand how transformers and BERT work.
Transformers
Transformers are a type of deep learning architecture introduced in the paper "Attention is All You Need" by Vaswani et al. in 2017. This architecture relies entirely on self-attention mechanisms to draw global dependencies between input and output, offering improved computational efficiency by allowing parallelization.
BERT Architecture
BERT builds on the transformer architecture and consists of:
- A stack of transformer encoder layers.
- Inputs are processed in parallel and transformed through these layers.
- BERT is designed to pre-train deep bidirectional representations by jointly conditioning on both left and right context in all layers.
BERT involves two training phases:
- Pre-Training: In which it is trained on a large corpus of text by predicting missing (masked) words and the next sentence.
- Fine-Tuning: The pre-trained model is then fine-tuned on downstream tasks like sentence classification or named entity recognition.
Why Use [CLS]
Token for Classification?
Purpose of Special Tokens
The [CLS]
token is a special token added at the beginning of every input sequence. Its embedding is used as an aggregate representation of the entire sequence, tailored for classification tasks.
The Technical Motivation
- Focus on Aggregate Meaning: In contrast to averaging, using a dedicated token allows BERT to focus on capturing an aggregate representation specifically tailored for summarizing the entire input sequence context. The model learns during pre-training to utilize this token as a summary representation, suitable for downstream tasks.
- Directed Learning: During fine-tuning, weights of layers adapt such that all relevant semantic information needed for classification is concentrated in the
[CLS]token. This focused learning allows[CLS]to act as a more computationally salient summary of the input than could be achieved by averaging. - Dimensional Stability: Averaging embeddings introduces the risk of dimming individual words’ distinctive features, particularly in cases involving long sequences or specific token importance. Instead, the
[CLS]token maintains dimensional stability by focusing on the learned representations.
Example: Sentiment Classification
Consider a sentiment classification task where the goal is to determine whether a given review is positive or negative. During fine-tuning, the transformation through BERT's layers will focus the necessary details and sentiment nuances into the [CLS]
token. This directed attention avoids information dilution, which might occur if an average of all token embeddings were used.
Key Differences with Average Pooling
To summarize and contrast, let's consider the main distinctions between using the [CLS]
token and averaging token embeddings for classification:
| Method | Approach | Advantages & Drawbacks |
**[CLS] | ||
| Token** | Uses the special token at the start of the sequence. | + Tailored for capturing sequence meaning. + More adaptable during fine-tuning. - Dependent on pre-trained weights. |
| Averaging | Computes mean across all token embeddings. | + Simple to compute. + Aggregates broad meaning. - Prone to information loss. - May not capture salient details effectively. |
Additional Observations
Performance Implications
Fine-tuning efficacy usually increases with the use of [CLS]
because it offers a more concentrated and learned aggregation point for task-specific information, shortening the gap between pre-training and task-specific performance.
Research Comparisons
Despite BERT's success with [CLS]
, alternative aggregation methods, such as pooling or concatenation of representations, are also being explored for different architectures. These could potentially mitigate some limitations of a single representation bottleneck and encourage new model adaptations.
Closing Thoughts
The design choice to use a [CLS]
token in BERT exemplifies a paradigm shift in how we aggregate and interpret text representations for classification tasks. By focusing on a dedicated token, the model can leverage nuanced, context-rich information, which would otherwise be diffused through averaging techniques. This strategic focus contributes to the robustness and accuracy seen in BERT's performance across multiple NLP challenges.

