Transformer架构

基于编码器-解码器架构来处理序列对

跟使用注意力的seq2seq不同，Transformer是纯基于注意力（或者说是纯基于self.attention的架构，里面没有RNN了）。

seq2seq里面加了个注意力进来

现在是：把里面的RNN全部换成transformer

最后一层的输出作为之后层的输入来完成信息传递

---

多头注意力

对同一key、value、query，希望抽取不同的信息。
    例如：短距离关系和长距离关系

多头注意力使用h个独立的注意力池化
    合并各个头输出得到最终输出

query（q∈$R^{d_q}$）

key（k∈$R^{d_k}$）

value（v∈$R^{d_v}$）

头i可学习参数${W_i}^{(q)}$∈$R^{d_q\times d_q}$、${W_i}^{(k)}$∈$R^{d_k\times d_k}$、${W_i}^{(v)}$∈$R^{d_v\times d_v}$

头i的输出$h_i$=f(${W_i}^{(q)}$q,${W_i}^{(k)}$k,${W_i}^{(v)}$v)∈$R^{p_v}$

输出的可学习参数$W_o$∈$R^{p_o\times p_v}$

多头注意力的输出

---

有掩码的多头注意力

解码器对序列中一个元素输出时，不应该考虑该元素之后的元素。
（因为attention没有时间信息的）

解决方法：
    可以通过掩码来实现
        也就是计算$x_i$输出时，假装当前序列长度为i。（也就是把i后面的用掩码掩藏掉，就是在算softmax上的时候不会算它的权重，不会对后面的key、value给权重）

---

基于位置的前馈网络

其实就是一个全连接层

1、将输入形状由(b,n,d)变换成(bn,d)
b：批量大小
n：序列长度
d：特征维度（每个元素的嵌入向量维度）

之前在卷积的时候，把nd换成一维{变成了(b,nd)}

2、作用两个全连接层
3、输出形状由(bn,d)变化成(b,n,d)
4、等价于两层核窗口为1的一维卷积层

---

层归一化

批量归一化对每个特征/通道里元素进行归一化
    不适合序列长度会变的NLP应用
层归一化对每个样本里元素进行归一化

---

信息传递

编码器中的输出$y_1$,…,$y_n$

将其作为解码中第i个Transformer块中多头注意力的key和value
    它的query来自目标序列
意味着编码器和解码器中块的个数和输出维度都是一样的

（中间这个）

---

预测

预测第t+1个输出时
解码器中输入前t个预测值
    在自注意力中，前t个预测值作为key和value，第t个预测值还作为query

---

总结

    ①Transformer是一个纯使用注意力的编码-解码器。

    ②编码器和解码器都有n个transformer块

    ③每个块里使用多头(自)注意力，基于位置的前馈网络，和层归一化。

---

代码

多头注意力

import math
import os
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l

---

使用多个头并行计算

形状X从[batch_size，查询或者“键－值”对的个数，num_hiddens]
➡[batch_size，查询或者“键－值”对的个数，num_heads，num_hiddens/num_heads]
➡[batch_size，num_heads，查询或者“键－值”对的个数，num_hiddens/num_heads]
➡[batch_size*num_heads，查询或者“键－值”对的个数，num_hiddens/num_heads]

通过将 batch_size 和 num_heads 维度合并，可以在一个批量中同时计算所有头的注意力，提高了计算效率。

def transpose_qkv(X, num_heads):"""为了多注意力头的并行计算而变换形状"""# 2，4，100➡2，4，5, 20# 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)# 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，num_hiddens/num_heads)# 这么做的目的：是为了在多头注意力机制中，将隐藏维度分成多个头，每个头处理一部分特征。X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)# 将第二维度和第三维度交换一下顺序# 在多头自注意力机制中，通常将num_heads维度提前，以便在每个头上独立执行注意力计算，提高计算效率# 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,num_hiddens/num_heads)# 2，5，4，20X = X.permute(0, 2, 1, 3)# 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,num_hiddens/num_heads)# 10，4，20return X.reshape(-1, X.shape[2], X.shape[3])# 在多头注意力机制的最后一步，将经过多头注意力计算后的张量恢复到原始形状。这样可以确保输出张量的形状与输入张量的形状一致
def transpose_output(X, num_heads):"""逆转transpose_qkv函数的操作"""# 10，4，20➡2，5，4, 20# 将 batch_size 和 num_heads 分离，以便后续操作。X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])# 2，5，4, 20➡2，4，5, 20# 恢复 num_heads 维度的位置，使其与 transpose_qkv 函数的逆操作一致。X = X.permute(0, 2, 1, 3)# 2，5，4, 20➡2，4，100return X.reshape(X.shape[0], X.shape[1], -1)

