深入解析class token在Transformer架构中的核心作用与实现机制

1次阅读

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

在视觉 Transformer（ViT）模型中，class token 是一个特殊的设计，它承担着全局特征聚合和分类信号载体的双重职责。与 CNN 不同，Transformer 处理的是序列化的图像块（patches），缺乏天然的全局视角。class token 就像是一个虚拟的 ” 观察者 ”，通过自注意力机制与所有图像块交互，最终输出作为整个图像的分类特征。

计算效率：class token 通过单次注意力计算完成特征聚合，而 CNN 的全局平均池化（GAP）是纯线性操作。当特征图较大时，GAP 计算量更小
位置敏感性：GAP 完全丢失空间信息，class token 通过位置编码保留相对位置关系
特征表达能力：class token 通过注意力权重动态调整各区域贡献，比 GAP 的固定权重更灵活

import torch
import torch.nn as nn

class ClassToken(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # 可学习的 class token 参数
        self.token = nn.Parameter(torch.randn(1, 1, dim))  # [1, 1, D]

    def forward(self, x):
        # x 形状: [B, N, D]
        b, n, _ = x.shape
        # 扩展 class token 到 batch 维度
        cls_tokens = self.token.expand(b, -1, -1)  # [B, 1, D]
        # 拼接序列首部
        return torch.cat([cls_tokens, x], dim=1)  # [B, N+1, D]

标准 ViT 做法：class token 与图像块共享相同的位置编码矩阵
改进方案：为 class token 单独设计位置编码（如零初始化）

独立 token：每个任务使用不同的 class token（内存消耗大）
共享 token：所有任务共用基础特征，最后接任务特定头（推荐）

# 使用梯度检查点技术
from torch.utils.checkpoint import checkpoint

class MemoryEfficientAttention(nn.Module):
    def forward(self, x):
        return checkpoint(self._attention, x)

    def _attention(self, x):
        # 实现注意力计算
        return x @ x.transpose(-2, -1)

对 class token 的输出做 LayerNorm 后再进入分类头
设置梯度缩放器避免下溢

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    output = model(inputs)
    loss = criterion(output, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()