# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Tensorflow中的Attention机制源码解析
#
# `tensorflow`中的`attention_wrapper`源码解析。
#

# ## Attention机制
#
# 首先，我们回顾一下`attention mechanism`的定义。
#
# 假设我们的source sequence为长度为$\text{n}$的序列$\mathbf{x}$，target sequence为长度为$\text{m}$的序列$\mathbf{y}$，则：
#
# $$\mathbf{x}=[x_0,x_1,\dots,x_n]$$
# $$\mathbf{y}=[y_0,y_1,\dots,y_m]$$
#
# 在第$i$个time step，source sequence产生一个`hidden state`：
#
# $$\boldsymbol{h}_i = [\overrightarrow{\boldsymbol{h}}_i^\top; \overleftarrow{\boldsymbol{h}}_i^\top]^\top, i=1,\dots,n$$
#
# 其中，$\overrightarrow{\boldsymbol{h}}_i^\top$和$\overleftarrow{\boldsymbol{h}}_i^\top$表示的是`Bidirectional RNN`产生的两个hidden state，如果不是使用BiRNN，则只有一个。
#
# 在第$t$个time step，decoder产生的`hidden state`表示为：
#
# $$\boldsymbol{s}_t=f(\boldsymbol{s}_{t-1}, y_{t-1}, \mathbf{c}_t),t=1,\dots,m$$
#
# 因为要产生target sequence的t时刻的输出，需要接收来自三个方向的信息，所以上面的公式很好理解。分别为：
#
# * $\boldsymbol{s}_{t-1}$，前一个时刻的状态
# * $y_{t_1}$，前一个时刻的输入序列
# * $\mathbf{c}_t$，当前时刻的context vector（上下文）
#
# 其中，$\mathbf{c}_t$就是`context vector`，表示`alignment scores`和source sequence的`hiddens states`的加权和：
#
# $$
# \begin{aligned}
# \mathbf{c}_t &= \sum_{i=1}^n \alpha_{t,i} \boldsymbol{h}_i & \small{\text{; Context vector for output }y_t}\\
# \alpha_{t,i} &= \text{align}(y_t, x_i) & \small{\text{; How well two words }y_t\text{ and }x_i\text{ are aligned.}}\\
# &= \frac{\text{score}(\boldsymbol{s}_{t-1}, \boldsymbol{h}_i)}{\sum_{i'=1}^n \text{score}(\boldsymbol{s}_{t-1}, \boldsymbol{h}_{i'})} & \small{\text{; Softmax of some predefined alignment score.}}.
# \end{aligned}
# $$
#
# 也就是说：
#
# $$\mathbf{c}_t = \sum_{i=1}^n \frac{\text{score}(\boldsymbol{s}_{t-1}, \boldsymbol{h}_i)}{\sum_{i'=1}^n \text{score}(\boldsymbol{s}_{t-1}, \boldsymbol{h}_{i'})} \boldsymbol{h}_i$$
#
# 根据公式可以看到：**$t$时刻的`context vector`需要计算所有的输入序列的隐状态**!
#
# $\alpha_{t,i}$表示的是第$t$个时刻的输入和第$i$个输出之间的关系，那么所有的$t$和$i$组成的$\alpha$就是**整个输入序列和输出序列之间的关系**！
#
# 所以，这个$\alpha$是一个二维矩阵！一般我们称呼它为`alignment matrix`！下面有一张示意图：
#
# <img src="images/attention_alignment_matrix.png" width=300 height=300>
#
#
# 那么，输入序列和输出序列的各个时间步之间的关系有多强烈呢？这就要用上面公式中的`score`函数来计算。
#
# `score`函数有很多选择，以下是常用的几种：
#
# * $\text{score}(\boldsymbol{s}_t, \boldsymbol{h}_i) = \mathbf{v}_a^\top \tanh(\mathbf{W}_a[\boldsymbol{s}_t; \boldsymbol{h}_i])$，加性注意力（additive attention）
# * $\text{score}(\boldsymbol{s}_t, \boldsymbol{h}_i) = \boldsymbol{s}_t^\top\mathbf{W}_a\boldsymbol{h}_i$，通用注意力（General）
# * $\text{score}(\boldsymbol{s}_t, \boldsymbol{h}_i) = \boldsymbol{s}_t^\top\boldsymbol{h}_i$，点积注意力（dot-product attention），属于乘性注意力（multiplicative）
# * $\text{score}(\boldsymbol{s}_t, \boldsymbol{h}_i) = \frac{\boldsymbol{s}_t^\top\boldsymbol{h}_i}{\sqrt{n}}$，上面的缩放版本
#
# 上面的所有$\text{W}_\alpha$和$\mathbf{v}_a^\top$都是可训练的权值矩阵。
#
# [《Attention?Attention!》](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html#a-family-of-attention-mechanisms)一文有一张比较全面的表格，展示了多种attention机制：
# ![attention_mechanism_family](images/attention_mechanism_table.png)

