Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

马昊宇作者

引入特征空间,显著降低计算量:双边局部注意力ViT性能媲美全局注意力

在这项研究中,来自百度研究院和香港大学的研究者重新思考了局部自注意力机制,提出了特征空间局部注意力(feature-space local attention或简称FSLA)。

Vision Transformer 舍弃了 ConvNet 先验信息,通过引入自注意力机制对远距离特征依赖进行建模,提升了模型的表征能力。然而 Vision Transformer 的自注意力机制在图像分辨率较高时,计算复杂度过高。为了克服这个问题,研究人员使用局部窗口计算自注意力,在此称之为图像空间局部注意力(image-space local attention 或简称 ISLA)。尽管基于窗口的图像空间局部注意力显著提升了效率,但仍面临难以捕捉远距离特征依赖的问题。


在这项研究中,来自百度研究院和香港大学的研究者重新思考了局部自注意力机制,提出了特征空间局部注意力(feature-space local attention 或简称 FSLA)。这种局部注意力从图像内容出发,把特征相似的 token 聚成类,并且只在每类特征的内部计算自注意力,相比全局自注意力显著降低了计算量,同时基本保留了原始的全局自注意力机制对远距离特征依赖的建模能力。

为了将特征空间局部注意力与图像空间局部注意力相结合,本文作者进一步提出了双边局部注意力 ViT (简称 BOAT),把特征空间局部注意力模块加入到现有的基于窗口的局部注意力视觉 Transformer 模型中,作为图像空间局部注意力的补充,大大提升了针对远距离特征依赖的建模能力,在几个基准数据集上的大量实验表明结合了特征空间局部注意力的模型明显优于现有的 ConvNet 和 ViT 模型。


  • 论文地址:https://arxiv.org/pdf/2201.13027v1

  • 代码地址:https://github.com/mahaoyuHKU/pytorch-boat


创新动机

为了保持更高的特征图分辨率,同时不会带来过高的运算复杂度,现有的图像空间局部注意力机制将一个图像划分为多个局部窗口,自注意力只在同一窗口的 token 间运算。这是一个合理的设计,因为一个 token 很可能与空间上邻近的 token 相关联。因此,局限于局部窗口的自注意力很可能不会显著降低性能,但是可以显著降低计算量。

本文重新思考了局部自注意力,从特征角度而非空间角度,实现了对局部窗口的划分。具体来说,图像空间局部自注意力的依据是:在空间上邻近的 token 很可能对彼此施加更大的影响(图 1 左);而本文提出的特征空间局部自注意力机制的依据是:即使在图像空间距离较远但在特征空间距离较近的 token 同样会对彼此有很大的影响,因此它在特征空间对 token 进行聚类,并且只在每类特征的内部如同空间局部窗口一样计算自注意力(图 1 右)。


本文提出的特征空间局部自注意力仅计算特征空间内最近邻的特征向量间的注意力,将距离较远的特征向量间的影响直接设为 0。这本质上定义了一个分段相似度函数,将相似度小的特征向量间的注意力近似为 0,降低了运算复杂度。与图像空间局部自注意力相比,特征空间局部自注意力在 ViT 模型中运用的较少。特征空间局部自注意力关注的是相似度较高的特征向量间的注意力,而不考虑空间上两者的邻近程度。因此,它是图像空间局部自注意力的很好补充,能对因跨越空间局部窗口而被遗漏的远距离特征依赖进行建模。

方法概述

本文中的 ViT 采用了和 Swin 和 CSWin 相同的层次化金字塔架构,由一个 patch embedding 模块和若干个双边局部注意力(bilateral local attention)模块组成。


本文与 Swin 和 CSWin 的主要区别是其中的 local attention 模块被替换成了下图所示的 bilateral local attention。而 patch embedding,position encoding 等设计皆和 Swin/CSWin 保持一致,所以接下来对 bilateral local attention 进行详细介绍。

Bilateral Local Attention

本文提出的 bilateral local attention 在基于窗口的图像空间局部注意力(ISLA)模型中添加了特征空间局部注意力(FSLA)模块。FSLA 模块根据 ISLA 模块的输出计算在特征空间彼此邻近的 token 之间的注意力:


最后,将 FSLA 模块的输出送入另一个归一化层和一个 MLP 模块进行处理,再通过一个短路连接得到整个 bilateral local attention 模块的输出:


FSLA 的重点是如何对特征进行聚类操作,并且在各个类内部计算自注意力。最直觉的方法是使用 K-means 聚类,但 K-means 聚类不能确保分组结果大小相同,这使得在 GPU 平台上难以有效地实现并行加速,同时也可能对自注意力计算的有效性产生负面影响。

因此本文提出均衡层次聚类,它进行 k 层聚类。在每一层,它进行均衡二分聚类,将上一层的各个类组均衡地划分为两个更小的类组。如下图所示,所有 token 分成了 token 数量相同的 8 个类组,然后在每组内部计算自注意力,具体的自注意力参数和图像空间局部注意力保持一致。


