双向SSM: Vision Mamba Encoder

Vision Mamba Encoder

Vision Mamba的编码器部分，也位于Vim模型的中间和主要部分。由多个Mamba块堆叠而成，VisionMamba的Mamba块是在原始论文MambaBlock上修改，特别的地方在于其双向SSM机制。双向与数据流动方向无关，并不是指网络中存在反馈回路，而是等价的扫描方向有两种。

初始化

输入映射

首先，还是一个标准的输入映射，这一点没有更改，输入映射用来得到门控变量z和主干变量x，其中x的维度d_model扩充到2 * d_inner。

self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs)

序列变换

通过一个1D卷积进行序列变换。

        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            bias=conv_bias,
            kernel_size=d_conv,
            groups=self.d_inner,
            padding=d_conv - 1,
            **factory_kwargs,
        )

        self.activation = "silu"
        self.act = nn.SiLU()

参数映射

参数映射是一个简单的线性映射，为了得到输入依赖的矩阵参数B，C还有$\Delta$参数

BC参数映射

d_state*2属于B,C参数，dt_rank属于delta参数的原始维度

 self.x_proj = nn.Linear(
            self.d_inner, self.dt_rank + self.d_state * 2, bias=False, **factory_kwargs
        )

delta参数映射

delta参数的给出：x->x_proj ->split -> dt_proj ->delta

输入x经过x_proj映射得到数据依赖的三个参数$B,C,\Delta$，其中$\Delta$ 得到的维度是dt_rank，还需要进行一个(dt_rank, d_inner)的线性映射

self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True, **factory_kwargs)

SSM参数初始化

在这里初始化非输入依赖的SSM参数包括A矩阵和D矩阵，还包括步长delta参数dt的初始化

A , D矩阵初始化

参数	维度
A	[d_state] -> [d_inner, d_state]
D	[d_inner]

		A = repeat(
            torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
            "n -> d n",
            d=self.d_inner,
        ).contiguous()
        A_log = torch.log(A)  # Keep A_log in fp32
        self.A_log = nn.Parameter(A_log)
        self.A_log._no_weight_decay = True

        # D "skip" parameter
        self.D = nn.Parameter(torch.ones(self.d_inner, device=device))  # Keep in fp32
        self.D._no_weight_decay = True

delta参数初始化

$$dt=e^{\alpha (\cdot log(dt\_max)-log(dt\_min))+log(dt\_min)}$$

其中$\alpha$属于0到1的均匀分布，因此$dt$的取值为$e^{logdt\_min}$到$e^{logdt\_max}$。即$dt\_min$到$dt\_max$

