多头自注意力

使用 PyTorch 的 Transformer 模型

James Chapman

Curriculum Manager, DataCamp

Transformer 的多头注意力

论文《Attention Is All You Need》中展示的 Transformer 架构。

自注意力机制

自注意力机制剖析

自注意力机制

自注意力机制剖析

Q：指示每个 token 在其他 token 中"寻找"什么
K：表示每个 token 的内容
V：要聚合或加权的实际内容

自注意力机制

自注意力机制剖析

Q：指示每个 token 在其他 token 中"寻找"什么
K：表示每个 token 的内容
V：要聚合或加权的实际内容

注意力分数：Q-K 相似度 → _点积_

自注意力机制

自注意力机制剖析

Q：指示每个 token 在其他 token 中"寻找"什么
K：表示每个 token 的内容
V：要聚合或加权的实际内容

注意力分数：Q-K 相似度 → 点积
注意力权重：softmax 缩放

自注意力机制

自注意力机制剖析

注意力分数：Q-K 相似度 → 点积
注意力权重：softmax 缩放

	Orange	is	my	favorite	fruit
查询：	Orange
注意力权重：	.21	.03	.05	.31	.40

自注意力机制

自注意力机制剖析

注意力分数：Q-K 相似度 → 点积
注意力权重：softmax 缩放

	Orange	is	my	favorite	fruit
查询：	Orange
注意力权重：	.21	.03	.05	.31	.40

自注意力机制

自注意力机制剖析

注意力分数：Q-K 相似度 → 点积
注意力权重：softmax 缩放

	Orange	is	my	favorite	fruit
查询：	Orange
注意力权重：	.21	.03	.05	.31	.40

多头注意力

多头自注意力

多头注意力

多头自注意力

多头注意力

多头自注意力

import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):

    def __init__(self, d_model, num_heads):
        super().__init__()

        assert d_model % num_heads == 0, "d_model must be divisible by num_heads."

        self.num_heads = num_heads
        self.d_model = d_model
        self.head_dim = d_model // num_heads

        self.query_linear = nn.Linear(d_model, d_model, bias=False)
        self.key_linear = nn.Linear(d_model, d_model, bias=False)
        self.value_linear = nn.Linear(d_model, d_model, bias=False)

        self.output_linear = nn.Linear(d_model, d_model)

num_heads：注意力头数量，每个处理 head_dim 大小的嵌入
bias=False：在降低复杂度的同时不影响性能（仅对输入）

    def split_heads(self, x, batch_size):

        seq_length = x.size(1)
        x = x.reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        return x.permute(0, 2, 1, 3)


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

        scores = torch.matmul(query, key.transpose(-2, -1)) / (self.head_dim ** 0.5)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attention_weights = F.softmax(scores, dim=-1)
        return torch.matmul(attention_weights, value)


    def combine_heads(self, x, batch_size):

        x = x.permute(0, 2, 1, 3).contiguous()
        return x.view(batch_size, -1, self.d_model)

compute_attention()：使用 F.softmax() 计算注意力权重
torch.matmul(attention_weights, value)：对 value 的加权求和

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        query = self.split_heads(self.query_linear(query), batch_size)
        key = self.split_heads(self.key_linear(key), batch_size)
        value = self.split_heads(self.value_linear(value), batch_size)

        attention_weights = self.compute_attention(query, key, value, mask)


        output = self.combine_heads(attention_weights, batch_size)

        return self.output_linear(output)

self.output_linear()：拼接并投影各头的输出

Passons à la pratique !

使用 PyTorch 的 Transformer 模型