PyTorch快餐教程2019 (2) - Multi-Head Attention-阿里云开发者社区

开发者社区> 人工智能> 正文

PyTorch快餐教程2019 (2) - Multi-Head Attention

简介: # PyTorch快餐教程2019 (2) - Multi-Head Attention 上一节我们为了让一个完整的语言模型跑起来,可能给大家带来的学习负担过重了。没关系,我们这一节开始来还上节没讲清楚的债。 还记得我们上节提到的两个Attention吗? ![两种Attention机制](https://upload-images.jianshu.io/upload_images/

PyTorch快餐教程2019 (2) - Multi-Head Attention

上一节我们为了让一个完整的语言模型跑起来,可能给大家带来的学习负担过重了。没关系,我们这一节开始来还上节没讲清楚的债。

还记得我们上节提到的两个Attention吗?
两种Attention机制

上节我们给大家一个印象,现在我们正式开始介绍其原理。

Scaled Dot-Product Attention

首先说Scaled Dot-Product Attention,其计算公式为:
$
Attention(Q,K,V)=softmax(frac{QK^T}{sqrt{d_k}})V
$

Q乘以K的转置,再除以$d_k$的平方根进行缩放,经过一个可选的Mask,经过softmax之后,再与V相乘。
用代码实现如下:

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

Multi-Head Attention

有了缩放点积注意力机制之后,我们就可以来定义多头注意力。

$
MultiHead(Q,K,V)=concat(head_1,...,head_n)W^O
$
其中,$head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)$
这个Attention是我们上面介绍的Scaled Dot-Product Attention.

这些W都是要训练的参数矩阵。
$
W_i^Qin mathbb{R}^{d_{model} times d_k},
W_i^Kinmathbb{R}^{d_{model} times d_k}, W_i^Vinmathbb{R}^{d_{model} times d_v}, W_oinmathbb{R}^{hd_v times d_{model}}
$
h是multi-head中的head数。在《Attention is all you need》论文中,h取值为8。
$d_k=d_v=d_{model}/h=64$
这样我们需要的参数就是d_model和h.

大家看公式有点要晕的节奏,别怕,我们上代码:

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "初始化时指定头数h和模型维度d_model"
        super(MultiHeadedAttention, self).__init__()
        # 二者是一定整除的
        assert d_model % h == 0
        # 按照文中的简化,我们让d_v与d_k相等
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

其中,clones是复制几个一模一样的模型的函数,其定义如下:

def clones(module, N):
    "生成n个相同的层"
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

Attention的逻辑主要分为4步。第一步是计算一下mask。

    def forward(self, query, key, value, mask=None):
        "实现多头注意力模型"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

第二步是将这一批次的数据进行变形 d_model => h x d_k

        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]

第三步,针对所有变量计算scaled dot product attention

        x, self.attn = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)

最后,将attention计算结果串联在一起,其实对张量进行一次变形:

        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

再看一种写法巩固一下

上面引用的代码来自:http://nlp.seas.harvard.edu/2018/04/03/attention.html

为了加深印象,我们再看另一种写法。
这个的命名更偏工程,d_model叫做hid_dim,h叫做n_heads,但是意思是一回事。

class SelfAttention(nn.Module):
    def __init__(self, hid_dim, n_heads, dropout, device):
        super().__init__()

        self.hid_dim = hid_dim
        self.n_heads = n_heads
        
        # d_model // h 仍然是要能整除,换个名字仍然意义不变
        assert hid_dim % n_heads == 0

        self.w_q = nn.Linear(hid_dim, hid_dim)
        self.w_k = nn.Linear(hid_dim, hid_dim)
        self.w_v = nn.Linear(hid_dim, hid_dim)

        self.fc = nn.Linear(hid_dim, hid_dim)

        self.do = nn.Dropout(dropout)

        self.scale = torch.sqrt(torch.FloatTensor([hid_dim // n_heads])).to(device)

下面是处理数据的过程:

    def forward(self, query, key, value, mask=None):

# Q,K,V计算与变形:

        bsz = query.shape[0]

        Q = self.w_q(query)
        K = self.w_k(key)
        V = self.w_v(value)

        Q = Q.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        K = K.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)
        V = V.view(bsz, -1, self.n_heads, self.hid_dim //
                   self.n_heads).permute(0, 2, 1, 3)

# Q, K相乘除以scale,这是计算scaled dot product attention的第一步

        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale

# 如果没有mask,就生成一个

        if mask is not None:
            energy = energy.masked_fill(mask == 0, -1e10)

# 然后对Q,K相乘的结果计算softmax加上dropout,这是计算scaled dot product attention的第二步:

        attention = self.do(torch.softmax(energy, dim=-1))

# 第三步,attention结果与V相乘

        x = torch.matmul(attention, V)

# 最后将多头排列好,就是multi-head attention的结果了

        x = x.permute(0, 2, 1, 3).contiguous()

        x = x.view(bsz, -1, self.n_heads * (self.hid_dim // self.n_heads))

        x = self.fc(x)

        return x

第二种实现取自:https://github.com/bentrevett/pytorch-seq2seq

版权声明:本文内容由阿里云实名注册用户自发贡献,版权归原作者所有,阿里云开发者社区不拥有其著作权,亦不承担相应法律责任。具体规则请查看《阿里云开发者社区用户服务协议》和《阿里云开发者社区知识产权保护指引》。如果您发现本社区中有涉嫌抄袭的内容,填写侵权投诉表单进行举报,一经查实,本社区将立刻删除涉嫌侵权内容。

分享:
人工智能
使用钉钉扫一扫加入圈子
+ 订阅

了解行业+人工智能最先进的技术和实践,参与行业+人工智能实践项目

其他文章