【Datawhale X 李宏毅苹果书 AI夏令营】Task3 自注意力机制

【Datawhale X 李宏毅苹果书 AI夏令营】Task1 Chapter3.1&3.2：
【Datawhale X 李宏毅苹果书 AI夏令营】Task2 Chapter3.3&4%5

1.从0.5开始入门

• 本blog仅作为夏令营过程中学习出现的问题和任务进行记录和学习留档，针对苹果书内的知识点不作详细记录（例如完全复制公式，思维导图等形式）。

• 本文根据夏令营进度，从3.1开始。

• 网络上有众多的学习资料，如果你恰好点了进来，对机器学习、深度学习或者仅对Datawhale感兴趣，可以点击下面的链接来详细了解。

  • **李宏毅深度学习教程LeeDL-Tutorial（苹果书）**开源地址：https://github.com/datawhalechina/leedl-tutorial
  • 李宏毅《机器学习/深度学习》2021课程B站视频：https://www.bilibili.com/video/BV1JA411c7VT?p=1&vd_source=7f728b80e21aaffa0f2781c650cbe2ce
  • 本人参与的**Datawhale AI夏令营（第五期）**地址（虽然不能报名，但仍可以进入学习说明和Task文件进行自学）：https://linklearner.com/activity/16
    • 附官网网址：Datawhale (linklearner.com)
    • 附B站官号（有直播回放）：二次元的Datawhale的个人空间-二次元的Datawhale个人主页-哔哩哔哩视频 (bilibili.com)
  • 本人自己之前接触过的一篇实践性学习框架Approaching (Almost) Any Machine Learning Problem（简称AAAMLP）：
    • 原文github地址（作者Abhishek Thakur）：https://github.com/abhishekkrthakur/approachingalmost
    • 中译版出处：https://ytzfhqs.github.io/AAAMLP-CN/

2. 自注意力机制

2.1 定义

Self-Attention，自注意力机制，又称内部注意力机制，顾名思义，是一种将单个序列的不同位置关联起来以计算同一序列的表示的注意机制。

通过对注意力机制的学习我们知道，在一般任务的Encoder-Decoder框架中，输入Source和输出Target内容是不一样的，比如对于英-中机器翻译来说，Source是英文句子，Target是对应的翻译出的中文句子，Attention机制发生在Target的元素Query和Source中的所有元素之间（即Attention机制与自身还有关注对象都有关系）。

而Self-Attention顾名思义，指的不是Target和Source之间的Attention机制，而是Source内部元素之间或者Target内部元素之间发生的Attention机制，也可以理解为Target=Source这种特殊情况下的注意力计算机制。（即Self-Attention只关注输入本身or只关注关注对象本身）

2.2 优缺点

1. 优点：可以建立全局的依赖关系，扩大图像的感受野。相比于CNN，其感受野更大，可以获取更多上下文信息。

2. 缺点：自注意力机制是通过筛选重要信息，过滤不重要信息实现的，这就导致其有效信息的抓取能力会比CNN小一些。这样是因为自注意力机制相比CNN，无法利用图像本身具有的尺度，平移不变性，以及图像的特征局部性（图片上相邻的区域有相似的特征，即同一物体的信息往往都集中在局部）这些先验知识，只能通过大量数据进行学习。这就导致自注意力机制只有在大数据的基础上才能有效地建立准确的全局关系，而在小数据的情况下，其效果不如CNN。

2.3 与CNN的对比

在全局建模能力上，自注意力机制具有明显的优势，因为它可以显式地捕捉序列中任意两个元素之间的关系（无论它们之间的距离）。这使得自注意力机制在处理长距离依赖和全局信息方面非常强大。
而CNN则在局部特征提取方面非常有效，但在全局建模能力上可能不如自注意力机制。然而，通过设计特定的网络结构（比如使用全局池化层或多尺度卷积），CNN也可以在一定程度上捕捉全局信息。

2.4 自注意力公式

计算过程：