softplus函数为$Softplus(x)=\frac{1}{\beta }\ast log(1+exp(\beta \ast x))$​

	 # Initialize special dt projection to preserve variance at initialization
        dt_init_std = self.dt_rank**-0.5 * dt_scale
        if dt_init == "constant":
            nn.init.constant_(self.dt_proj.weight, dt_init_std)
        elif dt_init == "random":
            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
        else:
            raise NotImplementedError
    
    # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max
        dt = torch.exp(
            torch.rand(self.d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
            + math.log(dt_min)
        ).clamp(min=dt_init_floor)
        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        with torch.no_grad():
            self.dt_proj.bias.copy_(inv_dt)
        # Our initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
        self.dt_proj.bias._no_reinit = True

双向SSM初始化

参数初始化

对于标准Mamba块来说，仅限于前向分支，而后向分支是不存在的，可以看到后向分支是前向分支的复制。在初始化阶段，双向SSM只是额外定义并初始化了一个A矩阵名为A_b。对于v1版本仅仅是多初始化一个矩阵A，而v2版本除此之外，还初始化了标准Mamba所需的全部参数，如D矩阵，参数映射。简单来说，v1版本的双向SSM除A矩阵以外，其他参数是公用的。

参数	维度
A_b	[d_state] -> [d_inner, d_state]

        # bidirectional
        if bimamba_type == "v1":
            A_b = repeat(
                torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
                "n -> d n",
                d=self.d_inner,
            ).contiguous()
            A_b_log = torch.log(A_b)  # Keep A_b_log in fp32
            self.A_b_log = nn.Parameter(A_b_log)
            self.A_b_log._no_weight_decay = True
        elif bimamba_type == "v2":
            A_b = repeat(
                torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
                "n -> d n",
                d=self.d_inner,
            ).contiguous()
            A_b_log = torch.log(A_b)  # Keep A_b_log in fp32
            self.A_b_log = nn.Parameter(A_b_log)
            self.A_b_log._no_weight_decay = True 
            self.conv1d_b = nn.Conv1d(
                in_channels=self.d_inner,
                out_channels=self.d_inner,
                bias=conv_bias,
                kernel_size=d_conv,
                groups=self.d_inner,
                padding=d_conv - 1,
                **factory_kwargs,
            )

            self.x_proj_b = nn.Linear(
                self.d_inner, self.dt_rank + self.d_state * 2, bias=False, **factory_kwargs
            )
            self.dt_proj_b = nn.Linear(self.dt_rank, self.d_inner, bias=True, **factory_kwargs)

            self.D_b = nn.Parameter(torch.ones(self.d_inner, device=device))  # Keep in fp32
            self.D_b._no_weight_decay = True

前向

参数	维度
输入x	[b, l, d]
xz	[b, 2 * d, l]
x_dbl	[b,dt_rank + d_state * 2 ]

SSM参数	shape	来源
状态矩阵A	(d_in, n)	在初始化中定义，非数据依赖
输入矩阵B	(b, l, n)	由x_db1切分而来，因此数据依赖
输出矩阵C	(b, l, n)	由x_db1切分而来，因此数据依赖
直接传递矩阵D	(d_in)	在初始化中定义，非数据依赖
数据依赖步长$\Delta$	(b, l, d_in)	由x_db1切分而来，因此数据依赖

维度约定	说明
B / b	batch size
L / l	length
D / d	d_inner

输入映射

输入映射把输入x映射为两个分支xz，主分支x和门控分支z。

def forward(self, hidden_states, inference_params=None):
        """
        hidden_states: (B, L, D)
        Returns: same shape as hidden_states
        """
        batch, seqlen, dim = hidden_states.shape

        conv_state, ssm_state = None, None
        if inference_params is not None:
            conv_state, ssm_state = self._get_states_from_cache(inference_params, batch)
            if inference_params.seqlen_offset > 0:
                # The states are updated inplace
                out, _, _ = self.step(hidden_states, conv_state, ssm_state)
                return out

        # We do matmul and transpose BLH -> HBL at the same time
        xz = rearrange(
            self.in_proj.weight @ rearrange(hidden_states, "b l d -> d (b l)"),
            "d (b l) -> b d l",
            l=seqlen,
        )
        if self.in_proj.bias is not None:
            xz = xz + rearrange(self.in_proj.bias.to(dtype=xz.dtype), "d -> d 1")

fast_path

在之后通过use_fast_path分为两种

use_fast_path

在这里通过bimamba类别又可分为三类，v1，v2和其它

bimamba_type == v1

在v1版本中，调用的函数是bimamba_inner_fn 在后面专门介绍。

       if self.use_fast_path and inference_params is None:  # Doesn't support outputting the states
            if self.bimamba_type == "v1":
                A_b = -torch.exp(self.A_b_log.float())
                out = bimamba_inner_fn(
                    xz,
                    self.conv1d.weight,
                    self.conv1d.bias,
                    self.x_proj.weight,
                    self.dt_proj.weight,
                    self.out_proj.weight,
                    self.out_proj.bias,
                    A,
                    A_b,
                    None,  # input-dependent B
                    None,  # input-dependent C
                    self.D.float(),
                    delta_bias=self.dt_proj.bias.float(),
                    delta_softplus=True,
                )    
            

bimamba_type == v2

在v2版本中，调用的函数是mamba_inner_fn_no_out_proj在后面专门介绍。可以看到，在这里不同于v1，v2版本因为新增了一套SSM参数，因此也得到了额外的输出out_b，最后的输出也有两种模式，一是两者的简单平均，注意到因为反向SSM方向与正向方向相反，因此反向的输出要先翻转后再相加。而是直接翻转后相加。

			elif self.bimamba_type == "v2":
                A_b = -torch.exp(self.A_b_log.float())
                out = mamba_inner_fn_no_out_proj(
                    xz,
                    self.conv1d.weight,
                    self.conv1d.bias,
                    self.x_proj.weight,
                    self.dt_proj.weight,
                    A,
                    None,  # input-dependent B
                    None,  # input-dependent C
                    self.D.float(),
                    delta_bias=self.dt_proj.bias.float(),
                    delta_softplus=True,
                )
                out_b = mamba_inner_fn_no_out_proj(
                    xz.flip([-1]),
                    self.conv1d_b.weight,
                    self.conv1d_b.bias,
                    self.x_proj_b.weight,
                    self.dt_proj_b.weight,
                    A_b,
                    None,
                    None,
                    self.D_b.float(),
                    delta_bias=self.dt_proj_b.bias.float(),
                    delta_softplus=True,
                )
                if not self.if_devide_out:
                    out = F.linear(rearrange(out + out_b.flip([-1]), "b d l -> b l d"), self.out_proj.weight,self.out_proj.bias)
                else:
                    out = F.linear(rearrange(out + out_b.flip([-1]), "b d l -> b l d") / 2, self.out_proj.weight, self.out_proj.bias)

其他

如果选择了双向模式，却没有定义模式，则使用Mamba默认的mamba_inner_fn

		else:
                out = mamba_inner_fn(
                    xz,
                    self.conv1d.weight,
                    self.conv1d.bias,
                    self.x_proj.weight,
                    self.dt_proj.weight,
                    self.out_proj.weight,
                    self.out_proj.bias,
                    A,
                    None,  # input-dependent B
                    None,  # input-dependent C
                    self.D.float(),
                    delta_bias=self.dt_proj.bias.float(),
                    delta_softplus=True,
                )

no use_fast_path

和原始论文一致，如果不选择use_fast_path，则会在这里计算完整个流程，而不是定位到selective_scan_interface中定义的函数，而是计算出SSM参数后再调用selective_scan_interface中定义的selective_scan_fn()，SSM数据依赖的参数有参数映射x_proj得到x_db1，然后切分得到B， C，delta参数。

 else:
            x, z = xz.chunk(2, dim=1)

            if conv_state is not None:
                conv_state.copy_(F.pad(x, (self.d_conv - x.shape[-1], 0)))  # Update state (B D W)
            if causal_conv1d_fn is None:
                x = self.act(self.conv1d(x)[..., :seqlen])
            else:
                assert self.activation in ["silu", "swish"]
                x = causal_conv1d_fn(
                    x=x,
                    weight=rearrange(self.conv1d.weight, "d 1 w -> d w"),
                    bias=self.conv1d.bias,
                    activation=self.activation,
                )

            x_dbl = self.x_proj(rearrange(x, "b d l -> (b l) d"))  # (bl d)
            dt, B, C = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
            dt = self.dt_proj.weight @ dt.t()
            dt = rearrange(dt, "d (b l) -> b d l", l=seqlen)
            B = rearrange(B, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
            C = rearrange(C, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
            assert self.activation in ["silu", "swish"]
            y = selective_scan_fn(
                x,
                dt,
                A,
                B,
                C,
                self.D.float(),
                z=z,
                delta_bias=self.dt_proj.bias.float(),
                delta_softplus=True,
                return_last_state=ssm_state is not None,
            )
            if ssm_state is not None:
                y, last_state = y
                ssm_state.copy_(last_state)
            y = rearrange(y, "b d l -> b l d")
            out = self.out_proj(y)
        if self.init_layer_scale is not None:
                out = out * self.gamma    
        return out

双向SSM

v1

对于v1版本双向SSM在前向时首先定义到bimamba_inner_fn，然后调用BiMambaInnerFn

def bimamba_inner_fn(
    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
    out_proj_weight, out_proj_bias,
    A, A_b, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
    C_proj_bias=None, delta_softplus=True
):
    return BiMambaInnerFn.apply(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
                              out_proj_weight, out_proj_bias,
                              A, A_b, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)

前向

去掉和原始论文中的MambaInnerFn相同的部分，在forward前向过程中，不同在于定义了两个输出，分别为out_zf和out_zb，out_zf对应于原来的前向输出，out_zb则是新增的反向输出，最终的out_z是两者翻转相加。

        out_f, scan_intermediates_f, out_z_f = selective_scan_cuda.fwd(
            conv1d_out, delta, A, B, C, D, z, delta_bias, delta_softplus
        )
        assert not A_b.is_complex(), "A should not be complex!!"
        out_b, scan_intermediates_b, out_z_b = selective_scan_cuda.fwd(
            conv1d_out.flip([-1]), delta.flip([-1]), A_b, B.flip([-1]), C.flip([-1]), D, z.flip([-1]), delta_bias, delta_softplus,
        )

        out_z = out_z_f + out_z_b.flip([-1])

后向

去掉和原始论文中的MambaInnerFn相同的部分，在backward后向过程中。对应的，定义复制了一套新参数，参数对应如下

原参数	新增后向参数
dz	dz_b
dconv1d_out	dconv1d_out_f_b
ddelta	ddelta_f_b
dA	dA_b
dB	dB_f_b
dC	dC_f_b
dD	dD_b

        dz_b = torch.empty_like(dz)
        dconv1d_out_f_b, ddelta_f_b, dA_b, dB_f_b, dC_f_b, dD_b, ddelta_bias_b, dz_b, out_z_b = selective_scan_cuda.bwd(
            conv1d_out.flip([-1]), delta.flip([-1]), A_b, B.flip([-1]), C.flip([-1]), D, z.flip([-1]), delta_bias, dout_y.flip([-1]), scan_intermediates_b, out_b, dz_b,
            ctx.delta_softplus,
            True  # option to recompute out_z
        )
        

根据这些新定义的参数，我们和前向参数相加来重定义原始的参数。我们得到新的dconv1d_out，ddelta等参数，最终保持与原始SSM一致

		dconv1d_out = dconv1d_out + dconv1d_out_f_b.flip([-1])
        ddelta = ddelta + ddelta_f_b.flip([-1])
        dB = dB + dB_f_b.flip([-1])
        dC = dC + dC_f_b.flip([-1])
        dD = dD + dD_b
        ddelta_bias = ddelta_bias + ddelta_bias_b
        dz = dz + dz_b.flip([-1])
        out_z = out_z_f + out_z_b.flip([-1])

v2

对于v2版本双向Mamba在前向时首先定义到mamba_inner_fn_no_out_proj，然后调用MambaInnerFnNoOutProj。在v2版本，因为定义了两套SSM参数，因此双向的修改相比于v1要简单。

def mamba_inner_fn_no_out_proj(
    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
    C_proj_bias=None, delta_softplus=True
):
    return MambaInnerFnNoOutProj.apply(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
                              A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)

前向

mamba_inner_fn_no_out_proj 即相比于原始的mamba_inner_fn缺少了输出映射。

  return F.linear(rearrange(out_z, "b d l -> b l d"), out_proj_weight, out_proj_bias)

后向

相应的，在其中修改掉和out_proj_weight相关的部分。