假如某个类组原先有 2m 个 token,均衡二分聚类后得到的每组的 token 数量为 m。与 K-means 类似,均衡二分聚类是一个迭代算法并且依赖于聚类中心。如以下算法所示,在每次迭代对所有 token 进行分组时,先计算每个 token 到两个聚类中心的距离比值,然后把所有 token 按距离比值的递减顺序排序,最后将排序列表前半部分 m 个 token 赋给第一组,后半部分 m 个 token 赋给第二组。


需要注意的是,这样进行无重叠的均衡二分聚类可能会导致两个处于排序列表中段位置的、特征比较相似的 token 被分配到两个不同的类组中,从而无法计算它们之间的相互影响。因此在实际计算中,为了避免遗漏邻近特征间的影响,会保留一定程度的类间重叠,也就是把排序列表的最前面 m+n 个 token 赋给第一组,最后 m+n 个 token 赋给第二组。这样两组之间就存在 2n 个重用的 token,这样的类间重叠会导致额外的运算,因此实际只在层次聚类的最后一层进行有重叠的均衡二分聚类。完成聚类以后,在每组 token 内部按照常规操作进行自注意力机制的计算即可。

值得注意的是,本文中的所有聚类都是临时计算的,不包含任何可学习的参数,因此不存在对聚类算法本身进行梯度回传的问题。此外,所有聚类运算都用 GPU 进行了加速,对模型的整体计算量影响不大。

实验结果

BOAT 遵循与其它 ViT 相同的训练策略。本文使用 ImageNet-1K 的训练集训练模型,输入图像使用 224×224 分辨率,并且没有外部数据。

具体来说,训练 300 个 epochs,使用 AdamW 优化器、余弦学习速率调度器和一个线性预热过程。BOAT 在多个数据集上都取得了 SOTA 的效果。比如,在 ImageNet-1K 测试集上,BOAT-CSWin-T 取得了 83.7 的 Top-1 分类准确率;在 ADE20K 语义分割测试集上,BOAT-CSWin-T 的 mIoU 达到了 50.5。


理论ViT注意力机制
相关数据
调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源,如线程、进程或数据流;也可以指硬件资源,如处理器、网络连接或扩展卡。 进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态,允许多位用户有效地同时共享系统资源,或达到指定的服务质量。 see planning for more details

自注意力技术

自注意力(Self-attention),有时也称为内部注意力,它是一种涉及单序列不同位置的注意力机制,并能计算序列的表征。自注意力在多种任务中都有非常成功的应用,例如阅读理解、摘要概括、文字蕴含和语句表征等。自注意力这种在序列内部执行 Attention 的方法可以视为搜索序列内部的隐藏关系,这种内部关系对于翻译以及序列任务的性能非常重要。

基准技术

一种简单的模型或启发法,用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集(或特征)的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的,这让「神经网络在执行预测任务时可以更多关注输入中的相关部分,更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时,源句子中仅有少部分是相关的;因此,可以应用一个基于内容的注意力机制来根据源句子动态地生成一个(加权的)语境向量(context vector), 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

准确率技术

分类模型的正确预测所占的比例。在多类别分类中,准确率的定义为:正确的预测数/样本总数。 在二元分类中,准确率的定义为:(真正例数+真负例数)/样本总数

层次聚类技术

层次聚类通过对数据集在不同层次进行划分,从而形成树形的聚类结构。数据集的划分可采用“自底向上”的聚合(agglomerative)策略,也可采用“自顶向下”的分拆(divisive)策略。“自底而上”的算法开始时把每一个原始数据看作一个单一的聚类簇,然后不断聚合小的聚类簇成为大的聚类。“自顶向下”的算法开始把所有数据看作一个聚类,通过不断分割大的聚类直到每一个单一的数据都被划分。

语义分割技术

语义分割,简单来说就是给定一张图片,对图片中的每一个像素点进行分类。图像语义分割是AI领域中一个重要的分支,是机器视觉技术中关于图像理解的重要一环。

局部注意力机制技术

每生成单个目标词都需要考虑原语句子的所有词在计算上是非常昂贵的,也不是必需的。为了减少这种问题,Luong et al. [20] 提出了一种仅关注固定窗口大小 2D+1 的局部注意力机制

优化器技术

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

百度网讯机构

北京百度网讯科技有限公司是百度旗下公司,拥有网页搜索、hao123、百度推广等多条业务线,致力于为用户提供“简单可依赖”的互联网搜索产品及服务。

http://www.baidu.com
聚类技术

将物理或抽象对象的集合分成由类似的对象组成的多个类的过程被称为聚类。由聚类所生成的簇是一组数据对象的集合,这些对象与同一个簇中的对象彼此相似,与其他簇中的对象相异。“物以类聚,人以群分”,在自然科学和社会科学中,存在着大量的分类问题。聚类分析又称群分析,它是研究(样品或指标)分类问题的一种统计分析方法。聚类分析起源于分类学,但是聚类不等于分类。聚类与分类的不同在于,聚类所要求划分的类是未知的。聚类分析内容非常丰富,有系统聚类法、有序样品聚类法、动态聚类法、模糊聚类法、图论聚类法、聚类预报法等。

推荐文章
暂无评论
暂无评论~