Bert不完全手册9. 长文本建模 BigBird & Longformer & Reformer & Performer( 三 )


  • LSH Attention:近似计算,针对l,只计算注意力中高权重的部分
  • 可逆网络:时间换空间,针对\(n_l\),只存储最后一层的参数
  • 分块计算:时间换空间,针对\(d_{ffn}\),对FFN层做分块计算
后两个方案和长文本无关这里我们简单过,重点是LSH Attention部分的创新~
Bert不完全手册9. 长文本建模 BigBird & Longformer & Reformer & Performer

文章插图
  1. LSH Attention
Local Sensitentive Hashing Attention是Reformer的主要贡献,也就是最初分类中的可学习pattern注意力机制 。考虑Attention的结果是被高权重的key主导的,因此每个token的注意力权重可以被部分高权重的token近似,只计算局部注意力从而避免计算\(L^2\)的注意力矩阵 。难点转换成了如何更高效的找到高权重的key,也就是和query token向量空间更相似的key token来进行局部交互,这里作者使用了LSH,一种在高维数据中快速近似查找的算法 。
LSH使用哈希函数对高位空间的向量x计算哈希函数h(x),\(h(x)\)满足在高维空间中更近的向量有更高的概率落在相同的哈希桶中 , 反之在高维空间中距离更远的向量有更低的概率会落在相同的哈希桶中 。LSH有很多种算法,这里作者使用的是基于角距离的局部敏感哈希算法 。随机初始化向量R维度是\(d_{model} * bucket/2\),哈希结果为旋转(xR)之后最近的一个正或者负的单位向量\(h(x) = argmax([xR;-xR])\)
使用LSH计算Attention会存在几个问题
  • query和key的hashing不同:为了解决这个问题作者把计算注意力之前query和key各自的线性映射统一成了一个,\(k_j=\frac{q_j}{||q_j||}\),这样二者的哈希也会相同,只需要对key进行计算就得到token的哈希分桶 。例如上图(b)长度为6的序列被分成3个桶[1,2,4],[3,6],[5]
  • 哈希的误差:哈希只是使得相似的向量落入相同桶的概率更高,为了进一步提高这个概率,可以进行多次不同的哈希函数对输出结果取交际,进一步降低近似带来的信息损失 。也就是用更多的时间和空间来换取更好的近似效果
  • 每个序列哈希分桶的大小可能不尽相同,无法进行batch计算:这里作者又做了一步近似 。根据以上的哈希结果对token进行重排序,相同哈希的token放一起,桶内按原始位置排序,按固定长度m进行切分,每个chunk的query对当前chunk和前一个chunk的key计算注意,也就是位于[m,2m]的query对[0,2m]的key计算注意力,这里m和哈希桶数反向相关\(m=\frac{l}{n_{bucket}}\) , 也就是和平均哈希桶的大小正相关 。实际上LSH只是用来排序,提高固定长度内注意力权重占整个序列的比例 , 从而通过有限长度的注意力矩阵近似全序列的注意力结果 。同样是固定窗口,LSH使得该窗口内的token权重会高于以上Longformer,BigBird这类完全基于位置的固定窗口的注意力机制,不过LSH的搜索和排序也会进一步提高时间复杂度
  1. 可逆残差网络
可逆残差的概念是来自The reversible residual network: Backpropagation without storing activations(Gomez2017) 。通过引入可逆变换,RevNet使得模型不需要存储中间层的参数计算梯度,每一层的参数可以由下一层通过可逆运算得到 。属于时间换空间的方案,因为反向传播计算梯度时需要先还原本层的参数,因此时间上会增加50%左右~ 细节我们就不多展开想看math的往苏神这看可逆ResNet:极致的暴力美学, 简单易懂的往大师兄这看可逆残差网络RevNet
  1. 分块计算
分块主要针对FFN层 。因为Feed Forward一般会设置几倍于Attention层的hidden size,通过先升维再降维的操作提高中间层的信息表达能力,优化信息的空间分布,以及抵消Relu带来的信息损失 。但是过大的hidden size会带来极高的空间占用 。因为是在embedding维度进行变换每个位置之间的计算独立,因此可以分块进行计算再拼接 , 用时间来换空间
效果评测部分我们在下面的performer里一起讨论
Performer
  • paper: Rethinking Attention with Performers
  • github: https://github.com/google-research/google-research/tree/master/performer
  • Take Away: 提出核函数使得QK变换后的内积可以近似注意力矩阵 , 配合乘法结合律把复杂度从平方降低到线性

Bert不完全手册9. 长文本建模 BigBird & Longformer & Reformer & Performer

文章插图

推荐阅读