# Transformer：Attention is all your need流程+代码详解

## 1.Attention

### 2.1 Scaled Dot-Product Attention

A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V ∈ R q × d v Attention(Q,K,V)=softmax(\frac{QK^{T}}{\sqrt{d_{k}}})V\in \mathbb{R}^{q\times d_{v}}

class ScaledDotProductAttention(nn.Module):
''' Scaled Dot-Product Attention '''

def __init__(self, temperature, attn_dropout=0.1):
super().__init__()
self.temperature = temperature
self.dropout = nn.Dropout(attn_dropout)

def forward(self, q, k, v, mask=None):
attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

attn = self.dropout(F.softmax(attn, dim=-1))
output = torch.matmul(attn, v)

return output, attn


H e a d i = A t t e n t i o n ( Q W i Q , K W i K , V W i V ) ∈ R q × d v Head_{i}=Attention(QW^{Q}_{i},KW^{K}_{i},VW^{V}_{i})\in \mathbb{R}^{q\times d_{v}}
W i Q ∈ R d m o d e l × d k , W i K ∈ R d m o d e l × d k , W i V ∈ R d m o d e l × d v W^{Q}_{i}\in \mathbb{R}^{d_{model}\times d_{k}},W^{K}_{i}\in \mathbb{R}^{d_{model}\times d_{k}},W^{V}_{i}\in \mathbb{R}^{d_{model}\times d_{v}} 都是线性转换矩阵。在实现中， d m o d e l = d k = d v d_{model}=d_{k}=d_{v} 。将 h h 个head得到的输出进行concat，得到 c o n c a t ( h e a d 1 , . . . , h e a d h ) ∈ R q × h d v concat(head_{1},...,head_{h})\in \mathbb{R}^{q\times hd_{v}} 。采用一个线性转换矩阵 W O ∈ R h d v × d m o d e l W^{O}\in \mathbb{R}^{hd_{v}\times d_{model}} 得到输出：
M u l t i H e a d ( Q , K , V ) = c o n c a t ( h e a d 1 , . . . , h e a d h ) W O ∈ R q × d m o d e l MultiHead(Q,K,V)=concat(head_{1},...,head_{h})W^{O}\in \mathbb{R}^{q\times d_{model}}

class MultiHeadAttention(nn.Module):

def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
super().__init__()

self.d_k = d_k
self.d_v = d_v

self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)

self.dropout = nn.Dropout(dropout)
self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

def forward(self, q, k, v, mask=None):

sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

residual = q

# Pass through the pre-attention projection: b x lq x (n*dv)
# Separate different heads: b x lq x n x dv
q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

# Transpose for attention dot product: b x n x lq x dv
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

# Transpose to move the head dimension back: b x lq x n x dv
# Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
q = self.dropout(self.fc(q))
q += residual

q = self.layer_norm(q)

return q, attn


## 2. Transformer整体架构

### Positional Encoding

P E ( p o s , 2 i ) = s i n ( p o s / 1000 0 2 i / d m o d e l ) PE_{(pos,2i)}=sin(pos/10000^{2i/d_{model}})

P E ( p o s , 2 i + 1 ) = c o s ( p o s / 1000 0 2 i / d m o d e l ) PE_{(pos,2i+1)}=cos(pos/10000^{2i/d_{model}})

s i n ( p o s + k ) = s i n ( p o s ) c o s ( k ) + c o s ( p o s ) s i n ( k ) sin(pos+k)=sin(pos)cos(k)+cos(pos)sin(k)

class PositionalEncoding(nn.Module):

def __init__(self, d_hid, n_position=200):
super(PositionalEncoding, self).__init__()

# Not a parameter
self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))

def _get_sinusoid_encoding_table(self, n_position, d_hid):
''' Sinusoid position encoding table '''
# TODO: make it with torch instead of numpy

def get_position_angle_vec(position):
return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

def forward(self, x):
return x + self.pos_table[:, :x.size(1)].clone().detach()


### Encoder

• 第一个子层是multi-head attention层。在Encoder阶段， Q , K , V Q,K,V
3个矩阵是完全相同的，输出为 R n × d m o d e l \mathbb{R}^{n\times d_{model}}
• 第二个子层是Position-wise Feed-Forward Networks层，实质就是2个FC层中间有一个ReLU函数。输入维度是 R n × d m o d e l \mathbb{R}^{n\times d_{model}} ，中间隐藏维度为 R n × d f f \mathbb{R}^{n\times d_{{ff}}} ，输出维度不变为 R n × d m o d e l \mathbb{R}^{n\times d_{model}}

Position-wise Feed-Forward Networks层代码实现如下：

class PositionwiseFeedForward(nn.Module):
''' A two-feed-forward-layer module '''

def __init__(self, d_in, d_hid, dropout=0.1):
super().__init__()
self.w_1 = nn.Linear(d_in, d_hid) # position-wise
self.w_2 = nn.Linear(d_hid, d_in) # position-wise
self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
self.dropout = nn.Dropout(dropout)

def forward(self, x):

residual = x

x = self.w_2(F.relu(self.w_1(x)))
x = self.dropout(x)
x += residual

x = self.layer_norm(x)

return x


Encoder的整体代码如下：

class Encoder(nn.Module):
''' A encoder model with self attention mechanism. '''

def __init__(
self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
d_model, d_inner, pad_idx, dropout=0.1, n_position=200, scale_emb=False):

super().__init__()

self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
self.dropout = nn.Dropout(p=dropout)
self.layer_stack = nn.ModuleList([
EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
for _ in range(n_layers)])
self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
self.scale_emb = scale_emb
self.d_model = d_model

enc_slf_attn_list = []

# -- Forward
enc_output = self.src_word_emb(src_seq)