---

选择缩放点积注意力作为每一个注意力头

class MultiHeadAttention(nn.Module):"""多头注意力"""# 100、100、100、100、5、0.5def __init__(self, key_size, query_size, value_size, num_hiddens, num_heads, dropout, bias=False, **kwargs):super(MultiHeadAttention, self).__init__(**kwargs)# 5self.num_heads = num_heads# 用的是缩放点积注意力所以里面不需要学习w了self.attention = d2l.DotProductAttention(dropout)# (2, 4, 100)self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)# (2, 6, 100)self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)# (2, 6, 100)self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)def forward(self, queries, keys, values, valid_lens):# queries，keys，values的形状:#                            (batch_size，查询或者“键－值”对的个数，num_hiddens)# valid_lens　的形状:#                            (batch_size，)或(batch_size，查询的个数)# 线性变换：self.W_q、self.W_k和self.W_v的作用是将输入的查询、键和值张量映射到一个新的隐藏空间，学习新的表示，并投影到统一的隐藏空间。# 经过transpose_qkv变换后，输出的queries，keys，values　的形状:#                            (batch_size*num_heads，查询或者“键－值”对的个数，num_hiddens/num_heads)# (10, 4, 20)queries = transpose_qkv(self.W_q(queries), self.num_heads)# (10, 6, 20)keys = transpose_qkv(self.W_k(keys), self.num_heads)# (10, 6, 20)values = transpose_qkv(self.W_v(values), self.num_heads)if valid_lens is not None:# 在轴0，将第一项（标量或者矢量）复制num_heads次，确保每个头都有对应的有效长度。# valid_lens：torch.Size([2])➡torch.Size([10])valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)# 根据查询（queries）、键（keys）和值（values）计算注意力权重，并据此加权求和得到输出。# output的形状:(batch_size*num_heads，查询的个数，num_hiddens/num_heads)# (10, 4, 20)output = self.attention(queries, keys, values, valid_lens)# 逆转 transpose_qkv 函数的操作，将输出恢复到原始的批次和序列维度。# output_concat的形状:(batch_size，查询的个数，num_hiddens)output_concat = transpose_output(output, self.num_heads)return self.W_o(output_concat)

---

测试

num_hiddens, num_heads = 100, 5
# 100\100\100\100\5\0.5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens, num_hiddens, num_heads, 0.5)
attention.eval()
batch_size, num_queries = 2, 4
# 包含了两个整数元素3和2，所以形状是(2,)
num_kvpairs, valid_lens = 6, torch.tensor([3, 2])
# X：2，4，100
# Y：2，6，100
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
# 输出形状[2, 4, 100]
print(attention(X, Y, Y, valid_lens).shape)

---

该部分总代码

import torch
from torch import nn
from d2l import torch as d2l# 重塑和转置
def transpose_qkv(X, num_heads):"""为了多注意力头的并行计算而变换形状"""X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)X = X.permute(0, 2, 1, 3)return X.reshape(-1, X.shape[2], X.shape[3])# 在多头注意力机制的最后一步，将经过多头注意力计算后的张量恢复到原始形状。这样可以确保输出张量的形状与输入张量的形状一致
def transpose_output(X, num_heads):"""逆转transpose_qkv函数的操作"""X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])X = X.permute(0, 2, 1, 3)return X.reshape(X.shape[0], X.shape[1], -1)class MultiHeadAttention(nn.Module):"""多头注意力"""def __init__(self, key_size, query_size, value_size, num_hiddens, num_heads, dropout, bias=False, **kwargs):super(MultiHeadAttention, self).__init__(**kwargs)# 5self.num_heads = num_headsself.attention = d2l.DotProductAttention(dropout)self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)def forward(self, queries, keys, values, valid_lens):queries = transpose_qkv(self.W_q(queries), self.num_heads)keys = transpose_qkv(self.W_k(keys), self.num_heads)values = transpose_qkv(self.W_v(values), self.num_heads)if valid_lens is not None:valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)output = self.attention(queries, keys, values, valid_lens)output_concat = transpose_output(output, self.num_heads)return self.W_o(output_concat)num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens, num_hiddens, num_heads, 0.5)
attention.eval()
batch_size, num_queries = 2, 4
num_kvpairs, valid_lens = 6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
print(attention(X, Y, Y, valid_lens).shape)

