深入解析NLP中的[cls] token：从BERT到实际应用的技术实现

1次阅读

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

在 Transformer 架构中，[CLS]（Classification）token 是一个特殊的标记，最初被设计用于分类任务。它的位置固定在输入序列的最前面，通过自注意力机制聚合整个序列的信息，最终输出一个固定长度的向量表示，用于下游任务如文本分类、情感分析等。

[CLS] token 的设计初衷是为了解决变长输入序列的统一表示问题。传统的序列模型如 RNN 或 LSTM 通过最后一个隐藏状态来表示整个序列，但 Transformer 没有这种顺序结构，因此需要一个明确的标记来承担这一角色。

位置与初始化：在 BERT 等模型中，[CLS] token 被添加在每个输入序列的开头。它的初始嵌入由三个部分组成：
Token 嵌入：一个特殊的标记，表示分类任务
位置嵌入：位置 0 的嵌入向量
段嵌入（如果使用）：通常为段 A 的嵌入
注意力机制中的处理：
在自注意力层中，[CLS] token 可以关注序列中的所有其他 token
通过多头注意力机制，它能捕获不同层次的语义信息
最终输出包含了整个序列的全局表示
输出表示：
最后一层的[CLS] token 隐藏状态通常作为整个序列的表示
这个 768 维的向量（在 BERT-base 中）会被送入分类头进行预测

下面是一个使用 PyTorch 和 HuggingFace Transformers 库实现文本分类的完整示例：

from transformers import BertTokenizer, BertForSequenceClassification
from transformers import AdamW
import torch
from torch.utils.data import Dataset, DataLoader

# 1. 数据准备
class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,  # 自动添加 [CLS] 和[SEP]
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )

        return {'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'label': torch.tensor(label, dtype=torch.long)
        }

# 2. 模型初始化
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=2  # 二分类任务
)

# 3. 训练循环
def train(model, data_loader, optimizer, device, epochs):
    model = model.to(device)
    model.train()

    for epoch in range(epochs):
        for batch in data_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['label'].to(device)

            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels
            )

            loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            print(f'Epoch: {epoch}, Loss: {loss.item()}')

# 4. 推理示例
def predict(text, model, tokenizer, device, max_len=128):
    encoding = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=max_len,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt'
    )

    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)

    logits = outputs.logits
    probabilities = torch.softmax(logits, dim=1)

    return probabilities.cpu().numpy()