阿里开源新一代通用神经网络架构: UniNeXt

Title: UniNeXt: Exploring A Unified Architecture for Vision Recognition
Paper: https://arxiv.org/pdf/2304.13700.pdf
Code: https://github.com/jianlong-yuan/UniNeXt/tree/main

导读

本文集中探讨了计算机视觉任务中的通用架构问题。最近的一些研究集中于通过精心设计的空间令牌混合器(Spatial Token Mixer, STM)来提高性能。但是，作者认为一个设计良好的通用架构可以显著提高整个骨干网络的性能，而不论配备哪种空间标记混合器。因此，本文提出了UniNeXt，一种改进的通用架构。

为了验证其有效性，本文采用了各种经典和现代化设计实例化了 STM。实现结果表明，与它们最初提出的架构相比，所提架构可以稳定地提高所有 STM 的性能，并缩小它们之间的性能差距。令人惊讶的是，当配备最原始的局部窗口注意力的UniNeXt甚至优于之前的最新技术，这表明优秀的 STM 可能会因为通用架构的次优性而受到压制，这进一步显示了对视觉骨干网络通用架构的研究的重要性。

可以明显地看出，本文核心在于 STM 的通用架构设计。那么，如何理解 STM 呢？

在计算机视觉任务中，输入的图像通常是一组特征映射Feature Maps，每个特征映射都由一组特征向量Feature Vectors组成。在 Vision Transformers 中，特征向量被视为一个 Token，并且令牌或称为标记之间存在依赖关系，这些依赖关系可以通过 STM 来处理。因此，STM 旨在学习在特征映射中所有 Token 之间的空间关系，并将它们组合成更有意义的表示。这种操作可以在卷积或注意力的形式下进行。

方法

Overall Architecture

在本文中，UniNeXt的主要设计思想是通过三种方法增加归纳偏差：

1. 添加并行的 EC 分支到 STM；
2. 在通道 MLP 后添加 OC 模块；
3. 在 FFN 中添加 3×3 深度可分离卷积。

归纳偏差inductive bias早先是指在机器学习模型中，因为先验知识或者假设导致模型对某些假设具有更高的概率，从而更容易学习到某些模式或规律。这里大家可以这样简单理解下：CNN 具有天然的局部性，即强的归纳偏置，因此通常我们仅需少量数据就能训练出一个不错的模型出来，但上限不高。另一方面，Transformer 因为没有了归纳偏差的“束缚”，所以一般需要比较 strong 的数据增强和较大的数据量才能训出一个不错的模型，但它的潜力很大。

UniNeXt 的整体架构图如下所示：

如上所示，UniNeXt 整体依旧沿用了常见的金字塔结构，包含四个分层阶段，每个阶段都包含一个下采样层和多个Unified Blocks。在下采样层中，每个阶段的空间下采样比例为 2，通道数扩展两倍。统一块中的 TOken 数量保持一致。最后，本文方法应用全局平均池化(Global Average Pooling, GAP)和全连接层来执行图像分类任务。

值得注意的是，UniNeXt 的设计灵感来源于Swin Transformer、ConvNeXt和PVT等经典的网络架构，并且通过添加自定义的 EC 分支、OC 模块和 3x3 dwconv 等模块，进一步增强了模型的表示能力。下面我们就分别就这几个模块展开并详细介绍一下。

High-dimensional Convolution (HdC)

首先，我们先看下HdC，这是一种轻量的 深度可分离卷积，它可以将高维特征进行局部融合，从而编码高维隐式特征，提高模型的性能。作者继承和扩展了这种卷积嵌入机制，首先使用 MLP 的第一个线性层将特征维度映射到高维特征 F，然后使用深度可分离卷积来进行局部融合，从而编码高维隐式特征，提高模型的效率。通过这种方法，UniNeXt 可以更好地建模高维特征，并且在保持模型轻量化的同时，获得更好的性能。

Embedded Convolution (EC)

其次，本文进一步提出了Embedded Convolution，主要有两个方面的动机：

1. EC 增强了模型的归纳偏差，这对于学习和泛化至关重要。
2. EC 与所有现有的 Token Mixer 兼容，提供了跨各种架构的灵活性和易于实现性。

因此，EC 能够增强模型对于空间信息的建模，从而提高模型的性能。此外，由于 EC 是一种通用的卷积操作，能够兼容各种 STM，因此可以在各种视觉任务中进行灵活使用。

Post Convolution (PC)

最后，为了进一步增强卷积嵌入偏差和增强局部表示，本文精心设计了一种卷积结构——Post Convolution。其中，PC 能够通过引入轻量级的3×3深度卷积，在特别是密集预测任务中显著提高模型性能。具体而言，首先将 Token 矩阵转换为 2D 图像表示，然后进行深度卷积以进行局部上下文融合，紧接着进行 flatten 操作，最后使用残差连接以防止权重过度缩放。通过这种方法能够有效增强模型对于图像局部信息的建模，进而提高模型的性能。

Unified Mixer

总的来说，Unified Mixer是一种模型架构，通过在架构中实例化不同的 Token Mixer，例如无参数、卷积和注意力等，模型性能稳步提高。此外，使用仅纯局部窗口注意力的框架，在性能上超越了先前的 sota 结果。

在上图中，作者首先对特征图进行非重叠窗口分区，对于图像分类使用窗口大小为 7，对于下游任务（例如检测和分割）使用窗口大小为 11。需要注意的是，与 ViT 和 Swin 不同，本文方法不使用额外的绝对位置编码或相对位置编码，因为所提出的卷积嵌入技术足以表示标记之间的位置关系。

实验

总结

这篇论文提出了一种名为UniNeXt的统一架构，通过将空间令牌混合器与其他组件相结合，增加了网络的归纳偏差，包括在 STM 中添加并行的 EC 分支、在通道 MLP 后面添加 OC 模块以及在 FFN 中添加 3×3 的深度可分离卷积。最后，在广泛的实验中，作者验证了该框架的有效性，并呼吁研究人员关注宏观架构设计，而非仅仅是一些精心设计的模块（魔改？）。

作者：派派星
文章来源：CVHub

推荐阅读

• 在BEV、3D检测、2D检测、Lane检测的应用，量化与加速
• 全新YOLO模型YOLOCS | 改进YOLOv5的Backbone/Neck/Head
• 特斯拉FSD技术优化改进分析
• 小目标检测技巧 | 全局上下文自适应稀疏卷积CEASA | 助力微小目标检测涨点
• 车道线模型落地技巧 | LGAD注意力蒸馏让模型更鲁棒
• 基于YOLOv5改进再设计 | M2S全面提升小目标精度

更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。