zoukankan      html  css  js  c++  java
  • [NLP] Reformer: The Efficient Transformer

    1.现状

    Transformer模型目前被广泛应用,但目前存在以下几个问题:

    (1) 模型层数加深

    (2) 模型参数量变大

    (3) 难以训练

    (4) 难以fine-tune

    2. 单层参数量和占用内存分析

    参数设置参数量与占用内存
    1 layer 0.5Billion 0.5Billion * 4Byte = 2GB
    embedding layer

    64K tokens 

    1024 emb_size

    8 batch_size

    参数量 64K × 1K × 8 = 0.5B 

    内存 2GB

    3. Transformer 模型内存占用的问题以及Reformer相应解决方案

    从以上来看,似乎训练一个单层的Transformer在内存开销上是可以接受的,即使训练大型BERT也只是大概需要17G,但是为什么目前训练Transformer如此困难呢?

    因为2中只考虑了单层的内存消耗,而对于N层的网络结构,其对内存资源的消耗将大增。其中主要消耗资源的部分如下:

    (1) N层网络比N个1层网络消耗更多内存。原因:为了反向传播时梯度的计算,在正向传播时,需要保留每一层的输入Activations的值。

    (2) FeedForward层的维度 $d_{ff}$ 远大于 $d_{model}$,在N层网络中,FeedForward层参数量巨大,大量占用内存。

    (3) Attention机制中涉及到的序列L的emb两两点积运算,当序列长度非常长时,$QK^T$这个是 [batch_size, length, length] 大小,对于batch_size = 8, length = 64k,单纯这个矩阵参数将占据:

    8 * 64k * 64k * 4byte = 16GB。除了这个很大的存储开销外,还有两两点积的计算开销。这些占用了大量的内存。

    针对以上问题,Reformer中提出了相应的解决方案,归纳如下。采用Reformer的创新技术点后,相对于原模型的影响见"对模型的影响"。

    Transformer 内存占用的问题Reformer采用的方案实现的效果对模型的影响 
    memory(N layer) > N * memory(1 layer)  → 需要为反向传播保留值

    Reversible layers 

    保留N层activations → 保留1层

    negligible 微不足道

     
    feed-forward 层 $d_{ff}$ >> $d_{model}$, 占据大量内存

    split activation &&

    process in chunks

    $d_{ff}$ → 每个chunk中处理的维度

    numerically identical 

    数值等价,计算方式不同

     
    Attention 运算: $O(L^2)$:存储 && 计算开销都很大

    Locality-Sensitive Hashing  

    局部敏感哈希,简称 LSH

    $O(L^2)$  → $O(L log L)$

    可以更好的处理长序列

    major change 

    但是找到了合适的参数

    使得与原模型效果相近

     

    4. Attention → LSH

    4.1 Recall Attention && Thoughts

    (1) dot-product attention && multi-head attention

    Q, K, V [batch_size, length, d_model]

    $QK^T$ → [batch_size, length, length] 

    ex. length = 64K  → 64K * 64K * 4Byte = 16GB → 该部分占用内存过大

    (2) 对公式 (1) 的优化整体思路:

    (a) 优化QK^T  

    → 4.2 把大矩阵运算降低为向量矩阵运算;

    → 4.3 Q = K

    (b) 优化 softmax(x)  

    → 4.4 softmax更关注x中数值大的那些元素,在这里也就是只找和 $q_i$ 相似的那些 $k_j$,它们点积值更大,而其他元素被忽略不参与计算节省开销。

    → 4.5 但是如何判断哪些 $k_j$ 和 $q_i$ 更相近呢? 这里引入LSH的思想。

    → 4.6 如何将LSH应用于Attention? LSH Attention

    → 4.7 用一轮hash,有些相近元素并不能hash到一个桶里 →  多来几轮hash取并集 Multi-round LSH attention

    4.2 优化 $QK^T$ → Memory-efficient attention

    $QK^T$ → $q_iK^T$   => 对每个query单独计算,而无需计算整个大矩阵乘积

    4.3 优化 $QK^T$ → Shared-QK Transformer (Q = K)

    TransformerReformer (LSH attention)

    A → linear projection 1 → Q

    A → linear projection 2 → K

    A → linear projection 3 → V

    A → linear projection 1 → Q

    A → linear projection 1 → K

    A → linear projection 3 → V

    4.4 优化softmax(x) → Hashing Attention

    $QK^T$只是一个中间值,我们最关注的还是softmax之后的这个结果。

    softmax 更关注大数值的元素,对于length = 64K, 我们找到和q_i 最相近的32/64个 $k_j$,即可 大大降低计算消耗。

    4.5 优化softmax(x) → Locality sensitive hashing (star)(star)(star)

    (1) LSH

    核心思想:相近的向量vectors被hash到同一个bucket的概率大。

     

    ~CS369G: Algorithmic Techniques for Big Data. Lecture 16: Approximate Nearest Neighbor Search. Spring 2015-2016

    (2) hash函数如何选?

    为了得到b个hash值,我们首先固定一个随机矩阵R [$d_k$, b/2] 

    定义:h(x) = arg max([xR; −xR]) 

    xR的物理意义:一个矩阵乘以一个矩阵,在物理意义上是对一个矩阵施加另一个矩阵表示的线性变化。结合上图相当于,将原向量映射到的圆上进行不同角度的顺时针/逆时针的旋转变换。

    Figure1 上方图所示,向量x, y映射到圆上的初始位置是最左图所示。

    右边彩色图1是相对最左原始图位置,顺时针旋转$ heta_0$度,其中x点移动到了图中红圈位置,对应区域的hash值是0;

    右边彩色图2是相对最左原始图位置,逆时针旋转$ heta_1$度,其中x点移动到了图中红圈位置,对应区域的hash值是2;

    右边彩色图3是相对最左原始图位置,逆时针旋转$ heta_2$度,其中x点移动到了图中红圈位置,对应区域的hash值是1。

    综上,x的hash值是 0, 2, 1; 同理分析y的hash值分别是 3 2 0. 可知二者差别较大,是不相近向量。

    同理分析Figure1 下方图,可知x, y hash值均相同,可知x, y是相近的向量。

    这里将比较两个向量是否相近转化为两个向量进行hash后hash值是否相同的比较。如果向量的hash值相同,则向量大概率相近

    ~ LSH scheme  Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya P. Razenshteyn, and Ludwig Schmidt. Practical and optimal LSH for angular distance. CoRR, abs/1509.02897, 2015. URL http://arxiv. org/abs/1509.02897.  

    4.6 优化softmax(x) → LSH Attention (star)(star)(star)(star)(star)

    有了LSH机制,如何应用于Attention?

     

    → 

    → 

    其中,,表示和$q_i$ 的hash值相同的k_j的集合。

    $z(i, P_i)$可以看做是softmax中归一化的操作和scale的操作等。

    (2) 中只取了和$q_i$相近的向量$k_j$进行计算,但是为了方便对一个batch的计算,可能会引入一些和$q_i$不相近的向量,对于这部分向量,可以用加mask的方式使其不参与计算。$m(j, Pi)$表示对这部分向量加mash的计算。

    Figure 2

    (a) Normal → sparse

    (b) Bucketed → q, k 根据它们hash值是否相同进行排序。颜色相同的hash值相同。

    因为hash值相同的vector接近,所以将(a)中对所有vector两两点积可以近似替换为只对bucket内部vector计算点积。

    (c) Q = K 

    → 问题:每个bucket中q, k 个数可能不相等,不好计算。

    → 解决方法:

    首先通过设定 $k_j = q_j / ||q_j||$  来保证:$h(k_j) = h(q_j) $

    根据bucket number 来sort bucket

    在bucket 内部根据 seqence position 来sort。$i$ → $s_i$

    → 效果:bucket中的pair 集中在对角附近。

    (d) Chunked

    如何chunk?

    chunk个数: $m = 2l / n_{buckets}$ 其中l是序列长度

    每个bucket平均size = $l / n_{buckets}$ → 认为size变为2倍大不太容易 → 不容易超过chunk大小

    4.7 优化softmax(x) → Multi-round LSH attention

    用一轮hash,有些相近元素并不能hash到一个桶里 →  多来几轮hash取并集

    并行来做

    4.8 Causal masking for shared-QK attention

    Transformer中原有的mask: mask后面的元素

    LSH attention中元素顺序变化,公式(3)mask如何实现?

    为每个query, key加入position index,根据与原来一样的排列重新排序,用一个比较的操作来实现mask.

    4.9 复杂度分析

    4l?

    4.10 实验效果

    5. Reversible Transormer

    5.1 Recall RevNets   (star)(star)(star)(star)(star)

    反向传播求导:$F = AX^T$

    F对A求梯度:X  → 每层的输入

    所以获取每层的输入值很重要~~

     ~ The Reversible Residual Network: Backpropagation Without Storing Activations 

    residual layer 类型输入输出form反向传播计算梯度时输入值来源
    Normal residual layer x → y y = x + F(x) 正向传播时每层保留
    reversible residual layer (x1, x2) → (y1, y2)

    y1 = x1 + F(x2)

    y2 = x2 + G(y1)

    根据上一层activations推理获得

    x2 = y2 −G(y1) 

    x1 = y1 − F(x2) 

    5.2 Reversible Transformer  (star)(star)(star)(star)(star)

    如何将reversible residual layer思想应用于Transformer?

     ~ Attention is all you need

    reversible residual layerTransformer
    y1 = x1 + F(x2) Y1 = X1 + Attention(X2)
    y2 = x2 + G(x1) Y2 = X2 + FeedForward(Y1)

    5.3 Chunking  

    → 问题:5.2 是对各个残差层进行的优化,但是在每个残差层内部,FeedForward占用内存也很高。

    ex. Attention is all you need 中 $d_{ff}$ = 2048 >> $d_{model}$ = 512

    → 解决方法:由于FeedForward层对于sequence的各个position是独立的,所以可以分解为c个chunk

    在5.2中FeedForward层:

    → 这里对FeedForward层进行chunk,相应的也要对reversible的正向和反向计算都要chunk。

    → 对于vocab非常大的情况,对output的log概率计算也chunk,并且会计算那个section sequence的loss.

    5.4 复杂度分析

    6. 他山之石

    Reformer中并没有提出一种全新的算法,但是却对以往算法提出了适合Transformer模型的应用。

    LSH本来用于大规模向量中求相近向量,reformer中用于softmax中 找与$q_i$ 相近的 $k_j$向量,并根据attention计算的需要,设计了LSH attention。

    Reversible resnet 本来用于ResNet, 这里替换了F, G函数为 Attention 和 FeedForward;

    Chunk思想也广为人知,这里用于 FeedForward 矩阵的分解。

    综合起来,这些技术对于Transformer模型降低存储与计算开销有非常大的帮助。

    7. 参考文献

    (1) Reformer:Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. "Reformer: The efficient transformer.arXiv preprint arXiv:2001.04451 (2020).

    (2) LSH相关:Andoni, Alexandr, et al. "Practical and optimal LSH for angular distance.Advances in neural information processing systems. 2015.

    (3) LSH相关:CS369G: Algorithmic Techniques for Big Data. Lecture 16: Approximate Nearest Neighbor Search. Spring 2015-2016

    (4) LSH相关:CS369G: Algorithmic Techniques for Big Data. Lecture 17: Locality Sensitive Hashing and Dimension Reduction. Spring 2015-2016

    (5) Reversible ResNet: Gomez, Aidan N., et al. "The reversible residual network: Backpropagation without storing activations.Advances in neural information processing systems. 2017.

    (6) Transformer: Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems. 2017.

    (7) Reformer Pytorch Github: https://github.com/lucidrains/reformer-pytorch

  • 相关阅读:
    Dropdown 下拉菜单
    Breadcrumb 面包屑
    Tabs 标签页
    NavMenu 导航菜单
    Notification 通知
    MessageBox 弹框
    Message 消息提示
    Loading 加载
    Alert 警告
    Badge 标记
  • 原文地址:https://www.cnblogs.com/shiyublog/p/13222231.html
Copyright © 2011-2022 走看看