if self.scale_emb:
enc_output *= self.d_model ** 0.5
enc_output = self.dropout(self.position_enc(enc_output))
enc_output = self.layer_norm(enc_output)

for enc_layer in self.layer_stack:
enc_slf_attn_list += [enc_slf_attn] if return_attns else []

if return_attns:
return enc_output, enc_slf_attn_list
return enc_output,


### Decoder

Decoder层，包括了 N = 6 N=6 个相同的层。每一个层包括3个子层。

• 第一个子层是multi-head attention层。输入的目标tensor的 Q , K , V ∈ R m × d m o d e l Q,K,V\in \mathbb{R}^{m\times d_{model}} ，这里 m m 表示目标句子的单词个数，而encoder端的输入 n n 表示输入句子的单词个数。3个矩阵是完全相同的，输出为 R m × d m o d e l \mathbb{R}^{m\times d_{model}}
• 第二个子层是multi-head attention层。 Q ∈ R m × d m o d e l Q\in \mathbb{R}^{m\times d_{model}} 是上一层输出的tensor矩阵， K , V ∈ R n × d m o d e l K,V\in \mathbb{R}^{n\times d_{model}} 是相同的，来自于encoder的矩阵，输出为 R m × d m o d e l \mathbb{R}^{m\times d_{model}}
• 第三个子层是Position-wise Feed-Forward Networks层，实质就是2个FC层中间有一个ReLU函数。输入维度是 R m × d m o d e l \mathbb{R}^{m\times d_{model}} ，中间隐藏维度为 R m × d f f \mathbb{R}^{m\times d_{{ff}}} ，输出维度不变为 R m × d m o d e l \mathbb{R}^{m\times d_{model}}
class Decoder(nn.Module):
''' A decoder model with self attention mechanism. '''

def __init__(
self, n_trg_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
d_model, d_inner, pad_idx, n_position=200, dropout=0.1, scale_emb=False):

super().__init__()

self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
self.dropout = nn.Dropout(p=dropout)
self.layer_stack = nn.ModuleList([
DecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
for _ in range(n_layers)])
self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
self.scale_emb = scale_emb
self.d_model = d_model

dec_slf_attn_list, dec_enc_attn_list = [], []

# -- Forward
dec_output = self.trg_word_emb(trg_seq)
if self.scale_emb:
dec_output *= self.d_model ** 0.5
dec_output = self.dropout(self.position_enc(dec_output))
dec_output = self.layer_norm(dec_output)
for dec_layer in self.layer_stack:
dec_output, dec_slf_attn, dec_enc_attn = dec_layer(
dec_slf_attn_list += [dec_slf_attn] if return_attns else []
dec_enc_attn_list += [dec_enc_attn] if return_attns else []

if return_attns:
return dec_output, dec_slf_attn_list, dec_enc_attn_list
return dec_output,


### 整体流程

class Transformer(nn.Module):
''' A sequence to sequence model with attention mechanism. '''

def __init__(
d_word_vec=512, d_model=512, d_inner=2048,
n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1, n_position=200,
trg_emb_prj_weight_sharing=True, emb_src_trg_weight_sharing=True,
scale_emb_or_prj='prj'):

super().__init__()

# In section 3.4 of paper "Attention Is All You Need", there is such detail:
# "In our model, we share the same weight matrix between the two
# embedding layers and the pre-softmax linear transformation...
# In the embedding layers, we multiply those weights by \sqrt{d_model}".
#
# Options here:
#   'emb': multiply \sqrt{d_model} to embedding output
#   'prj': multiply (\sqrt{d_model} ^ -1) to linear projection output
#   'none': no multiplication

assert scale_emb_or_prj in ['emb', 'prj', 'none']
scale_emb = (scale_emb_or_prj == 'emb') if trg_emb_prj_weight_sharing else False
self.scale_prj = (scale_emb_or_prj == 'prj') if trg_emb_prj_weight_sharing else False
self.d_model = d_model

self.encoder = Encoder(
n_src_vocab=n_src_vocab, n_position=n_position,
d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,

self.decoder = Decoder(
n_trg_vocab=n_trg_vocab, n_position=n_position,
d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,

self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)

assert d_model == d_word_vec, \
'To facilitate the residual connections, \
the dimensions of all module outputs shall be the same.'

if trg_emb_prj_weight_sharing:
# Share the weight between target word embedding & last dense layer
self.trg_word_prj.weight = self.decoder.trg_word_emb.weight

if emb_src_trg_weight_sharing:
self.encoder.src_word_emb.weight = self.decoder.trg_word_emb.weight

def forward(self, src_seq, trg_seq):

seq_logit = self.trg_word_prj(dec_output)
if self.scale_prj:
seq_logit *= self.d_model ** -0.5

return seq_logit.view(-1, seq_logit.size(2))


1 1 1 1
1 1 1 1
1 1 1 1
1 1 1 1


1 0 0 0
1 1 0 0
1 1 1 0
1 1 1 1


### 输入句子的预处理

0: UNK_WORD = '<unk>'
2: BOS_WORD = '<s>'
3: EOS_WORD = '</s>'


• Decoder输入：对于目标英文句子 It is a dog. 预处理后的有效信息为[‘BOS’,‘it’，‘is’，‘a’，‘dog’, ‘EOS’], 编码之后为trg=[2,36,253, 433,3]。对于decoder的输入trg[:-1]=[2, 36,253, 433], 对应的ground-truth是trg[:1]=[36,253, 433,3]。因为我们期望当前时刻的信息能够预测下一时刻的单词。例如，利用2得到36，利用[2,36]得到253，… 。

10-22 9346
11-21 9957

05-17 1465
02-03 147
07-18 771
07-10 4万+
09-29
04-07 614