架构图:

步骤：通过将多头注意力的计算转化为单头注意力的计算，可以简化实现逻辑

逻辑图：

数据形状步骤图：

最终线性变换的作用：

    1、self.W_o 的作用是将这些拼接后的结果重新映射回一个统一的隐藏空间，确保输出的维度与输入的维度一致。
    2、这个线性变换可以看作是一个权重矩阵，用于将多头注意力的结果进行加权组合，生成最终的输出。
    3、可以学习到更复杂的表示，从而更好地捕捉输入数据的特征和模式。

---

Transformer

import math
import os
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l

---

基于位置的前馈网络

pytorch的densel默认的实现：当输入不是二维的时候，把前面的维度都当作样本维，后面的维度当成特征维。

基于位置的前馈网络：实际上是一个单隐藏层的MLP。

class PositionWiseFFN(nn.Module):def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):super(PositionWiseFFN, self).__init__(**kwargs)# 第一个全连接层，将输入维度从ffn_num_input映射到ffn_num_hiddensself.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)# ReLU激活函数，增加非线性特性self.relu = nn.ReLU()# 第二个全连接层，将输入维度从ffn_num_hiddens映射到ffn_num_outputsself.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)def forward(self, X):# 前向传播过程# 通过第一个全连接层，将输入X映射到隐藏层# 应用ReLU激活函数，增加非线性特性# 通过第二个全连接层，将隐藏层映射到输出层return self.dense2(self.relu(self.dense1(X)))

---

改变张量的最里层维度的尺寸

其实就是把4变成8

# 创建一个PositionWiseFFN实例，输入维度为4，隐藏层维度为4，输出维度为8
ffn = PositionWiseFFN(4, 4, 8)
# 设置模型为评估模式，不进行训练
ffn.eval()
# 创建一个形状为(2, 3, 4)的张量，所有元素都设置为1
# 输入张量通过前馈网络进行前向传播，得到输出张量
# 取第一个样本的输出张量
X = torch.ones((2, 3, 4))
print(ffn(X)[0])

---

对比不同维度的层归一化和批量归一化的效果

layer norm：把每个样本变成均值为0方差为1

batch norm：每个特征变成均值为0方差为1

# 创建一个对最后一个维度进行层归一化的实例
ln = nn.LayerNorm(2)
# 创建一个对最后一个维度进行批量归一化的实例
bn = nn.BatchNorm1d(2)
# 创建一个形状为(2, 2)的张量
X = torch.tensor([[1, 2], [2, 3]], dtype=torch.float32)
# 对张量进行层归一化
# 对张量进行批量归一化
# 打印层归一化和批量归一化的结果
print('layer norm:', ln(X), '\nbatch norm:', bn(X))

---

使用残差连接和层归一化

class AddNorm(nn.Module):def __init__(self, normalized_shape, dropout, **kwargs):super(AddNorm, self).__init__(**kwargs)# 定义一个dropout层，用于随机丢弃部分神经元self.dropout = nn.Dropout(dropout)# 定义一个层归一化层，对输入进行归一化self.ln = nn.LayerNorm(normalized_shape)def forward(self, X, Y):# 前向传播过程# 使用残差连接：将dropout(Y)与X相加# 应用层归一化：对残差进行归一化return self.ln(self.dropout(Y) + X)

---

加法操作后输出张量的形状相同

# 创建一个AddNorm实例，输入的归一化维度为[3, 4]，dropout率为0.5
add_norm = AddNorm([3, 4], 0.5)
# 设置模型为评估模式，不进行训练
add_norm.eval()
# 创建一个形状为(2, 3, 4)的张量作为X
# 创建一个形状为(2, 3, 4)的张量作为Y
# 输入X和Y通过AddNorm进行前向传播
# 打印输出结果的形状
print(add_norm(torch.ones((2, 3, 4)), torch.ones((2, 3, 4))).shape)

该部分总代码

