MAP

发表于 更新于 分类于 阅读次数： 本文字数： 2.3k 阅读时长 ≈ 8 分钟

MAP: Unleashing Hybrid Mamba-Transformer Vision Backbone's Potential with Masked Autoregressive Pretraining[1]

作者是来自清华叉院和上海AI Lab、QiZhi 研究院的Yunze Liu和Li Yi，论文引用[1]:Liu, Yunze and Li Yi. “MAP: Unleashing Hybrid Mamba-Transformer Vision Backbone's Potential with Masked Autoregressive Pretraining.” ArXiv abs/2410.00871 (2024): n. pag.

Time

  • 2025.Mar

Key Words

  • masked Autoregressive Pretraining
  • 一句话总结：结合了MAE of Transformer的local features 和AR of Mamba的long context modeling

总结

  1. 混合的Mamba-Transformer网络最近受到了很多的关注，这些网络利用Transformer的可扩展性和Mamba的long-context modeling和高效计算。然而，有效地预训练这样的混合网络仍然是一个open question，现有的方法，例如MAE 或者自回归 pretraining，主要聚焦于single-type network 架构，相比之下，对于Mamba和Transformer的混合结构，预训练策略必须有效，基于此，作者提出了Masked Autoregressive pretraining，以统一的范式，提高了Mamba和Transformer modules的性能。

  1. Mamba-Transformer backbone 最近吸引了广泛的关注，利用Transformer的scalability和Mamba在long-context language modeling下的strong capabilities的又是，然而，如何scale up Mamba-Transformer vision backbones，一个有效的pretraining策略对于最大化Mamba和Transformer的结合的能力是重要的。作者的工作是这个方向上的first step. 尽管MAE在预训练的Transformer中被验证为有效的，对于Mamba的预训练方法仍然是没有被探索的，最优的approach是unclear的，另外，混合架构要求一个pretraining策略能够兼容两者的computation blocks。这个很有挑战，因为SSM capture visual features的方式和Transformers不同。

    为了解决这个挑战，作者进行了实验性的研究，发现了三个key observations。**首先，现有的Transformer的pretraining方法，例如MAE和Contrastive Learning，在Mamba上没有产生满意的结果；另外，自回归pretraining对于Mamba-based vision backbones是有效的，用合适的scanning pattern和token masking ratio, 第三，对于Mamba或者Transformer有效的预训练策略，对于另外一个没有效果。

    基于以上的观察，作者开发了一个新的pretraining策略，对于Mamba-Transformer vision backbone有效果，称之为Masked Autoregressive Pretraining，或者简称为MAP。这个key是hierarchical pretraining objective，局部的MAE被利用来学习good local attention for Transformer blocks，然而，global autogressive pretraining使得Mamba blocks能够学习有意义的contextual information。特别地，预训练方法是有两个key designs，首先，利用local MAE使得hybrid framework，特别是Transformer module，学习Local bidirectional connectivity，这要求hybrid network在感知local bidirectional information之后，在一个local region中预测所有的tokens，其次，对于每个local region，自回归地产生tokens, 使得hybrid network，特别是Mamba Module能够学习rich contextual information。这要求network能够自回归地产生subsequent local regions，基于之前的decoded tokens。

    注意到，因为这篇paper的focus在于hybrid network的pretraining策略，而不是hybrid network本身，作者选择了一个densely mixed hybrid scheme，作为default hybrid framework，称之为HybridNet，作者在HybridNet上的实验表明，MAP预训练的方法，在2D、3D vision tasks上显著地超过了其它的pretraining策略。

    贡献如下:

    • 提出了预训练Mamba-Transformer的预训练方法
    • 对于diagnostic purposes，作者用自回归预训练，对pure-Mamba的Key components进行了深度的分析，表明effective取决于pretraining order和Mamba scanning order之间的consistency，第三，作者展示了，它们的方法提高了很多。

  2. Vision Mamba是visual representation learning的有效的model，利用双向state space blocks，超过了传统的vision transformers，例如DeiT。VMamba架构，用了Visual State-Space blocks和2D Selective Scanning，通过平衡效率和精度，在visual perpcetion tasks上表现突出。自回归进一步提高了vision Mamba的性能，使其实现了superior accuracy，相比于传统的监督的model，加速了training。然而，为什么自回归对于Vision Mamba是有效的，key factors是什么还是悬而未决的问题。在本文中，作者探索了Mamba 自回归预训练背后的success的critical design elements，提出了MAP和ARM之间的key difference在于，在AR 过程中， 用了random masking，network基于masked local information，解码next complete local information，这个masking mechanism增强了local information的modeling capability和relevance of context。ViT采用了transformer架构，将image patches视为sequential tokens，Swin Tranformer用shifted window引入了hierarchical design，有效地capturing both local和global contextual information。MAE对于Mamba没有效果。尽管用了MAP的混合架构很难超越相同条件下的用MAE的Transformer，它提供了计算成本和性能之间的平衡，还有保持scalability的同时，引入了longer contexts，这个property在video domain和large language models中特别重要。

  3. 自监督视觉表征学习: 自监督视觉表征学习是一个机器学习的方式，从大量的无标注数据中提取meaningful visual features，这个方法依赖于pretext tasks，通常作为一个不需要explicit labels的方式，来学习representations。GTP-style AR models，通过给定的previous parts，来预测an image or sequence的next part，使得model能够理解spatial or temporal dependecies within data。MAE是随机mask random patches，训练model来重建masked regions，使得model能够学习contextual information和global representations。Contrastive Learning涉及contrasting positive和negative samples，来学习discriminative features，通常涉及creating pairs of positive and negative examples，训练model来区分。然而，作者发现，现有的pretraining策略，不能unlock hybrid framework的潜力，激励作者来探索一个新的pretraining paradigm。由于MAP的有效地model local features和associations between local regions的能力，在pure transformer架构上，尽管MAP没有超过MAE的性能，实现了comparable results，这是因为local MAE降低了receptive field，local regions之间的autoregressive modeling增强了capture local relationships的能力。

  4. 通过实验表明：MAE更适合Transformer pretraining，AR更适合Mamba pretraining。对于ViT，用MAE是重要的，它能够establish tokens之间的bidirectional associations，提高了性能，相比之下，对于Vim，model preceding 和suceeding tokens之间的continuity更重要，因此，对于hybrid network，需要一个策略，能够利用MAE在bidirectional modeling上的优势，同时保持AR在context modeling上的长处，对于AR，pretraining design(scanning order) 和Mamba scanning order一致，效果最好，同时Masking ratio也是重要的。

    MAP利用AR来model local regions和local MAE来 model local regions之间的internal features，同样采用masking 策略，这对于MAE和AR同样重要，来增强networks的representations，这个策略结合了MAE在local feature modeling的长处，还有AR in context modeling的长处，继承了both，MAP提高了hybrid architecture的pretraining性能。

  5. 本文聚焦于如何预训练Mamba-Transformer frameworks，而不是设计hybrid frameworks。作者训练了很多Hybrid Mamba-Transformer Vision backbones，和从头开始训练进行对比。结果显示，用MMMTMMMT表型最好。和Mamba-R* with MMMMMMTT相比，作者发现，在Mamba后面增加一个Transformer，增强了它的long-context modeling capabilities，提高了性能。然而，将MMMMMMTT和TTMMMMMM相比，作者观察到，在Mamba后面简单地添加Transformers，不能完全利用网络架构的潜力，这表明：在一开始引入Transformers，对于提取sufficient local features是重要的。作者相信，MMMTMMMT方法，有效地平衡了local feature extraction和contextual modeling enhancement。 作者的MAP pretraining是在一个给定的image上进行random masking，然后region-wise autoregressively重建它的original image。image的每个row的tokens被同时预测，rows之间的tokens通过自回归预测。

    注意到，将每个row作为sub-region，是因为，Mamba的AR order应该和scanning order一致，大多数Mamba的scannning order是row-first。为了确保和现有方法的公平对比，将每个row作为prediction unit，在改变Mamba的scannning order之后，大多数复杂的clustering 策略期望能产生更好的结果。

    • 如图所示，HybridNet首先将randomly masked image 映射到feature space，然后用一个Transformer Decoder来decoder original image row-wise。总的loss function是negative log-likelihoods of the predicted tokens。

    \[\mathcal{L} = - \sum_{i=1}^{M} \sum_{j \in \mathbf{M}_i} \log p(\mathbf{x}_{ij} \mid \mathbf{x}_{i,j \notin \mathbf{M}_i}, \mathbf{r} < i)\]

    • 对于pure Transformer，MAP没有超过MAE的性能，但是实现了comparable results，这是因为local MAE将i的了receptive field，autoregressive modeling between local regions增强了capture local relationships的能力，使得MAP保持了strong performance。

    • 对于Masking，实验发现，random masking的效果比sequential、diagonal masking的效果都好。用一个masked transformer decoder来重建，因为Transformer decoder能够基于encoder's features重建region-wise，而Mamba decoder由于其单向的scanning nature，很难同时重建整个local region，作者的decoder采用了row-wise decoding 策略，使得同时autoregressive decoder of one row of tokens，增强了网络capture local features和contextual relationships的能力。local MAR是指将MAE的receptive field限制在image的每个row。

    • 重建target：和MAE一致，作者将重建normalized original pixels作为target，采用MSE loss，受MAR的启发，用重建的输出作为conditional signal，用于diffusion models，来提高generation quality。探索了用diffusion loss能够增强性能，然而，这个方法没有产生significant improvements，表明重建的images的quality和encoder pretraining success不直接相关。

MAP \(Fig.1^{[1]}\) MAP

\(Fig.2^{[1]}\) MAP首先randomly mask input image，然后以row-by-row autoregressively 的方式，重建original image。预训练方法展示了significant advantages in modeling contextual features of local characteristics and correlations between local features.

Framework \(Fig.1^{[1]}\)

comparison \(Fig.2^{[1]}\)