深入解析CLS Token在Mamba架构中的关键作用与实现机制

1次阅读

没有评论

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

传统 Transformer 的 CLS Token 作为全局特征聚合器存在两个显著缺陷：

固定位置编码限制：静态位置编码难以适应不同长度的序列
二次方注意力开销：处理长序列时计算成本急剧上升

Mamba 架构通过状态空间模型 (SSM) 引入两个关键改进：

动态状态更新：CLS Token 的状态随序列推进而演化
线性复杂度 ：SSM 的卷积模式实现 O(L) 计算复杂度

Mamba 将 CLS Token 建模为连续状态系统，其演化过程可表示为：

$$
\begin{aligned}
h_t &= A h_{t-1} + B x_t \
y_t &= C h_t
\end{aligned}
$$

其中：

$A \in \mathbb{R}^{N×N}$ 是状态转移矩阵
$B,C \in \mathbb{R}^{N×d}$ 是投影矩阵
$h_t$ 是时刻 t 的隐状态

Mamba 通过门控机制实现权重动态分配：

class DynamicCLS(nn.Module):
    def __init__(self, dim, n_states):
        super().__init__()
        self.dim = dim
        self.A = nn.Parameter(torch.randn(n_states, n_states))
        self.B = nn.Parameter(torch.randn(dim, n_states))
        self.C = nn.Parameter(torch.randn(n_states, dim))
        self.gate = nn.Linear(dim, 1)

    def forward(self, x):  # x: [batch, seq_len, dim]
        batch, seq_len, dim = x.shape
        h = torch.zeros(batch, self.A.shape[0]).to(x.device)
        outputs = []

        for t in range(seq_len):
            # 门控权重计算
            g = torch.sigmoid(self.gate(x[:, t]))
            # 状态更新
            h = (1-g) * h + g * (self.A @ h + self.B @ x[:, t])
            outputs.append(self.C @ h)

        return torch.stack(outputs, dim=1)  # [batch, seq_len, dim]

import torch
import torch.nn as nn

class MambaCLS(nn.Module):
    def __init__(self, vocab_size=30000, dim=512, n_states=64):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim)
        self.dynamic_cls = DynamicCLS(dim, n_states)
        self.proj = nn.Linear(dim, 1)

    def forward(self, x):
        # x: [batch, seq_len]
        x = self.embed(x)  # [batch, seq_len, dim]
        cls_states = self.dynamic_cls(x)  # [batch, seq_len, dim]
        # 取最后时刻状态作为分类特征
        return self.proj(cls_states[:, -1])

在 GLUE 的 MRPC 任务上（Tesla V100 32GB 环境）：