深度学习基础模型之Mamba

news2024/11/24 10:34:44

Mamba模型简介

问题：许多亚二次时间架构(运行时间复杂度低于O(n^2)，但高于O(n)的情况)（例如线性注意力、门控卷积和循环模型以及结构化状态空间模型（SSM））已被开发出来，以解决 Transformer 在长序列上的计算效率低下问题，但此类模型的一个关键弱点是它们无法执行基于内容的推理

1. 模型架构

模型简单理解(特殊的门控RNN网络)：线性层+门控+选择性SSM的组合

在这里插入图片描述

2. 模型特点

2.1 选择性机制

在这里插入图片描述

$\Delta$ 、A、B、C应该是SSM中的可学习参数

根据输入参数化 SSM 参数来设计一种简单的选择机制，这使得模型能够过滤掉不相关的信息并无限期地记住相关信息。
这里作者认为（研究动机）：‘序列建模的一个基本问题是将上下文压缩成更小的状态。事实上，我们可以从这个角度来看待流行序列模型的权衡。例如，注意力既有效又低效，因为它明确地根本不压缩上下文。自回归推理需要显式存储整个上下文（即KV缓存），这直接导致Transformers的线性时间推理和二次时间训练缓慢。’
序列模型的效率与有效性权衡的特征在于它们压缩状态的程度：高效模型必须具有较小的状态，而有效模型必须具有包含上下文中所有必要信息的状态。反过来，我们提出构建序列模型的基本原则是选择性：或关注或过滤掉序列状态输入的上下文感知能力。

2.2 硬件算法

算法通过扫描而不是卷积来循环计算模型，但不会具体化扩展状态，计算速度比所有先前的 SSM 模型提升三倍。

代码调用

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
print(x.shape)
print(y.shape)
assert y.shape == x.shape

总结

这项基础性模型研究旨在解决transformer模型的长序列数据计算效率低的问题，其解决方法的动机：利用选择性机制实现有效特征的提取。个人理解为通过有效特征信息的选择实现知识提取（信息压缩），这让我联想到，最初的VGG语义分割网络结构设计其实类似于模拟知识特征的压缩与抽取，但后来发现这种方式会损失边缘信息，因此提出了U-net架构，再进一步卷积的方式无法有效估计全局上下文信息的联系，进而提出注意力机制来解决这一问题。
从技术与文章写作的角度来看，问题的发展似乎从知识压缩->细节特征提取->全局信息整合，到Mamba貌似是在全局信息整合基础上在进行一次有效信息的抽取，进而使模型从数据中提取根据代表性的特征。整体突出一点：深度学习也是一个特征工程，利用模型来替换原有的手工设计的特征