抛弃注意力,类Transformer新模型实现新SOTA


选自OpenReview
机器之心编译
机器之心编辑部
Transformer 用在图像识别上会存在计算效率过低的挑战 , 最近一篇向 ICLR 2021 大会提交的论文似乎很好地解决了这一问题 , 其研究的 PyTorch 代码也已在 GitHub 上开源 。
在自然语言处理领域(NLP)成为重要力量的 Transformer 技术最近已经开始在计算机视觉领域展现自己的实力 。 不过到目前为止 , 由于 Transformer 注意力机制对内存的需求是输入图像的二次方 , 所以这一方向还存在一些挑战 。
【抛弃注意力,类Transformer新模型实现新SOTA】近日 , LambdaNetworks 的出现提供了一种解决此问题的方法 , 人们可以无需建立昂贵的注意力图即可捕捉长距离交互 。 这一方法在 ImageNet 上达到了新的业界最佳水平(state-of-the-art) 。
抛弃注意力,类Transformer新模型实现新SOTA
本文插图
论文链接:https://openreview.net/pdf?id=xTJEN-ggl1b
GitHub链接:https://github.com/lucidrains/lambda-networks
对长程交互进行建模在机器学习中至关重要 。 注意力已成为捕获长程交互的一种常用范式 。 但是 , 自注意力二次方式的内存占用已经阻碍了其对长序列或多维输入(例如包含数万个像素的图像)的适用性 。 例如 , 将单个多头注意力层应用于一批 256 个64x64 (8 头)输入图像需要32GB的内存 , 这在实践中是不允许的 。
该研究提出了一种名为「lambda」的层 , 这些层提供了一种捕获输入和一组结构化上下文元素之间长程交互的通用框架 。
lambda 层将可用上下文转换为单个线性函数(lambdas) 。 这些函数直接单独应用于每个输入 。 研究者认为 , lambda 层可以作为注意力机制的自然替代 。 注意力定义了输入元素和上下文元素之间的相似性核 , 而 lambda 层将上下文信息汇总为固定大小的线性函数 , 从而避免了对内存消耗大的注意力图的需求 。 这种对比如图1所示 。
抛弃注意力,类Transformer新模型实现新SOTA
本文插图
研究者证明了 lambda 层的通用性 , 展示了它们的实现可以被用来捕获全局、局部或掩模上下文中基于内容和位置的交互 。 利用lambda生成的神经网络 LambdaNetwork 计算效率很高 , 能够以很小的内存成本建模长程依赖 , 因而可用于高分辨率图像等大型结构化输入 。
研究者在计算机视觉任务上评估了 LambdaNetwork , 在这些任务上 , 自注意力显示出了希望 , 但遇到了内存成本高昂和无法实际实现的问题 。 在 ImageNet 分类、COCO 目标检测和实例分割三方面的对照实验表明 , LambdaNetwork 显著优于基于卷积和注意力的同类方法 , 并且计算效率更高、运行速度更快 。
最后 , 研究者提出了 LambdaResNets , 它显著改善了图像分类模型的速度-准确性权衡 。 具体而言 , LambdaResNets 在实现 SOTA ImageNet 准确性的同时 , 运行速度是 EfficientNets 的4.5 倍左右 。
建模长程交互
在论文第二章开头 , 研究者定义了查询、上下文和交互等概念及各种用到的符号 , 此处不做赘述 。
注意力交互 。 首先缩小查询深度 , 创建查询和上下文元素之间的相似性核(注意力图) , 也就是所谓的注意力操作 。 该机制可以被看作激发查询、键和数值术语的可微内存的寻址 。 由于上下文位置 |m| 的数量越来越大 , 并且输入和输出维数 |k| 和 |v| 保持不变 , 所以在层输出是较小维数 |v| << |m| 的一个向量时 , 我们可以假设计算注意力图是不划算的 。
相反 ,由于 y_n = F((q_n, n), C) = λ(C, n)(q_n) 拟合一些线性函数 λ(C, n) , 所以通过线性函数将每个查询映射到输出可能更有效率 。 在这种场景下 , 上下文聚合到一个固定大小的线性函数 λ_n = λ(C, n) 。 每个 λ_n作为独立于上下文(计算后)存在的小的线性函数 , 并在应用于关联查询q_n之后被丢弃 。 该机制令人联想到与lambda这个术语相关的函数式编程和 λ 演算 。