Self-Attention with Relative Position Representations

ML_DL

Created At : 2023-10-20 20:08

Motivation
Recap: Why position embedding ? 回顾为什么需要 position embedding ?
Relation-aware Self-Attention 感知相对位置关系的自注意力机制
Relative Position Representations 相对位置表达
Core Code 核心代码分析
Reference

Motivation

1.针对 transformer 结构, 提出一种相对位置编码的方法, 能够在 self-attention 过程中将相对位置信息直接编码到单个 token 的表示学习中
2.提出高效实现融入相对位置编码方法下的 scaled dot self-attention 实现方式

Recap: Why position embedding ? 回顾为什么需要 position embedding ?

为什么要引入 position embedding ? 有时候一句话中相同的 token 因为具有不同的位置, 含义完全不同, 找几个极端的体会一下:

Can you can a can as a canner can can a can ? 你能像个罐头工人一样装罐头吗
一把把把把把住了
我也想过过过过过过的生活

稍微正常一点, 比如我们要对一句话 “I think therefore I am” 这句话编码, 第 1 个 token I 和第 4 个 token I, 在不同的上下文位置上因此有着不同的含义, 它俩因为不同的语境的位置, 应该有这不同的信息成分

Relation-aware Self-Attention 感知相对位置关系的自注意力机制

回顾 self-attention 操作, 先设定只有 1 个 head, 对于长度为 $n$ 的 token 序列 $x=(x_1,\ldots,x_n)$ , 其中 $x_i\in \mathbb R^{d_x}$, 计算一个等长度的结果 $z=(z_1,\ldots,z_n)$, 其中 $z_i\in \mathbb R^{d_z}$, 核心 attention 操作为

$\begin{align} z_i=&\sum_{j=1}^n\alpha_{ij}(x_jW^V) \\ \alpha_{ij}=&\frac{\exp e_{ij}}{\sum_{k=1}^n\exp e_{ik}} \\ e_{ij}=&\frac{(x_iW^Q)(x_jW^K)^T}{\sqrt{d_z}} \end{align}$

引入相对位置编码, 目的显式建模 pairwise 级别的 token 之间的相对位置关系, 假设在 $x_i$ 和 $x_j$ 之间的 edge (将 token 抽象成 graph 中的 nodes) 引入两套相对位置表达 $a_{ij}^V, a_{ij}^{K} \in \mathbb R^{d_a}$, 注意到相比传统的 self attention: 式子 (1) v.s. (4), 式子 (3) v.s. (6), 新增引入的相对位置编码分别影响到了 $V$ 和 $K$ 之上

$\begin{align} z_i=&\sum_{j=1}^n\alpha_{ij}(x_jW^V+a_{ij}^V) \\ \alpha_{ij}=&\frac{\exp e_{ij}}{\sum_{k=1}^n\exp e_{ik}} \\ e_{ij}=&\frac{(x_iW^Q)(x_jW^K+a_{ij}^K)^T}{\sqrt{d_z}} \end{align}$

Relative Position Representations 相对位置表达

所以这两套表达 $a_{ij}^V, a_{ij}^{K} \in \mathbb R^{d_a}$, 是怎么学出来的呢 ?

假设 $i$ 为我们研究的某个序列中间的位置, 它的左侧相对位置可以表示成 $i-k,\ldots,i-4, i-3, i-2, i-1$, 它的右侧侧相对位置可以表示成 $i+1, i+2, i+3, i+4, \ldots, i+k$

We hypothesized that precise relative position information is not useful beyond a certain distance

我们先确定一个最大的相对值 $k$, 我们假设相对位置如果超过 $k$ 的话就会距离太远没意义了, 所以我们根据相对值 $k$ 做个 clip, 在这种设定下, 我们可以保证至多有 $2k+1$ 范围个 edges

$a_{ij}^K=w_{\text{clip}(j-i,k)}^{K} \\ a_{ij}^V=w_{\text{clip}(j-i,k)}^{V} \\ \text{clip}(x,k)=\max(-k,\min(k,x))$

目标学出来 $w^K=(w_{-k}^K,\ldots,w_{k}^K)$ 和 $w^V=(w_{-k}^V,\ldots,w_{k}^V)$, 其中 $w_i^K,w_i^V\in \mathbb R^{d_a}$

在实现的时候, 式子 (6) 有个更高效的实现方式

$\begin{align} e_{ij}=&\frac{(x_iW^Q)(x_jW^K+a_{ij}^K)^T}{\sqrt{d_z}} \\ \Rightarrow e_{ij}&=\frac{(x_iW^Q)(x_jW^K)^T+(x_iW^Q)(a_{ij}^K)^T}{\sqrt{d_z}} \end{align}$

Core Code 核心代码分析

更新 scaled dot self attention 的实现方式

def _relative_attention_inner(x, y, z, transpose):
  # from https://github.com/tensorflow/tensor2tensor/blob/9e0a894034d8090892c238df1bd9bd3180c2b9a3/tensor2tensor/layers/common_attention.py#L1556-L1587
  """Relative position-aware dot-product attention inner calculation.
  This batches matrix multiply calculations to avoid unnecessary broadcasting.
  Args:
    x: Tensor with shape [batch_size, heads, len, length or depth].
    y: Tensor with shape [batch_size, heads, len, depth].
    z: Tensor with shape [len, length, depth].
    transpose: Whether to transpose inner matrices of y and z. 
    Should be true if last dimension of x is depth, not length.
  Returns:
    A Tensor with shape [batch_size, heads, length, length or depth].
  """
  batch_size = tf.shape(x)[0]
  heads = x.get_shape().as_list()[1]
  length = tf.shape(x)[2]

  xy_matmul = tf.matmul(x, y, transpose_b=transpose)                        # [b, heads, len, length or depth]

  x_t = tf.transpose(x, [2, 0, 1, 3])                                       # [len, b, heads, length or depth]

  x_t_r = tf.reshape(x_t, [length, heads * batch_size, -1])                 # [len, b * heads, length or depth]

  x_tz_matmul = tf.matmul(x_t_r, z, transpose_b=transpose)                  # [len, b * heads, length or depth]

  x_tz_matmul_r = tf.reshape(x_tz_matmul, [length, batch_size, heads, -1])  # [len, b, heads, length or depth]

  x_tz_matmul_r_t = tf.transpose(x_tz_matmul_r, [1, 2, 0, 3])               # [b, heads, len, length or depth]

  return xy_matmul + x_tz_matmul_r_t

Reference

[1]. Self-Attention with Relative Position Representations.
[2]. https://zhuanlan.zhihu.com/p/397269153.
[3]. https://medium.com/@_init_/how-self-attention-with-relative-position-representations-works-28173b8c245a

转载请注明来源, from goldandrabbit.github.io