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

1次阅读

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

在视觉 Transformer（ViT）模型中，class token 是一个特殊的设计，它的主要作用是作为全局特征的聚合点和分类信号的载体。与 CNN 不同，Transformer 处理的是序列化的图像块（patches），缺乏天然的全局视角。class token 的引入，使得模型能够学习到一个综合所有图像块信息的全局表示，从而更好地完成分类任务。

CNN 通过全局平均池化（GAP）对特征图进行降维，计算复杂度为 O(HWC)，其中 H、W 是特征图的高和宽，C 是通道数。
class token 的计算复杂度为 O(N*D)，其中 N 是序列长度（包括 class token），D 是嵌入维度。对于大分辨率图像，class token 通常更高效。

CNN 的 GAP 对位置信息不敏感，可能导致空间信息的丢失。
class token 通过与所有图像块交互，可以保留更多的空间关系信息。

import torch
import torch.nn as nn

class ClassToken(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        # 初始化一个可学习的 class token
        self.token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # 位置编码，假设已经实现
        self.pos_embed = PositionalEncoding(embed_dim)

    def forward(self, x):
        # x 的形状: (batch_size, num_patches, embed_dim)
        batch_size = x.shape[0]
        # 扩展 class token 到 batch size
        cls_tokens = self.token.expand(batch_size, -1, -1)
        # 将 class token 与图像块拼接
        x = torch.cat((cls_tokens, x), dim=1)
        # 添加位置编码
        x = self.pos_embed(x)
        return x

零初始化：初始时 class token 不携带任何信息，依赖训练过程中从图像块学习。收敛较慢，但稳定性高。
随机初始化：初始时 class token 带有随机信息，可能加速早期训练，但需要小心调整学习率以避免不稳定。

在多任务学习中，可以为每个任务分配独立的 class token，或者共享一个 class token 但通过不同的投影头（projection head）生成任务特定的特征。

对于长序列输入（如高分辨率图像），可以通过以下方法优化内存：

使用梯度检查点技术（gradient checkpointing）减少中间激活的内存占用。
采用混合精度训练，利用 FP16 加速计算并节省内存。

from torch.utils.checkpoint import checkpoint

# 在 forward 函数中使用 checkpoint
x = checkpoint(self.transformer_block, x)

使用 torch.cuda.amp 进行自动混合精度训练，避免数值不稳定：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    output = model(input)
    loss = criterion(output, target)

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