# ## Tensorflow中的Attention机制
#
# `AttentionMechanism`是一个抽象类，或者说是接口。
#
# 它的代码定义如下：
#

class AttentionMechanism(object):

  @property
  def alignments_size(self):
    raise NotImplementedError

  @property
  def state_size(self):
    raise NotImplementedError


# 从这个接口定义，我们可以看出，attention机制只关注两点：
#
# * alignments_size
# * state_size
#
# 在这里，我们看不出这两个属性是什么意思。那么我们先看看它的子类`_BaseAttentionMechanism`吧。
#
# 根据文档，`_BazeAttentionMechanism`是一个最基本的注意力机制，它提供几个通用的功能，如下：
#
# * 存储query_layer和memory_layer
# * 预处理并存储memory
#
# 什么是`query_layer`和`memory_layer`呢？我们留到后面解答。
#
# 首先看看`_BaseAttentionMechanism`的`__init__`函数，它的签名如下：
#
# |参数名                   |解释                                           |
# |:-----------------------|:---------------------------------------------|
# |query_layer             |`tf.layers.Layer`的子类                        |
# |memory                  |通常是RNN encoder的输出                          |
# |probability_fn          |概率计算函数，默认是`softmax`                     |
# |memory_sequence_length  |序列长度                                        |
# |memory_layer            |`tf.layers.Layer`的子类                         |
# |check_inner_dims_defined|是否检查内部维度                                 |
# |score_mask_value        |掩码值，默认是负无穷                              |
# |name                    |变量域                                          |
#
# 下面请看代码的注释，来理解初始化过程。
#

def __init__(self,
               query_layer,
               memory,
               probability_fn,
               memory_sequence_length=None,
               memory_layer=None,
               check_inner_dims_defined=True,
               score_mask_value=None,
               name=None):

    if (query_layer is not None　and not isinstance(query_layer, layers_base.Layer)):
      raise TypeError("query_layer is not a Layer: %s" % type(query_layer).__name__)
    if (memory_layer is not None　and not isinstance(memory_layer, layers_base.Layer)):
      raise TypeError("memory_layer is not a Layer: %s" % type(memory_layer).__name__)
    
    self._query_layer = query_layer
    self._memory_layer = memory_layer
    self.dtype = memory_layer.dtype
    if not callable(probability_fn):
      raise TypeError("probability_fn must be callable, saw type: %s" %　type(probability_fn).__name__)
    
    # 默认是负无穷
    if score_mask_value is None:
      score_mask_value = dtypes.as_dtype(self._memory_layer.dtype).as_numpy_dtype(-np.inf)
    
    self._probability_fn = lambda score, prev: (  # pylint:disable=g-long-lambda
        # _maybe_mask_score稍后解释
        probability_fn(_maybe_mask_score(score, memory_sequence_length, score_mask_value),prev))
    
    with ops.name_scope(
        name, "BaseAttentionMechanismInit", nest.flatten(memory)):
      # 稍后解释
      self._values = _prepare_memory(memory, memory_sequence_length,
                                     check_inner_dims_defined=check_inner_dims_defined)
      # shape = [B, T,...]
      self._keys = self.memory_layer(self._values) if self.memory_layer else self._values
      # 获取keys的第一个维度，即batch size
      self._batch_size = self._keys.shape[0].value or array_ops.shape(self._keys)[0]
      # 获取keys的第二个维度，即alignments size，这就是序列的最大长度
      self._alignments_size = self._keys.shape[1].value or array_ops.shape(self._keys)[1]



