cosFormer阅读笔记

论文（ICLR 2022）：cosFormer: Rethinking Softmax in Attention

• Q1 论文试图解决什么问题？
  • 过去的线性Transformer设计中，通常采用核方法近似Softmax，但是近似误差较大。
• Q2 这是否是一个新的问题？
  • 不是的。过去已经有了一些Linear Transformer的设计方案。
• Q3 这篇文章要验证一个什么科学假设？
  • 高效而精确地逼近Softmax是比较困难的，因此产生问题：“能否使用线性方法代替Softmax，同时保留它的主要特性？”
• Q4 有哪些相关研究？如何归类？谁是这一课题在领域内值得关注的研究员？
  • 关于Efficient Transformer，可以分成以下几类：
    1. 知识蒸馏、剪枝、量化等：比较通用，但是不降低模型的复杂度；
    2. 稀疏矩阵分解， $O(n\sqrt{n})$复杂度；
    3. 核方法：通过将Softmax去掉，并应用矩阵乘法的结合律，将 $O(n^2d)$ 的复杂度变为  $O(nd^2)$ 。代表性方法：Efficient Attention, Linear Attention。
• Q5 论文中提到的解决方案之关键是什么？
  • 关键之处是抓住Softmax的两个主要性质：
    1. Attention Map的非负性；
    2. 非线性加权方案使得Attention Map的分布较为集中

    并针对性地设计了基于ReLU的线性映射函数，和基于cos的重新加权方案。

• Q6 论文中的实验是如何设计的？
  • 使用Adaptive input representations for neural language modeling作为baseline，更改其中的Attention机制并进行微调，在速度、精度上进行对比。
• Q7 用于定量评估的数据集是什么？代码有没有开源？
  • 自回归语言模型中，在WikiText-103上进行了验证；在LRA任务上对比了多种任务；代码已经开源。
• Q8 论文中的实验及结果有没有很好地支持需要验证的科学假设？
  • 吞吐量能够接近Linear Transformer，且效果甚至能超越原始Transformer。很好地验证了假设。
• Q9 这篇论文到底有什么贡献？
  • 提出了cosFormer代替Softmax，实现了线性的注意力机制，在大大加快推理速度的同时，能够有接近甚至超越原始Transformer的表现。
• Q10 下一步呢？有什么工作可以继续深入？
  • 同样的方法能否迁移到视觉模型上？在视觉模型上是否能够同样表现很好？

0 Abstract

过去的Linear Transformer工作中，通常采用核方法近似Softmax。但是由于近似误差的存在，它们的性能下降较大。本文提出的cosFormer线性Transformer，基于Softmax的两个主要特点：

• Attention Map的非负性；
• 非线性加权方案使得Attention Map的分布较为集中

cosFormer通过线性操作和基于cosine的加权机制，满足了以上特性，实现了与Softmax相当的精度。

1 Introduction

已有的加速方法中：

• 通常对于Attention Map引入不切实际的假设；
• 或者只能在受约束的理论范围内有效近似Softmax；
• 还有一些方法是针对Cross Attention设计的，不能用于任意的Attention。

高效而精确地逼近Softmax是比较困难的，因此产生问题：“能否使用线性方法代替Softmax，同时保留它的主要特性？”

Linear Transformer解决了第一条，但是由于缺乏重新加权的方案，它在Long-Range Arena（LRA）上比其他Transformer表现地差一些。

2 Our Method

• 使用 $\phi_{\text{linear}}(x) = \text{ReLU}(\cdot)$ 作为映射函数（代替Linear Transformer中的 $\text{elu}(x)+1$）
• Attention即可使用以下公式表示：

$$O_i = \frac{\text{ReLU}(Q_i)\sum_{j=1}^N \text{ReLU}(K_j)^{\top}V_j}{\text{ReLU}(Q_i)\sum_{j=1}^N \text{ReLU}(K_j)^T}$$

一些对Softmax的感性理解与假设：

• 这种非线性的加权方法，能够使得Attention Score更为集中，使得训练更加稳定；
• 它可以惩罚长距离连接，并强制短距离连接：在NLP任务中，很大一部分上下文相关性来自于相邻的token。

基于cos的加权机制：

• • 托勒密定理确保cos权重可以分解为两个求和；
  • cos会给相邻的token增加更多的权重，可以加强locality；

$$\begin{aligned}Q_i’ K_j’^{\top} \cos\left(\frac{\pi}{2} \times \frac{i – j}{M}\right)  &=  Q_i’ K_j’^{\top} \left( \cos\left(\frac{\pi i}{2M}\right) \cos\left(\frac{\pi j}{2M}\right) + \sin\left(\frac{\pi i}{2M}\right) \sin\left(\frac{\pi j}{2M}\right) \right) \\ &= \left( Q_i’ \cos\left(\frac{\pi i}{2M}\right) \right) \left( K_j’ \cos\left(\frac{\pi j}{2M}\right) \right)^{\top} + \left( Q_i’ \sin\left(\frac{\pi i}{2M}\right) \right) \left( K_j’ \sin\left(\frac{\pi j}{2M}\right) \right)^{\top}\end{aligned}$$

$$\mathcal{O} = S(Q, K)V = \left( Q^{cos} K^{cos} + Q^{sin} K^{sin} \right) V = Q^{cos} (K^{cos} V) + Q^{sin} (K^{sin} V)$$

• cosFormer 还可以看作一种将相对位置偏差引入到高效 transformer 中的新方法。