智能家居科技|谷歌大改Transformer注意力,自己挖坑自己填( 二 )


智能家居科技|谷歌大改Transformer注意力,自己挖坑自己填
文章图片
标准注意力矩阵包括每一对entry的相似度系数 , 由query和key上的softmax计算组成 , 表示为q和k 。
常规的softmax注意力可以看作是由指数函数和高斯投影定义的非线性函数的一个特例 。 在这里我们也可以反向推理 , 首先实现一些更广义的非线性函数 , 隐式定义query-key结果中其他类型的相似性度量或核函数 。 研究者基于早期的核方法(kernelmethod) , 将其定义为广义注意力(generalizedattention) 。 尽管对于大多核函数来说 , 闭式解并不存在 , 但这一机制仍然可以应用 , 因为它并不依赖于闭式解 。
该研究首次证明了 , 任意注意力矩阵都可以通过随机特征在下游Transformer应用中实现有效地近似 。 实现这一点的的新机制是使用正随机特征 , 即原始query和key的正直非线性函数 , 这对于避免训练过程中的不稳定性至关重要 , 并实现了对常规softmax注意力的更准确近似 。
新算法FAVOR+:通过矩阵相关性实现快速注意力
上文描述的分解允许我们以线性而非二次内存复杂度的方式存储隐式注意力矩阵 。 我们还可以通过分解获得一个线性时间注意力机制 。 虽然在分解注意力矩阵之后 , 原始注意力机制与具有值输入的存储注意力矩阵相乘以获得最终结果 , 我们可以重新排列矩阵乘法以近似常规注意力机制的结果 , 并且不需要显式地构建二次方大小的注意力矩阵 。 最终生成了新算法FAVOR+ 。
智能家居科技|谷歌大改Transformer注意力,自己挖坑自己填
文章图片
左:标准注意力模块计算 , 其中通过执行带有矩阵A和值张量V的矩阵乘法来计算最终的预期结果;右:通过解耦低秩分解A中使用的矩阵Q′和K′以及按照虚线框中指示的顺序执行矩阵乘法 , 研究者获得了一个线性注意力矩阵 , 同时不用显式地构建A或其近似 。
上述分析与双向注意力(即非因果注意力)相关 , 其中没有past和future的概念 。 对于输入序列中没有注意前后token的单向(即因果)注意力而言 , 研究者稍微修改方法以使用前缀和计算(prefix-sumcomputation) , 它们只存储矩阵计算的运行总数 , 而不存储显式的下三角常规注意力矩阵 。
智能家居科技|谷歌大改Transformer注意力,自己挖坑自己填
文章图片
左:标准单向注意力需要mask注意力矩阵以获得其下三角部分;右:LHS上的无偏近似可以通过前缀和获得 , 其中用于key和值向量的随机特征图的外积(outer-product)前缀和实现动态构建 , 并通过query随机特征向量进行左乘计算 , 以在最终矩阵中获得新行(newrow) 。
性能
研究者首先对Performer的空间和时间复杂度进行基准测试 , 结果表明 , 注意力的加速比和内存减少在实证的角度上近乎最优 , 也就是说 , 这非常接近在模型中根本不使用注意力机制的情况 。
智能家居科技|谷歌大改Transformer注意力,自己挖坑自己填
文章图片
在以时间(T)和长度(L)为度量的双对数坐标轴中 , 常规Transformer模型的双向timing 。
研究者进一步证明 , 使用无偏softmax近似 , 该Performer模型在稍微进行微调之后可以向后兼容预训练Transformer模型 , 从而在提升推理速度的同时降低能耗 , 并且不需要从头训练预先存在的模型 。
智能家居科技|谷歌大改Transformer注意力,自己挖坑自己填
文章图片
在OneBillionWordBenchmark(LM1B)数据集上 , 研究者将原始预训练Transformer的权重迁移至Performer模型 , 使得初始非零准确度为0.07(橙色虚线) 。 但在微调之后 , Performer的准确度在很少的梯度步数之后迅速恢复 。