# ### `_maybe_mask_score`
#
# 输入的张量score的形状是[B,T,..]，并且这个张量是对齐的，也就是每个序列的T这个维度的值都是一样的。但是实际上，我们同一批序列中，序列的长短是不一定一样的，也就是说我们的score张量中的某些值是无效的。我们需要把这些无效的值的位置找出来，然后给这些位置设置一个特定的数值，使得在这些位置的影响很小，趋于0。一般来说，设置一个**负无穷**。根据原来的序列的真实长度张量，以及我们的score张量，这样就可以找出哪些位置是无效的。当然，这要结合`probability_fn`函数，因为只有根据这个函数的性质，我们才知道赋予什么值，会使得这个函数的值趋于0。
#
# 这也就说明了：如果我们的输入序列是已经对齐了的，那么我们可以直接设置`memory_sequence_length=None`，这样不就需要
# 进一步处理了！
#
# 这个函数的签名解释如下：
#
# |参数名                 |解释                                                |
# |:---------------------|:---------------------------------------------------|
# |score                 |分数张量，也就是我们上面提到的`score`函数的结果，形状是[B,T]|
# |memory_sequence_length|输入数据的长度张量，形状是[B]                           |
# |score_mask_value      |掩码值，默认是负无穷                                   |
#
#
# 这个函数的实现如下：
#

def _maybe_mask_score(score, memory_sequence_length, score_mask_value):
  # 如果序列长度未指定，则不进行掩码计算
  if memory_sequence_length is None:
    return score
  message = ("All values in memory_sequence_length must greater than zero.")
  with ops.control_dependencies(
      [check_ops.assert_positive(memory_sequence_length, message=message)]):
    # 进行sequence mask，得到的是一个Boolean类型（或者只有0和1）的张量，score元素为0的地方为False，其他位置为True
    score_mask = array_ops.sequence_mask(memory_sequence_length, maxlen=array_ops.shape(score)[1])
    # mask张量
    score_mask_values = score_mask_value * array_ops.ones_like(score)
    # 返回mask之后的张量，原来score位置为0的地方替换成了score_mask_value
    return array_ops.where(score_mask, score, score_mask_values)



# ### `_prepare_memory`
#
# 这个函数是把输入张量转化成掩码之后的张量。为什么这里需要进行掩码处理呢？
#
# 首先memory张量的形状是[B,T,...]，其中T是这一批次的数据中的最大长度。为了搞清楚这T个time step之间，哪些是有效位置，哪些是无效的（通常进行对齐添加的标记），我们需要把这些无效的位置找出来，然后数值上设为0。这样这些无效位置就不会影响后面的计算。
#
# 当然，如果你手动对齐了输入序列（在较短的序列后面补0），保证每个序列的长度都一样，那么你可以设置参数`memory_sequence_length=None`，这样就不需要额外处理了。同样的，即使你对齐处理过，你也可以传入一个`memory_sequence_length`张量，这样不会有影响。
#
# 签名如下：
#
# |参数名                   |解释                           |
# |:-----------------------|:-----------------------------|
# |memory                  |输入张量，形状为[B,T,...]        |
# |memory_sequence_length  |序列长度张量，形状为[B]           |
# |check_inner_dims_defined|Boolean，检查最外面两个维度是否定义|
#
# 注意，掩码只是作用于[B,T]两个维度，之后如果还有其他维度，应该处理使之不进行掩码计算。这个在代码中有很明显的处理。
#
# 下面，详细解释代码：
#

