本文最后也有部分代码,更详细的请参考官方链接https://github.com/tommyip/mamba2-minimal/blob/main/mamba2.py这段代码实现的是 Mamba-2 模型的核心算法 Structured State Space Duality (SSD),用于高效地处理序列数据。让我们来逐步拆解这段代码,理解其背后的原理和运作机制。
1. SSD 算法概述
-
SSD 是一种基于状态空间模型的序列处理方法,其核心思想是将序列分解成若干个块 (chunk),并在块内和块间进行高效的信息传递。
-
SSD 利用了矩阵的低秩分解和指数衰减特性,将原本复杂的序列建模问题转化为一系列高效的矩阵乘法运算,从而显著降低了计算复杂度。
2. 代码解析
segsum 函数解析:
-
功能: 计算一个特殊的累积和,用于生成一个 1-半可分矩阵(1-semiseparable matrix),该矩阵等效于一个标量状态空间模型 (SSM)。
-
输入:
-
x: 输入张量。
-
device: 计算设备。
-
-
步骤:
-
复制扩展: 将输入 x 沿着最后一个维度复制 T 次,生成一个新的张量。
-
生成掩码: 创建两个下三角掩码矩阵,分别用于控制累积和的范围。
-
计算累积和: 使用 torch.cumsum 函数计算累积和,并利用掩码矩阵控制计算范围。
-
填充负无穷: 将不在计算范围内的元素填充为负无穷,这是为了在后续计算指数时将其置零。
-
SSD 函数解析:
-
输入:
-
x: 输入序列,形状为 (batch, seqlen, n_heads, d_head)。
-
A: 控制状态转移的矩阵,形状为 (batch, seqlen, n_heads)。
-
B: 将输入映射到状态空间的矩阵,形状为 (batch, seqlen, n_heads, d_state)。
-
C: 将状态空间映射到输出的矩阵,形状为 (batch, seqlen, n_heads, d_state)。
-
chunk_size: 将序列分割成的块的大小。
-
initial_states: 初始状态,可选。
-
device: 计算设备,例如 CPU 或 GPU。
-
-
步骤:
-
数据重排 (Rearrange into chunks):
-
将输入序列、A、B、C 矩阵按照 chunk_size 分割成若干个块。
-
代码中使用了 rearrange 函数进行高效的数据重排操作。
-
-
块内计算 (Intra-chunk computation):
-
计算每个块内的输出 Y_diag:
-
L = torch.exp(segsum(A, device=device)) 计算状态转移矩阵的累积效应。
-
Y_diag 通过 C, B, L 和 x 的矩阵乘法得到。
-
-
计算每个块内的最终状态 states:
-
decay_states 计算状态的衰减情况。
-
states 通过 B, decay_states 和 x 的矩阵乘法得到。
-
-
-
块间循环 (Inter-chunk recurrence):
-
使用 initial_states 初始化状态,或使用默认的零向量。
-
decay_chunk 计算块间的衰减情况。
-
new_states 通过 decay_chunk 和 states 的矩阵乘法得到,表示更新后的状态。
-
final_state 保存最后一个块的最终状态。
-
-
状态到输出的转换 (State-to-output conversion):
-
state_decay_out 计算状态衰减对输出的影响。
-
Y_off 通过 C, states 和 state_decay_out 的矩阵乘法得到,表示块间信息传递对输出的贡献。
-
-
输出合并 (Output combination):
-
将块内输出 Y_diag 和块间输出 Y_off 相加,得到最终的输出 Y。
-
-
-
输出:
-
Y: 处理后的输出序列,形状为 (batch, seqlen, n_heads, d_head)。
-
final_state: 最后一个块的最终状态。
-
3. 代码亮点:
-
高效的矩阵运算: 代码大量使用了 torch.einsum 函数,这是一种高效的矩阵乘法运算方法,可以充分利用硬件资源,加速计算。
-
并行计算: 代码中提到,步骤 1、2 和 4 可以并行计算,这为进一步提升性能提供了空间。
总结:
-
segsum 函数是一个辅助函数,用于生成 SSD 算法中所需的特殊矩阵。
-
ssd 函数实现了 SSD 算法,通过将序列分解成块并利用矩阵的低秩分解和指数衰减特性,实现了高效的序列建模。
def segsum(x: Tensor, device: Device = None) -> Tensor:
"""Stable segment sum calculation.
`exp(segsum(A))` produces a 1-semiseparable matrix, which is equivalent to a scalar SSM.
Source: https://github.com/state-spaces/mamba/blob/219f03c840d5a44e7d42e4e728134834fddccf45/mamba_ssm/modules/ssd_minimal.py#L23-L32
"""
T = x.size(-1)
x = repeat(x, "... d -> ... d e", e=T)
mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=-1)
x = x.masked_fill(~mask, 0)
x_segsum = torch.cumsum(x, dim=-2)
mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=0)
x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
return x_segsum
def ssd(x, A, B, C, chunk_size, initial_states=None, device: Device = None):
"""Structed State Space Duality (SSD) - the core of Mamba-2
This is almost the exact same minimal SSD code from the blog post.
Arguments
x: (batch, seqlen, n_heads, d_head)
A: (batch, seqlen, n_heads)
B: (batch, seqlen, n_heads, d_state)
C: (batch, seqlen, n_heads, d_state)
Return
y: (batch, seqlen, n_heads, d_head)
Source
1. https://tridao.me/blog/2024/mamba2-part3-algorithm/
2. https://github.com/state-spaces/mamba/blob/219f03c840d5a44e7d42e4e728134834fddccf45/mamba_ssm/modules/ssd_minimal.py#L34-L78
"""
assert x.shape[1] % chunk_size == 0
# Rearrange into chunks
# Step 1, 2 and 4 of SSD can be computed in parallel for each chunk across devices (sequence parallel)
# This is not implemented and left as an exercise for the reader 😜
x, A, B, C = [
rearrange(m, "b (c l) ... -> b c l ...", l=chunk_size) for m in (x, A, B, C)
]
A = rearrange(A, "b c l h -> b h c l")
A_cumsum = torch.cumsum(A, dim=-1)
# 1. Compute the output for each intra-chunk (diagonal blocks)
L = torch.exp(segsum(A, device=device))
Y_diag = torch.einsum("bclhn, bcshn, bhcls, bcshp -> bclhp", C, B, L, x)
# 2. Compute the state for each intra-chunk
# (right term of low-rank factorization of off-diagonal blocks; B terms)
decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
states = torch.einsum("bclhn, bhcl, bclhp -> bchpn", B, decay_states, x)
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
# (middle term of factorization of off-diag blocks; A terms)
if initial_states is None:
initial_states = torch.zeros_like(states[:, :1])
states = torch.cat([initial_states, states], dim=1)
decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0)), device=device))
new_states = torch.einsum("bhzc, bchpn -> bzhpn", decay_chunk, states)
states, final_state = new_states[:, :-1], new_states[:, -1]
# 4. Compute state -> output conversion per chunk
# (left term of low-rank factorization of off-diagonal blocks; C terms)
state_decay_out = torch.exp(A_cumsum)
Y_off = torch.einsum("bclhn, bchpn, bhcl -> bclhp", C, states, state_decay_out)
# Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
Y = rearrange(Y_diag + Y_off, "b c l h p -> b (c l) h p")
return Y, final_state
最后附上关于mamba2块缝合到你的模型中的缝合方法链接