Mamba 永远不会忘记任何标记
DeMansia: Mamba Never Forgets Any Tokens
https://arxiv.org/pdf/2408.01986
Introduction
1. 研究背景
-
Transformer架构的兴起：
作者首先强调了Transformer架构在深度学习领域的重大影响，尤其是在自然语言处理（NLP）和计算机视觉（CV）等领域的广泛应用。Transformer模型通过自注意力（self-attention）机制，能够动态地权衡输入数据中不同部分的重要性，从而实现对上下文的精准理解和处理。 -
Transformer的局限性：
尽管Transformer模型取得了显著的成功，但其自注意力机制的计算复杂度与输入序列长度的平方成正比，这在处理大规模输入数据或资源受限的环境下变得尤为不利。传统的Transformer在处理长序列时容易遇到计算瓶颈，从而限制了其应用范围。
2. 现有解决方案的不足
-
计算复杂度的挑战：
为解决上述问题，近年来研究者们提出了一系列优化方案，如使用低秩矩阵近似、线性化点积注意力计算或稀疏Transformer等，但这些方法往往需要在计算效率和模型精度之间做出权衡。 -
模型复杂度与准确性的平衡：
这些优化方法虽然能够在一定程度上降低计算复杂度，但往往会牺牲模型的准确性或复杂度。因此，在保持高性能的同时降低计算开销，仍然是Transformer架构面临的一大挑战。
3. DeMansia模型的提出
-
核心思想：
在此背景下，作者提出了DeMansia模型，该模型集成了Mamba和Vision Mamba（ViM）的优势，并借鉴了LV-ViT的训练流程，旨在解决图像分类任务中的计算挑战，特别是在资源受限的环境下实现高性能。 -
创新点：
DeMansia模型将状态空间模型（SSM）与标记技术相结合，通过位置感知的状态空间模型和创新的标记应用，在保持计算效率的同时，不牺牲模型的上下文理解能力。这一设计使得DeMansia能够在处理大规模图像数据时保持高效且准确的性能。
4. 研究目的与贡献
-
目的：
本文旨在通过详细介绍DeMansia模型的发展过程，并通过与现有模型的比较，展示其在图像分类任务中的有效性。 -
贡献：
作者期望DeMansia模型能够为计算机视觉领域的广泛应用提供一个新的、有前景的解决方案。通过优化计算效率和模型精度之间的平衡，DeMansia有望在资源受限的环境下实现更广泛的应用。
Method
文件首先概述了当前Transformer架构在处理长序列时存在的局限性，并介绍了Mamba和Vision Mamba（ViM）模型作为前提条件模型。随后，提出了DeMansia模型，该模型集成了Mamba和ViM的优势，并结合LV-ViT的训练管道，以增强图像分类任务的性能。
1. DeMansia模型架构
DeMansia模型的架构紧密遵循ViM架构，主要通过以下几个步骤处理输入图像：
- 图像预处理：通过四层卷积网络将输入图像转换为一系列patch嵌入。每个嵌入都添加了一个可学习的位置嵌入，以保留空间信息。
- 嵌入与位置编码：在嵌入序列中加入一个可学习的类别token和位置嵌入，以保留全局信息和空间位置。
- Vision Mamba编码器：使用多个Vision Mamba（ViM）块进行编码。每个ViM块包含线性投影层、SSM块、1D卷积等，以双向方式处理输入序列。
- 预测与分类：最终，通过多层感知机（MLP）头部进行预测，将类别token和patch token的平均（加权）作为全局类别预测。
2. Vision Mamba（ViM）块
ViM块是DeMansia模型的核心组件，它通过以下方式工作：
- 双向SSM机制：输入序列首先经过线性投影，然后通过1D卷积和SSM分别在正向和反向方向上进行并行处理。这两个方向的输出随后被合并。
- 选择性SSM：SSM块在ViM中负责捕获长期依赖关系，并通过选择性机制保持未压缩的上下文信息。
3. 训练过程
在训练阶段，DeMansia模型采用了LV-ViT中的token标注技术，具体过程如下：
- Token Labeling：计算token标注损失，通过标注每个patch token的类别来优化模型。这有助于模型更好地理解图像的局部和全局信息。
- 优化器与学习率调度：使用RAdam优化器，并采用Cosine Annealing with Warm Restarts学习率调度策略。
- 混合精度训练：利用自动混合精度（AMP）在Bfloat16格式下进行计算，以加速训练过程，同时保持模型权重在Float32中。
4. 实验设置
- 数据集：实验主要在ImageNet-1k数据集上进行，该数据集包含约1.28百万张训练图像和5万张验证图像，分布在1000个类别中。
- 数据增强：对每张图像进行随机裁剪和翻转增强，并归一化到224x224的大小。
- 模型变体：实验中主要使用了DeMansia Tiny变体，并计划在未来探索更大版本的模型。
- 训练参数：DeMansia Tiny在单个RTX A6000 GPU上训练了310个epochs，batch size为768。