# +

def _prepare_memory(memory, memory_sequence_length, check_inner_dims_defined):
    memory = nest.map_structure(
        lambda m: ops.convert_to_tensor(m, name="memory"), memory)
    if memory_sequence_length is not None:
        memory_sequence_length = ops.convert_to_tensor(memory_sequence_length,
                                                       name="memory_sequence_length")

    # 检查innder dim
    if check_inner_dims_defined:
        def _check_dims(m):
            # 检查第三个维度及其之后的维度是否存在
            if not m.get_shape()[2:].is_fully_defined():
                raise ValueError(
                    "Expected memory %s to have fully defined inner dims, "
                    "but saw shape: %s" % (m.name, m.get_shape()))
            nest.map_structure(_check_dims, memory)

    if memory_sequence_length is None:
        seq_len_mask = None
    else:
        # 进行sequence_mask
        seq_len_mask = array_ops.sequence_mask(
            memory_sequence_length,
            maxlen=array_ops.shape(nest.flatten(memory)[0])[1],
            dtype=nest.flatten(memory)[0].dtype)
        # 获取batch size
        seq_len_batch_size = (memory_sequence_length.shape[0].value or
                              array_ops.shape(memory_sequence_length)[0])

    def _maybe_mask(m, seq_len_mask):
        # 获取维度数量，也就是阶
        rank = m.get_shape().ndims
        rank = rank if rank is not None else array_ops.rank(m)
        # 取出B,T之外的维度，设置成1，当计算掩码的时候，这些维度的值会保留原来的值
        extra_ones = array_ops.ones(rank - 2, dtype=dtypes.int32)
        # 获取batch size
        m_batch_size = m.shape[0].value or array_ops.shape(m)[0]
        if memory_sequence_length is not None:
            message = (
                "memory_sequence_length and memory tensor batch sizes do not match.")
            with ops.control_dependencies([
                check_ops.assert_equal(seq_len_batch_size, m_batch_size,
                                       message=message)]):
                # 将seq_len_mask重塑成[B,T,...]形状的张量
                seq_len_mask = array_ops.reshape(
                    seq_len_mask,
                    array_ops.concat((array_ops.shape(seq_len_mask), extra_ones), 0))
                # 掩码计算，并且返回结果
                return m * seq_len_mask
        else:
            return m

    return nest.map_structure(lambda m: _maybe_mask(m, seq_len_mask), memory)


# -

# 以上首先就可以回答`AttentionMechanism`类的第一个函数是什么这个问题。因为`_BaseAttentionMechansim`的实现很简单：
#

  @property
  def alignments_size(self):
    # __init__函数里面，self._alignments_size = self._keys.shape[1].value or array_ops.shape(self._keys)[1]
    # 实际上就是这一批序列的最大长度
    return self._alignments_size

# 那么第二个问题，`state_size`又是什么呢？在`_BaseAttentionMechanism`里面的实现，和上面一样：
#

  @property
  def state_size(self):
    return self._alignments_size

# 也就是说，默认情况下，`alignments_size`和`state_size`是一样的！
#
# 那么为什么这两个值默认是一样的呢？
#
# 回顾文章开始的attention机制的定义，我们发现知道，`alignment matrix`的两个维度的大小，分别是输入序列和输出序列的长度，并且，在编解码的过程中，每一个时间步都会产生一个`hidden state`，所以`hidden state`的数量也是和序列长度一样的。
#
# 而我们这里的`alinments_size`就是对应`alignment matrix`，`state`就是对应`hidden state`的！
#