import torch
from torch import nnclass AddNorm(nn.Module):def __init__(self, normalized_shape, dropout, **kwargs):super(AddNorm, self).__init__(**kwargs)self.dropout = nn.Dropout(dropout)self.ln = nn.LayerNorm(normalized_shape)def forward(self, X, Y):return self.ln(self.dropout(Y) + X)add_norm = AddNorm([3, 4], 0.5)
add_norm.eval()
print(add_norm(torch.ones((2, 3, 4)), torch.ones((2, 3, 4))).shape)

---

实现编码器中的一个层

ffn：前馈神经网络

class EncoderBlock(nn.Module):def __init__(self, key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,dropout, use_bias=False, **kwargs):super(EncoderBlock, self).__init__(**kwargs)# 多头注意力机制self.attention = d2l.MultiHeadAttention(key_size, query_size,value_size, num_hiddens,num_heads, dropout, use_bias)# 第一个残差连接和层归一化模块self.addnorm1 = AddNorm(norm_shape, dropout)# 位置前馈网络self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)# 第二个残差连接和层归一化模块self.addnorm2 = AddNorm(norm_shape, dropout)def forward(self, X, valid_lens):# 第一个残差连接和层归一化模块# 使用多头注意力机制计算X的自注意力输出，并将该输出与X相加（残差连接）# 然后进行层归一化Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))# 位置前馈网络：对第一个残差连接和层归一化模块的输出进行位置前馈网络操作# 第二个残差连接和层归一化模块：# 将位置前馈网络输出与第一个残差连接和层归一化模块的输出相加并进行归一化return self.addnorm2(Y, self.ffn(Y))

Transformer编码器中的任何层都不会改变其输入的状态

# 创建一个形状为(2, 100, 24)的张量作为输入X
X = torch.ones((2, 100, 24))
# 创建一个形状为(2,)的张量作为有效长度
valid_lens = torch.tensor([3, 2])
encoder_blk = EncoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5)
# 设置模型为评估模式，不进行训练
encoder_blk.eval()
# 输入X和有效长度通过EncoderBlock进行前向传播
print(encoder_blk(X, valid_lens).shape)

---

该部分总代码

import math
import os
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l# 基于位置的前馈网络
class PositionWiseFFN(nn.Module):def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):super(PositionWiseFFN, self).__init__(**kwargs)# 第一个全连接层，将输入维度从ffn_num_input映射到ffn_num_hiddensself.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)# ReLU激活函数，增加非线性特性self.relu = nn.ReLU()# 第二个全连接层，将输入维度从ffn_num_hiddens映射到ffn_num_outputsself.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)def forward(self, X):# 前向传播过程# 通过第一个全连接层，将输入X映射到隐藏层# 应用ReLU激活函数，增加非线性特性# 通过第二个全连接层，将隐藏层映射到输出层return self.dense2(self.relu(self.dense1(X)))# 使用残差连接和层归一化
class AddNorm(nn.Module):def __init__(self, normalized_shape, dropout, **kwargs):super(AddNorm, self).__init__(**kwargs)# 定义一个dropout层，用于随机丢弃部分神经元self.dropout = nn.Dropout(dropout)# 定义一个层归一化层，对输入进行归一化self.ln = nn.LayerNorm(normalized_shape)def forward(self, X, Y):# 前向传播过程# 使用残差连接：将dropout(Y)与X相加# 应用层归一化：对残差进行归一化return self.ln(self.dropout(Y) + X)# 实现编码器中的一个层
class EncoderBlock(nn.Module):def __init__(self, key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,dropout, use_bias=False, **kwargs):super(EncoderBlock, self).__init__(**kwargs)# 多头注意力机制self.attention = d2l.MultiHeadAttention(key_size, query_size,value_size, num_hiddens,num_heads, dropout, use_bias)# 第一个残差连接和层归一化模块self.addnorm1 = AddNorm(norm_shape, dropout)# 位置前馈网络self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)# 第二个残差连接和层归一化模块self.addnorm2 = AddNorm(norm_shape, dropout)def forward(self, X, valid_lens):# 第一个残差连接和层归一化模块# 使用多头注意力机制计算X的自注意力输出，并将该输出与X相加（残差连接）# 然后进行层归一化Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))# 位置前馈网络：对第一个残差连接和层归一化模块的输出进行位置前馈网络操作# 第二个残差连接和层归一化模块：# 将位置前馈网络输出与第一个残差连接和层归一化模块的输出相加并进行归一化return self.addnorm2(Y, self.ffn(Y))# Transformer编码器中的任何层都不会改变其输入的状态
# 创建一个形状为(2, 100, 24)的张量作为输入X
X = torch.ones((2, 100, 24))
# 创建一个形状为(2,)的张量作为有效长度
valid_lens = torch.tensor([3, 2])
encoder_blk = EncoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5)
# 设置模型为评估模式，不进行训练
encoder_blk.eval()
# 输入X和有效长度通过EncoderBlock进行前向传播
print(encoder_blk(X, valid_lens).shape)

