[论文解读] Linear Complexity Randomized Self-attention Mechanism

Zheng Lin, Chong Wang+1 more|arXiv (Cornell University)|Apr 10, 2022

Domain Adaptation and Few-Shot Learning被引用 8

一句话总结

本文提出线性复杂度自注意力机制 Linear Randomized Attention (LARA)，该机制结合了无偏随机特征注意力 (RA) 估计器的表达能力与随机特征近似的高效性。通过将随机特征注意力 (RFA) 重新解释为自归一化重要性采样器，作者推导出 RA，其使用查询特定的分布以提高近似保真度，并通过引入多个查询相关提议分布和多重重要性采样，实现线性复杂度，同时在视觉、自然语言处理和视频任务上显著优于现有方法。

ABSTRACT

Recently, random feature attentions (RFAs) are proposed to approximate the softmax attention in linear time and space complexity by linearizing the exponential kernel. In this paper, we first propose a novel perspective to understand the bias in such approximation by recasting RFAs as self-normalized importance samplers. This perspective further sheds light on an \emph{unbiased} estimator for the whole softmax attention, called randomized attention (RA). RA constructs positive random features via query-specific distributions and enjoys greatly improved approximation fidelity, albeit exhibiting quadratic complexity. By combining the expressiveness in RA and the efficiency in RFA, we develop a novel linear complexity self-attention mechanism called linear randomized attention (LARA). Extensive experiments across various domains demonstrate that RA and LARA significantly improve the performance of RFAs by a substantial margin.

研究动机与目标

• 为解决现有随机特征注意力 (RFA) 方法中因估计过程的自归一化而对 Softmax 注意力产生近似偏差的问题。
• 通过将 RFA 重新解释为自归一化重要性采样器，构建一个查询特定的正随机特征估计器，从而实现对完整 Softmax 注意力的无偏估计。
• 将无偏估计器 (RA) 的高表达能力与 RFA 的线性复杂度相结合，构建一种新型高效注意力机制。
• 在保持线性时间与空间复杂度的前提下，实现在多样化序列建模任务上的最先进性能。

提出的方法

• 将 RFA 重新解释为自归一化重要性采样器，以识别在近似 Softmax 注意力时偏差的来源。
• 提出随机注意力 (RA)，一种无偏估计器，通过使用查询特定的分布构造正随机特征，从而提升近似保真度。
• 通过引入多个查询相关提议分布，将重要性采样框架进行泛化，每个分布针对查询子集进行定制。
• 应用具有自适应加权函数的多重重要性采样，以组合来自多个提议分布的估计结果，实现查询特定的自适应能力。
• 将加权函数中的查询无关与查询相关组件解耦，以提升估计效率与准确性。
• 使用高斯或高斯混合分布对每个提议分布进行参数化，并在 Transformer 框架内端到端进行训练。

实验结果

研究问题

• RQ1 为何现有随机特征注意力方法尽管对指数核的估计是无偏的，仍表现出近似偏差？
• RQ2 我们能否通过将近似过程重新解释为重要性采样，构建出对完整 Softmax 注意力的无偏估计器？
• RQ3 如何在保持 RFA 计算效率的同时，提升 RA 的近似保真度？
• RQ4 多重查询相关提议分布与自适应加权函数对注意力机制性能与复杂度有何影响？
• RQ5 所提方法能否在保持线性复杂度的同时，在多样化任务上实现更优性能？

主要发现

• RA 在近似保真度上优于 RFA，且在 LRA 基准测试中 5 项任务中有 3 项表现优于标准 Softmax 注意力，平均 Top-1 准确率达 59.30%。
• LARA 作为 RA 的线性复杂度变体，在 LRA 基准测试中实现 59.12% 的平均 Top-1 准确率，优于 Performer (57.63%) 及其他基线模型，在全部五项任务中均表现更优。
• 在 ImageNet1k 上使用 DeiT-Tiny 模型，LARA 达到 71.48% 的 Top-1 准确率，显著优于 Performer (65.92%) 和单提议变体 (68.42%)。
• 消融实验表明，与耦合版本相比，采用多重提议分布和解耦加权函数可使性能提升超过 0.4% 的准确率。
• 使用简单的高斯参数化提议分布即可获得与更复杂的高斯混合模型相当的性能，表明方法具备鲁棒性与高效性。