# ### Query，Keys和Values
#
# 你可能已经发现了，在`_BaseAttentionMechanism`的`__init__`函数里，出现了`self._query_layer`、`self._memory_layer`、`self._keys`和`self._values`等成员变量。那么这些变量指的是什么呢？
#
# [Attention is all you need](https://arxiv.org/pdf/1706.03762.pdf)一文有一个对attention机制的解释：
#
# > An attention function can be described as mapping a query and a set of key-values pairs to an output, where the query, keys, values and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.
#
# 翻译一下就是：
#
# > 注意力函数可以描述为一个映射过程：把query和key-value键值对映射到一个输出。其中query，keys, values和output都是矢量。output是values的加权和，其中每一个value的权重都是通过一个关于key和query的兼容性函数计算得来的。
#
# 我再加一点解释：
#
# * 所谓的关于key和query的兼容性函数，就是我们的`score`函数！
#
# 所以关于attention机制，可以再精简一点，概括为：**Attention机制就是，从key-value键值对中，通过确定query和key的相似度，然后再通过value获得最终输出的映射过程。**
#
# 你已经发现了，上面的定义都是围绕`keys`、`query`和`values`进行的！！！
#
# 那么言归正传，什么是`keys`、`values`和`query`呢？
#
# 在`_BaseAttentionMechanism`里面，`keys`和`values`的关系是这样的：
#
# ```python
# self._values = _prepare_memory(
#           memory, memory_sequence_length,
#           check_inner_dims_defined=check_inner_dims_defined)
# self._keys = (
#           self.memory_layer(self._values) if self.memory_layer  # pylint: disable=not-callable
#           else self._values)
# ```
# `_prepare_memory`上面已经解释过了。所以：
#
# * `values`就是`memory或者masked memory`！（memory一般是RNN encoder的输出`encoder_output`！）
# * `keys`就是`values`或者`memory_layer处理过的values`！(keys和values可以认为是一样的！)
#
# 那么，这里为什么没有出现`query`呢？很简单，因为`query`是在调用的时候传入的参数，所以这里没有`query`！具体的可以从`_BaseAttentionMechanism`的子类的`__call__`函数解读出来。
#
# 总结一下：
#
# * `values`就是`encoder_output`或者`masked encoder_output`
# * `keys`就是`values`或者是经过`memory_layer`处理后的`values`
# * `query`是调用的时候传入的参数，一般与来自RNN decoder
#
# 疑问: **这里的`query`、`keys`和`values`怎么和`score`函数对应起来？**
# （提示，可以参考attention_wrapper代码。）
#
# 这个问题留到后面解释。
#

# ### memory_layer和query_layer
#
# 上面还有两个疑惑，`memory_layer`和`query_layer`是什么？
#
# `_BaseAttentionMechanism`没有解答，但是它的子类`LuongAttention`和`BahdanauAttention`可以找到答案。
#
# `LuongAttention`的`__init__`函数有定义：
#
# ```python
# query_layer=None,
# memory_layer=layers_core.Dense(
#             num_units, name="memory_layer", use_bias=False, dtype=dtype)
# ```
#
# 在`LuongAttention`里，`query_layer`是`None`，`memory_layer`只是一个**全连接层**！，并且激活函数默认是`None`，也就是说**memory_layer只是对memory做一个简单的线性变换！**
#
# `BahdanauAttention`是这样定义的：
#
# ```python
# query_layer=layers_core.Dense(
#             num_units, name="query_layer", use_bias=False, dtype=dtype)
# memory_layer=layers_core.Dense(
#             num_units, name="memory_layer", use_bias=False, dtype=dtype)
# ```
#
# 在`BahdanauAttention`里，`query_layer`和`memory_layer`都是一个全连接层，激活函数默认也是`None`，所以`query_layer`和`memory_layer`分别是对`query`和`memory`的线性变换。
#
# 现在`memory_layer`和`query_layer`已经清楚了，但是现在又有一个新的问题：**`LuongAttention`和`BahdanauAttention`有什么区别呢？**
#

