超越 Transformer：Mamba/SSM/RWKV 架构深度对比

h'(t) = A·h(t) + B·x(t)      # 状态方程
y(t)  = C·h(t) + D·x(t)      # 输出方程

h_t = Ā·h_{t-1} + B̄·x_t
y_t = C·h_t

# S4 的 HiPPO 矩阵初始化（简化版）
import numpy as np

def hippo_matrix(N):
    """HiPPO-LegS 矩阵"""
    A = np.zeros((N, N))
    for n in range(N):
        for k in range(n):
            A[n, k] = (2*n + 1)**0.5 * (2*k + 1)**0.5 / (n + 1)
        A[n, n] = -(n + 1)
    return A

class MambaLayer(nn.Module):
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = d_model * expand
        self.in_proj = nn.Linear(d_model, d_inner * 2)  # 投影到 x, z
        
        self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv, 
                                 padding=d_conv-1, groups=d_inner)
        
        # SSM 参数：B, C, Δ 都是输入依赖的
        self.x_proj = nn.Linear(d_inner, d_state * 2 + 1)  # B, C, Δ
        self.dt_proj = nn.Linear(1, d_inner)
        
        # A 是可学习的但不直接依赖输入
        A = torch.arange(1, d_state + 1).float().repeat(d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))
        self.D = nn.Parameter(torch.ones(d_inner))
        
        self.out_proj = nn.Linear(d_inner, d_model)
    
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        xz = self.in_proj(x)  # (batch, seq, 2*d_inner)
        x, z = xz.chunk(2, dim=-1)
        
        x = x.transpose(1, 2)  # (batch, d_inner, seq)
        x = self.conv1d(x)[:, :, :x.shape[2]].transpose(1, 2)
        x = F.silu(x)
        
        # 选择性 SSM：Δ, B, C 依赖输入
        dt_B_C = self.x_proj(x)  # (batch, seq, 2*d_state + 1)
        dt, B, C = dt_B_C.split([1, d_state, d_state], dim=-1)
        dt = F.softplus(self.dt_proj(dt))  # 确保 Δ > 0
        
        # 离散化 A
        A = -torch.exp(self.A_log)  # (d_inner, d_state)
        dA = torch.exp(A * dt)  # (batch, seq, d_inner, d_state)
        
        # 递推计算隐状态
        h = torch.zeros(x.shape[0], d_inner, d_state, device=x.device)
        ys = []
        for t in range(x.shape[1]):
            h = dA[:, t] * h + B[:, t].unsqueeze(1) * x[:, t].unsqueeze(-1)
            y = (C[:, t].unsqueeze(1) * h).sum(-1)
            ys.append(y)
        
        y = torch.stack(ys, dim=1)
        y = y * F.silu(z)  # 门控
        return self.out_proj(y)

# RWKV-7 的 time-mixing（简化）
w_t = sigmoid(可学习参数)  # 衰减因子
k_t = linear(x_t)          # key
v_t = linear(x_t)          # value
r_t = linear(x_t)          # receptance (gate)

# 状态递推（RNN 模式）
kv_t = w_t * kv_{t-1} + k_t * v_t  # 累积 KV
k_sum_t = w_t * k_sum_{t-1} + k_t   # 累积 K

# 输出
y_t = sigmoid(r_t) * (kv_t / k_sum_t)

# RWKV 并行训练（伪代码）
def rwkv_parallel(x, w, k, v, r):
    # x: (batch, seq, d)
    # 使用 parallel scan 计算累积
    log_w = torch.log(w)  # 转换到 log 空间
    cum_log_w = torch.cumsum(log_w, dim=1)  # 前缀和
    
    # 利用 log-space 累积避免数值溢出
    log_kv = log_softmax(cum_log_w + torch.log(k * v), dim=1)
    log_k  = log_softmax(cum_log_w + torch.log(k), dim=1)
    
    y = torch.exp(log_kv - log_k) * v  # 归一化输出
    return sigmoid(r) * y

Jamba Block = [Mamba Layer × N] + [Self-Attention × 1] + [MLP × 1]

Zamba Block = [Mamba2 Layer × N] + [Shared GQA × 1]