SiMBA：基于FFT变换和Mamba的图像分类模型算法解析

1 算法原理
- 1 EinFFT：
- - 1.1 频谱变换
  - 1.2 频域通道混频 & 频谱门控网络
  - 1.3 逆频谱变换
- 2 整体计算公式
- 3 SiMBA：
- 4 小结
2 代码大致逻辑
- EinFFT模块
- MambaBlock模块
- MambaLayer模块
- PVT2FFN和FFN模块
- ClassBlock和Block_mamba模块
- DownSamples模块
- Stem模块
- 定义SiMBA模型
- 3 一些环境配置问题

《SiMBA: Simplified Mamba-based Architecture for Vision and Multivariate Time series》

github： https://github.com/badripatro/Simba

这是一种引入EinFFT进行通道建模的新架构。SiMBA体系结构将Mamba用于序列建模，并引入EinFFT作为一种新的通道建模技术。有效地解决了在扩展到大型网络时在Mamba中观察到的不稳定性问题。该方法突出了基于卷积模型、transformers模型、mlp混频器、频谱混频器模型和状态空间方法的各种模型。

1 算法原理

1 EinFFT：

提出了一种新的信道建模技术，称为 EinFFT，它解决了 Mamba 中的稳定性问题。该方法使用具有非线性的傅里叶变换将特征值建模为负实数，从而解决了不稳定性。

SiMBA 的通道混频组件（EinFFT）包含三个主要组件：频谱变换、使用 Einstein 矩阵乘法的频谱门控网络、逆谱变换

1.1 频谱变换

$\mathcal{F}(k)=\int_Df(x)e^{-j2\pi kx}\mathrm{d}x=\int_Df(x)\cos(2\pi kx)\mathrm{d}x+j\int_Df(x)\sin(2\pi kx)\mathrm{d}x$
傅里叶变换用于将输入信号分解为其组成频率。
空间域中两个函数的卷积的傅里叶变换等于它们的傅里叶变换的逐点乘积。

1.2 频域通道混频 & 频谱门控网络

对于给定的输入 $\mathbf{X}\in\mathbb{R}^{N\times D}$ 及其相应的频域转换 $\mathcal{X}\in\mathbb{R}^{N\times D}$ ，使用傅里叶变换。卷积定理指出， $\mathbf{X}$ 上卷积的傅里叶变换及其内核位于空间域中，可以等效地表示为其频域表示的乘积运算。这种等价表示如下：
$\mathcal{F}(\mathbf{X}*W+B)=\mathcal{F}(\mathbf{X})\cdot\mathcal{F}(W)+\mathcal{F}(B)=\mathcal{X}\mathcal{W}+\mathcal{B}$
其中 ∗ 表示圆形卷积， $\mathcal{W}$ 和 $\mathcal{B}$ 表示频域中的复数权重和偏置，W和B表示空间域中的权重和偏置， $\mathcal{F}$ 表示离散傅里叶变换（DFT）。输出 DFT 运算是复数值 $\mathcal{X}\in\mathbb{R}^{N\times D}$ ，乘以复数权重矩阵 $\mathcal{W}\in\mathbb{R}^{D\times D}$ ，并使用频谱门控网络（SGN）加上复数偏置 $\mathcal{B}\in\mathbb{R}^D$ 。为了降低复杂性，我们执行 Einstein 矩阵乘法（EMM）以获得高效的块对角矩阵。SGN 由以下公式表示：
$h^\ell=\sigma(h^{\ell-1}\mathcal{W}^\ell+\mathcal{B}^\ell),h^0=\mathcal{X}$
拓展后的公式为：
$\begin{aligned} Re(h)^{\ell}& =EMM(Re(h^{\ell-1})\mathcal W_{r}^{\ell})-EMM(Im(h^{\ell-1})\mathcal W_{i}^{\ell})+\mathcal B_{r}^{\ell} \\ Im(h)^{\ell}& =EMM(Re(h^{\ell-1})W_i^\ell)+EMM(Im(h^{\ell-1})W_r^\ell)+B_i^\ell \\ h^{\ell}& =\sigma(Re(h)^\ell)+j\sigma(Im(h)^\ell) \end{aligned}$
EMM（爱因斯坦矩阵乘法，又叫爱因斯坦求和约定）（结果和普通矩阵乘法一样，只是表达形式不同）：
$\mathbf{Y}^{N\times C_b\times C_d}=\mathbf{I}^{N\times C_b\times C_d}\boxed{*}\mathbf{W}^{C_b\times C_d\times C_d}$
这里 $\boxed{*}$ 表示 Einstein 矩阵乘法频域中的信道混频，表示为 EinFFT，是通过分别计算频率分量的实部和虚部来执行的。随后，这些部分被堆叠成一个复数，最终结果是通过将实部和虚部连接起来得到的。这种方法增强了模式的提取，同时确保了计算资源的有效利用。