# ### LuongAttention和BahdanauAttention
#
# 不同的attention机制，最大和最关键的区别就在于它们各自的`score`函数。
#
# 我们首先来看看`LuongAttention`的`score`函数：
#
# ```python
# def _luong_score(query, keys, scale):
#   # query shape: [batch_size, depth]
#   # keys shape: [batch_size, max_time, depth]
#   depth = query.get_shape()[-1]
#   key_units = keys.get_shape()[-1]
#   if depth != key_units:
#     raise ValueError(
#         "Incompatible or unknown inner dimensions between query and keys.  "
#         "Query (%s) has units: %s.  Keys (%s) have units: %s.  "
#         "Perhaps you need to set num_units to the keys' dimension (%s)?"
#         % (query, depth, keys, key_units, key_units))
#   dtype = query.dtype
#
#   # 为后面的矩阵乘法做准备，在第二个维度上扩充一个维度
#   # [batch_size, depth] -> [batch_size, 1, depth]
#   query = array_ops.expand_dims(query, 1)
#
#   # [batch_size, 1, depth]*[batch_size, depth, max_time] -> [batch_size, 1, max_time]
#   score = math_ops.matmul(query, keys, transpose_b=True)
#   # 将第二个维度压缩掉，[batch_size, 1, max_time] -> [batch_size, max_time]
#   score = array_ops.squeeze(score, [1])
#   
#   # 缩放，可选项
#   if scale:
#     # Scalar used in weight scaling
#     g = variable_scope.get_variable(
#         "attention_g", dtype=dtype,
#         initializer=init_ops.ones_initializer, shape=())
#     score = g * score
#   return score
# ```
#
# `LuongAttention`计算`socre`的过程很简单：**query和keys做矩阵乘法而已！**
#
# 从数学上来说，这里的`score`计算公式如下：
#
# $$\text{score}(\boldsymbol{s}_t, \boldsymbol{h}_i) = \boldsymbol{s}_t^\top\boldsymbol{h}_i$$
#
# 所以，`LuongAttention`属于**multiplicative attention**！，具体一点，就是**dot-product attention**!
#
# 我们继续看看`BahdanauAttention`的`score`函数：
#
# ```python
# def _bahdanau_score(processed_query, keys, normalize):
#   # processed_query shape: [batch_size, depth]
#   # keys shape: [batch_size, max_time, depth]
#   dtype = processed_query.dtype
#   # 获取depth(num_units)，是第三个维度
#   num_units = keys.shape[2].value or array_ops.shape(keys)[2]
#   # 将query变形，[batch_size, depth] -> [batch_size, 1, depth]
#   processed_query = array_ops.expand_dims(processed_query, 1)
#   v = variable_scope.get_variable(
#       "attention_v", [num_units], dtype=dtype)
#   # 归一化，可选项
#   if normalize:
#     # Scalar used in weight normalization
#     g = variable_scope.get_variable(
#         "attention_g", dtype=dtype,
#         initializer=init_ops.constant_initializer(math.sqrt((1. / num_units))),
#         shape=())
#     # tanh激活函数的偏置
#     b = variable_scope.get_variable(
#         "attention_b", [num_units], dtype=dtype,
#         initializer=init_ops.zeros_initializer())
#     # normed_v = g * v / ||v||
#     normed_v = g * v * math_ops.rsqrt(
#         math_ops.reduce_sum(math_ops.square(v)))
#     # query和keys直接相加
#     return math_ops.reduce_sum(
#         normed_v * math_ops.tanh(keys + processed_query + b), [2])
#   else:
#     # query和keys直接相加
#     return math_ops.reduce_sum(v * math_ops.tanh(keys + processed_query), [2])
# ```
#
# 可以看出，`BahdanauAttention`属于**additive attention**！
#



# ## 参考文章
#
# 1.[Attention?Attention!](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html#a-family-of-attention-mechanisms)  
#

StarcoderdataJupyterScriptsDedupFiltered