将注意力机制中的softmax函数进行替换

来源：投稿 作者：摩卡
编辑：学姐

论文解读

Sparse Attention with Linear Units

论文：

https://arxiv.org/abs/2104.07012

代码：

https://github.com/rishikksh20/rectified-linear-attention

阅读笔记

Motivation

最近一些研究者在研究将注意力机制中的softmax函数进行替换，使得attention score变得稀疏。虽然近期有一些方法在做这项工作，但是目前的工作仍存在一些缺点：

(1)大多数工作采用固定的稀疏模式，这使得模型缺少可扩展性以及泛化能力。

(2)其次，稀疏softmax函数的变体往往需要复杂的支撑算法(e.g., sorting operation)。

Method

为了解决这一问题，本文提出了名为Rectified Linear Attention（ReLA）。

模型如下图所示：

本文将文中的softmax函数换为ReLU函数以此来实现稀疏attention score的操作（因为ReLU函数天生具有稀疏性，所有小于0的x值经过ReLU函数激活，y值都会变为0）。并且仅仅是添加了ReLU函数便可以实现动态稀疏的过程。最后，为了促进优化，经过ReLA模块的特征进行了Layer Normalization（采用均方差LN），公式如下：

由于上述改进使得模型在刚开始训练的时候出现梯度不稳定的问题，导致收敛到局部最优解。本文也针对该问题提出了两个解决方案: (1)使用xavier初始化方法初始化RMSNorm，(2)使用简易门控单元进行正则化，公式如下所示。

Result

最终与SOTA进行比较实验，实验结果如下图所示

实验结果表明，与softmax激活函数作用于attention score相比较，使用本文提出的ReLA方法加快了训练的速度。

Table3是ReLA与SOTA相比较的结果，结果表明ReLA在提升速度的前提下仍然能达到SOTA的效果。

个人体会

本文针对Transformer模块中attention score所运用的softmax激活函数进行改进，提出了ReLA模块，在保证能达到SOTA的情况下提升了训练的速度，并且还解决了在训练早期的梯度不稳定的问题。

但是觉得为什么要进行sparse表述的不明确，以及进行sparse的优势在abstract和introduction没有指出，得自己去总结。

180+篇经典AI论文解读，点击卡片免费领取