1.3 逆频谱变换

$f(x)=\int_{D}\mathcal{F}(k)e^{j2\pi kx}\mathrm{d}f=\int_{D}(Re(\mathcal{F}(k))+jIm(\mathcal{F}(k))e^{j2\pi kx}\mathrm{d}f$

2 整体计算公式

$\begin{aligned} \mathcal{X}^{:,(l)}=\mathcal{X}_{R}^{:,(l)}+j\mathcal{X}_{I}^{:,(l)}& =\mathrm{FFT}(\mathbf{X}_{c}^{:,(l)}) \\ Re(h)^\ell+jIm(h)^\ell & =\mathrm{EMM}(\mathcal{X}^{\therefore,(l)},\mathcal{W}^c,\mathcal{B}^c)\text{ as in eq-4} \\ \mathcal{Y}^{:,(l)}=\mathcal{Y}_{R}^{:,(l)}+j\mathcal{Y}_{I}^{:,(l)}& =\sigma(Re(h)^\ell)+j\sigma(Im(h)^\ell) \\ \mathcal{Z}^{:,(l)}=\mathcal{Z}_{R}^{:,(l)}+j\mathcal{Z}_{I}^{:,(l)}& =\mathrm{EMM}(\mathcal{Y}^{:,(l)},\mathcal{W}^c,\mathcal{B}^c)\text{ as in eq-4} \\ \mathbf{Z}_{c}^{:,(l)}& =\mathrm{IFFT}(\mathcal{Z}_{R}^{:,(l)}+j\mathcal{Z}_{I}^{:,(l)}) \end{aligned}$
在这里插入图片描述

3 SiMBA：

虽然我们将 Mamba 块应用于视觉任务，但与 S4 或 Hippo 等其他模型相比，我们面临着稳定性问题（损失收敛问题）。我们为不稳定性问题提供了一种类型的解决方案，即仅保留负特征值。为了实现这一点，我们需要一个额外的模块，我们称之为 channel mixing，这在 mamba 块中是缺失的。我们将通道混合模块与 Sequence 混合模块相结合，制作了一个简化的基于 Mamba 的架构（SiMBA）。
SiMBA 使用 Mamba 块进行序列建模，L 表示层数，Norm 代表归一化层，EinFFT 用于通道建模。该过程包括将 Mamba 块应用于前一个时间步 $X_{l−1}$ ，然后是残差连接和dropout（DP）。然后对生成的张量进行归一化（Norm）以获得最终的序列向量。随后的操作涉及应用 EinFFT，这是论文中提出的频域通道混频操作。最后，张量经历另一次 dropout 并被添加到之前的状态 $X_{l}$ 。

在这里插入图片描述

4 小结

SiMBA 说明了性能和可扩展性之间的重要权衡。Mamba 本身对于大型网络可能存在稳定性问题。Mamba 与 MLP 相结合进行通道混合，弥合了小型网络的性能差距，但对于大型网络来说，可能存在相同的稳定性问题。Mamba 与 EinFFT 相结合，解决了小型和大型网络的稳定性问题。对于大型网络，最先进的变压器可能仍然存在性能差距，我们将在未来的工作中解决这一问题。