architecture-details

# Mamba 架构详情 ## 选择性状态空间机制 Mamba 的核心创新是 **选择性 SSM (S6)** 层，它使状态空间模型的参数具有输入依赖性。 ### S6 的工作原理 **传统 SSM** (非选择性)： python # 对所有输入使用固定的 A, B, C 矩阵 h(t) = A * h(t-1) + B * x(t) # 状态更新 y(t) = C * h(t) # 输出 **Mamba 的选择性 SSM**： python # 输入依赖参数 B(t) = Linear_B(x(t)) # 选择机制 C(t) = Linear_C(x(t)) # 输出投影 Δ(t) = Linear_Δ(x(t)) # 离散化步长 # 选择性状态更新 h(t) = discretize(A, Δ(t)) * h(t-1) + Δ(t) * B(t) * x(t) y(t) = C(t) * h(t) ### 关键优势 **1. 基于内容的推理**： - 可以根据输入选择性地记忆或遗忘 - 解决了传统 SSM 在离散模态上的弱点 - 示例：记住重要 token，遗忘填充 token **2. 输入依赖的选择**： python # Mamba 决定每个 token 需要记住什么 if is_important(x(t)): Δ(t) = large_value # 保留在状态中 else: Δ(t) = small_value # 快速遗忘 **3. 无需注意力机制**： - 用 O(n) 的状态更新取代了 O(n²) 的注意力机制 - 状态维度保持恒定（通常为 16） ## 模型配置 ### 核心参数 python from mamba_ssm import Mamba model = Mamba( d_model=256, # 隐藏层维度 (256, 512, 768, 1024, 2048) d_state=16, # SSM 状态维度 (固定为 16 是最优的) d_conv=4, # 局部卷积宽度 (4 是标准值) expand=2, # 扩展因子 (1.5-2.0) dt_rank="auto", # Δ 投影的秩 (auto = d_model / 16) dt_min=0.001, # 最小 Δ 初始化 (控制遗忘率) dt_max=0.1, # 最大 Δ 初始化 dt_init="random", # Δ 初始化 (random, constant) dt_scale=1.0, # Δ 缩放因子 conv_bias=True, # 卷积中使用偏置 bias=False # 线性投影中使用偏置 ) ### 参数影响 **d_state** (SSM 状态维度)： - 标准：16 (消融实验得出的最优值) - 更小 (8)：更快但容量较低 - 更大 (32, 64)：提升微小，速度慢 2 倍 **expand** (块扩展)： - 标准：2.0 - 范围：1.5-2.0 - 控制内部维度 = expand * d_model **d_conv** (卷积宽度)： - 标准：4 - SSM 之前的局部上下文窗口 - 有助于位置信息 **dt_rank** (Δ 投影秩)： - Auto：d_model / 16 (推荐) - 控制 Δ 参数效率 - 秩越低 = 更高效