1. 计算Q和K的点积，得到形状为[n, n] 的相似度矩阵。
2. 对相似度矩阵应用softmax函数，得到权重矩阵。
3. 将权重矩阵与V相乘，得到加权后的输出矩阵。

给定查询矩阵 Q、键矩阵 K，则相似度矩阵S=QK`

2.5计算实例

2.6代码实例

import numpy as np
import torch
from torch import Tensor
from typing import Optional, Any, Union, Callable
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time


class MultiHeadedAttention(nn.Module):
    def __init__(self, 
                num_heads: int, 
                d_model: int, 
                dropout: float=0.1):
        super(MultiHeadedAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        # Assume v_dim always equals k_dim
        self.k_dim = d_model // num_heads
        self.num_heads = num_heads
        self.proj_weights = clones(nn.Linear(d_model, d_model), 4) # W^Q, W^K, W^V, W^O
        self.attention_score = None
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, 
                query:Tensor, 
                key: Tensor, 
                value: Tensor, 
                mask:Optional[Tensor]=None):
        """
        Args:
            query: shape (batch_size, seq_len, d_model)
            key: shape (batch_size, seq_len, d_model)
            value: shape (batch_size, seq_len, d_model)
            mask: shape (batch_size, seq_len, seq_len). Since we assume all data use a same mask, so
                  here the shape also equals to (1, seq_len, seq_len)
        
        Return:
            out: shape (batch_size, seq_len, d_model). The output of a multihead attention layer
        """
        if mask is not None:
            mask = mask.unsqueeze(1)
        batch_size = query.size(0)
        
        # 1) Apply W^Q, W^K, W^V to generate new query, key, value
        query, key, value \
            = [proj_weight(x).view(batch_size, -1, self.num_heads, self.k_dim).transpose(1, 2)
                for proj_weight, x in zip(self.proj_weights, [query, key, value])] # -1 equals to seq_len
        
        # 2) Calculate attention score and the out
        out, self.attention_score = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)
        
        # 3) "Concat" output
        out = out.transpose(1, 2).contiguous() \
             .view(batch_size, -1, self.num_heads * self.k_dim)

        # 4) Apply W^O to get the final output
        out = self.proj_weights[-1](out)
        
        return out



def clones(module, N):
        "Produce N identical layers."
        return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
    
    
def attention(query: Tensor, 
              key: Tensor, 
              value: Tensor, 
              mask: Optional[Tensor] = None, 
              dropout: float = 0.1):
    """
    Define how to calculate attention score
    Args:
        query: shape (batch_size, num_heads, seq_len, k_dim)
        key: shape(batch_size, num_heads, seq_len, k_dim)
        value: shape(batch_size, num_heads, seq_len, v_dim)
        mask: shape (batch_size, num_heads, seq_len, seq_len). Since our assumption, here the shape is
              (1, 1, seq_len, seq_len)
    Return:
        out: shape (batch_size, v_dim). Output of an attention head.
        attention_score: shape (seq_len, seq_len).

    """
    k_dim = query.size(-1)

    # shape (seq_len ,seq_len)，row: token，col: that token's attention score
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(k_dim)
        
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e10)

    attention_score = F.softmax(scores, dim = -1)

    if dropout is not None:
        attention_score = dropout(attention_score)
        
    out = torch.matmul(attention_score, value)
    
    return out, attention_score # shape: (seq_len, v_dim), (seq_len, seq_lem)


if __name__ == '__main__':
    d_model = 8
    seq_len = 3
    batch_size = 6
    num_heads = 2
    # mask = None
    mask = torch.tril(torch.ones((seq_len, seq_len)), diagonal = 0).unsqueeze(0)
    
    input = torch.rand(batch_size, seq_len, d_model)
    multi_attn = MultiHeadedAttention(num_heads = num_heads, d_model = d_model, dropout = 0.1)
    out = multi_attn(query = input, key = input, value = input, mask = mask)
    print(out.shape)