深入解析CLS Transformer：从核心原理到高效实现

1次阅读

共计 1588 个字符，预计需要花费 4 分钟才能阅读完成。

在传统的 Transformer 架构中，每个 token 都会生成独立的表示。但对于分类或句子级任务，我们需要一个能够代表整个序列的全局表征。CLS（Classification）Token 的设计初衷就是作为这种 ” 句向量 ” 的载体——它像一位会议主持人，在自注意力机制中汇总所有 token 的信息。

普通 Token：专注于局部上下文关系，比如 ”bank” 在 ”river bank” 和 ”bank account” 中会产生不同表示
CLS Token：通过多层 Transformer 的全局注意力，融合整个序列的语义。实验显示，在 BERT 中第 12 层 CLS Token 对分类任务的贡献度比平均池化高 37%

import torch
import torch.nn as nn

class CLSTransformer(nn.Module):
    def __init__(self, vocab_size=10000, d_model=512, nhead=8):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model))  # 可学习的 CLS
        self.transformer = nn.TransformerEncoderLayer(d_model, nhead)

    def forward(self, x):
        # x: [batch_size, seq_len]
        embeddings = self.embedding(x)  # [batch, seq, d_model]
        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.cat([cls_tokens, embeddings], dim=1)  # 添加 CLS 到序列开头

        # 位置编码（简化版）positions = torch.arange(x.size(1)).unsqueeze(0)
        pos_encoding = self.positional_encoding(positions, x.size(-1))
        x = x + pos_encoding

        return self.transformer(x)[:, 0]  # 只返回 CLS 位置的输出

关键点说明：
1. 可学习的 CLS Token 比固定零向量效果更好
2. 位置编码采用 sin/cos 函数的经典方案：

$$PE_{(pos,2i)} = \sin(pos/10000^{2i/d_{model}}))$$
$$PE_{(pos,2i+1)} = \cos(pos/10000^{2i/d_{model}}))$$