---

Transformer编码器

# Transformer编码器
class TransformerEncoder(d2l.Encoder):def __init__(self, vocab_size, key_size, query_size, value_size,num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, num_layers, dropout, use_bias=False, **kwargs):super(TransformerEncoder, self).__init__(**kwargs)# 隐藏层大小self.num_hiddens = num_hiddens# 词嵌入层self.embedding = nn.Embedding(vocab_size, num_hiddens)# 位置编码层self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)# 创建多个编码器块组成的序列# 创建了一个空的顺序容器self.blks，用于存储多个编码器块self.blks = nn.Sequential()# 通过循环迭代的方式，逐个添加编码器块到顺序容器self.blks中for i in range(num_layers):# 使用self.blks.add_module()方法将一个新的编码器块添加到顺序容器中self.blks.add_module("block" + str(i),EncoderBlock(key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, dropout, use_bias))def forward(self, X, valid_lens, *args):# 前向传播过程# 词嵌入层：对输入进行词嵌入操作，将输入转换为高维向量表示# 将嵌入后的向量×隐藏层维度的平方根，目的是缩放嵌入向量，使其位置编码的尺度相匹配# 然后对缩放后的嵌入向量进行编码以便模型能够区分不同位置的输入。"""这种方式一般是用来处理序列数据的"""X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))# 用于存储每个编码器块中的注意力权重self.attention_weights = [None] * len(self.blks)for i, blk in enumerate(self.blks):# 编码器块：将词嵌入结果传入编码器块进行处理X = blk(X, valid_lens)# 将每个编码器块的注意力权重存储到self.attention_weights列表中的对应位置self.attention_weights[i] = blk.attention.attention.attention_weightsreturn X

创建一个两层的Transformer编码器

# 创建一个两层的Transformer编码器，看倒数第二个参数
encoder = TransformerEncoder(200, 24, 24, 24, 24, [100, 24], 24, 48, 8, 2, 0.5)
# 将编码器设置为评估模式
encoder.eval()
# 对输入数据进行前向传播，获取输出的形状[批量大小，序列长度，隐藏层大小]
print(encoder(torch.ones((2, 100), dtype=torch.long), valid_lens).shape)

该部分总代码

import math
import os
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l# 基于位置的前馈网络
class PositionWiseFFN(nn.Module):def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):super(PositionWiseFFN, self).__init__(**kwargs)# 第一个全连接层，将输入维度从ffn_num_input映射到ffn_num_hiddensself.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)# ReLU激活函数，增加非线性特性self.relu = nn.ReLU()# 第二个全连接层，将输入维度从ffn_num_hiddens映射到ffn_num_outputsself.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)def forward(self, X):# 前向传播过程# 通过第一个全连接层，将输入X映射到隐藏层# 应用ReLU激活函数，增加非线性特性# 通过第二个全连接层，将隐藏层映射到输出层return self.dense2(self.relu(self.dense1(X)))# 使用残差连接和层归一化
class AddNorm(nn.Module):def __init__(self, normalized_shape, dropout, **kwargs):super(AddNorm, self).__init__(**kwargs)# 定义一个dropout层，用于随机丢弃部分神经元self.dropout = nn.Dropout(dropout)# 定义一个层归一化层，对输入进行归一化self.ln = nn.LayerNorm(normalized_shape)def forward(self, X, Y):# 前向传播过程# 使用残差连接：将dropout(Y)与X相加# 应用层归一化：对残差进行归一化return self.ln(self.dropout(Y) + X)# 实现编码器中的一个层
class EncoderBlock(nn.Module):def __init__(self, key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,dropout, use_bias=False, **kwargs):super(EncoderBlock, self).__init__(**kwargs)# 多头注意力机制self.attention = d2l.MultiHeadAttention(key_siz