Transformer代码详细解读
文章目录
- Transformer代码详细解读
- 简介
- 1. 数据准备
- 1.1 词表构建
- 1.2 数据构建
- 2. 模型整体架构
- 2.1 超参数设置
- 2.2 整体架构
- 2.2 模型训练
- 3. 编码器(Encoder)
- 3.1 编码器
- 3.2 单个编码层
- 3.3 Padding Mask
- 4. 解码器(Decoder)
- 4.1 解码器
- 4.2 单个解码层
- 4.3 Sequence Mask
- 5. 位置编码
- 6. 多头注意力机制(Muti-Head Attention)
- 6.1 多头注意力机制
- 6.2 点积缩放的注意力机制(ScaledDotProductAttention)
- 7. 前馈神经网络(Poswise-FeedForward)
- 7.1 实现方式1:Conv1d
- 7.2 实现方式2:Linear
简介
Transformer是当前自然语言处理乃至整个深度学习领域极其重要的经典模型,是当前大规模预训练语言模型如BERT、GPT、BART等的基本架构,读懂该模型的源码,对于深入了解其原理、后续微调、改进以该模型为基本架构的模型很有帮助,解读该代码重点在于:
- 搞清楚Transformer的模型框架、 各模块细节。可以参考另一篇博客【学习笔记】Transformer模型解读
- 从整体到局部,搞清楚数据输入输出流动形状(shape)。
另:
- 源码参考自链接1,可以通过该链接获取;
- 解读参考自B站up主链接2,对该博主解读的注释中的一些小错误进行修正,以及对一些新手难以理解的代码细节进行详细解读;
- 这篇博客的代码也可以参考链接3。
- 实现框架:Pytorch。
1. 数据准备
本文以一个简单的德语到英语的机器翻译任务Demo为例。
1.1 词表构建
词嵌入本身是一个look up查表的过程,因此需要构建词表:token及其索引。现在的实际任务中,一般使用Huggingface Transformers库的Tokenizer等API直接获取。
src_vocab = {'P': 0, 'ich': 1, 'mochte': 2, 'ein': 3, 'bier': 4}src_vocab_size = len(src_vocab)tgt_vocab = {'P': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, 'S': 5, 'E': 6}tgt_vocab_size = len(tgt_vocab)
1.2 数据构建
在实际任务中,应该从数据集中读取,然后构建DataLoader,本文为便于解读,仅实现一个toy过程。
‘S’(Start)表示开始字符,‘E’(End)表示结束字符,‘P’(Pad)表示填充字符。
输入文本是字符串类型的,需要将其转换成字符在词汇表中的索引,然后转换成Tensor类型。
sentences = ['ich mochte ein bier P', 'S i want a beer', 'i want a beer E']def make_batch(sentences):# 把文本转成词表索引input_batch = [[src_vocab[n] for n in sentences[0].split()]]output_batch = [[tgt_vocab[n] for n in sentences[1].split()]]target_batch = [[tgt_vocab[n] for n in sentences[2].split()]]# 把索引转成tensor类型return torch.LongTensor(input_batch), torch.LongTensor(output_batch), torch.LongTensor(target_batch)
2. 模型整体架构
2.1 超参数设置
部分重要的模型超参数设置,包括:
- 输入、输出的句子长度;
- 模型词嵌入大小;
- 前馈神经网络(FeedForward)层隐藏层维度;
- K(在自注意力中,也是Q的大小)、V向量大小;
- 编码器和解码器的层数;
- 多头自注意力的头数。
src_len = 5 # length of sourcetgt_len = 5 # length of target## 模型参数d_model = 512# Embedding Sized_ff = 2048# FeedForward dimensiond_k = d_v = 64# dimension of K(=Q), Vn_layers = 6# number of Encoder of Decoder Layern_heads = 8# number of heads in Multi-Head Attention
2.2 整体架构
Transformer的整个网络结构包括三部分:编码层、解码层、输出层。
流程
- 输入文本进行词嵌入和位置编码,作为最终的文本嵌入;
- 文本嵌入经过Encoder编码,得到注意力加权后输出的编码向量以及自注意力权重矩阵;
- 然后将编码向量和样本的Ground trurh共同输入解码器,经过注意力加权等操作后输出最终的上下文向量,然后映射到词表大小的线性层上进行解码生成文本;
- 最终返回代表预测结果的logits矩阵。
数据形状
enc_inputs:[batch_size,src_len]
dec_inputs:[batch_size,tgt_len]
enc_outputs:[batch_size,src_len,d_model]
enc_self_attns:[batch_size,n_heads,src_len,src_len]
dec_outputs:[batch_size,tgt_len,d_model]
dec_self_attns:[batch_size,n_heads,tgt_len,tgt_len]
dec_enc_attns:[batch_size,n_heads,tgt_len,src_len]
dec_logits:[batch_size,tgt_len,tgt_vocab_size]
class Transformer(nn.Module):def __init__(self):super(Transformer, self).__init__()# 编码器self.encoder = Encoder()# 解码器self.decoder = Decoder()# 输出层,d_model是解码层每个token输出的维度大小,之后会做一个tgt_vocab_size大小的softmax# 因为解码输出过程相当于是在一个词表大小级别上的分类self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False)def forward(self, enc_inputs, dec_inputs):# 输入输出部分具体要输入和返回什么参数,可以根据自己的任务和改进需要进行自定义修改,内部的执行过程是不变的# enc_outputs就是编码器的输出,enc_self_attns是QK转置相乘之后softmax之后的注意力矩阵,代表的是每个单词和其他单词相关性;# 由于多头注意力机制会分头计算注意力,所以注意力权重矩阵是个四维向量,# 即[batch_size,n_heads,src_len,src_len]enc_outputs, enc_self_attns = self.encoder(enc_inputs)# dec_self_attns类比于enc_self_attns,是查看每个单词对decoder中输入的其余单词的相关性;# dec_enc_attns是decoder中每个单词对encoder中每个单词的相关性,# 即Cross_Attention输出的注意力权重矩阵,形状为[batch_size,n_heads,tgt_len,src_len]# 注意,这里的xxx_attns都是list类型,因为在Encoder、Decoder中将每一层的注意力权重都保存下来添加到列表中返回了dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs)# dec_outputs映射到词表大小dec_logits = self.projection(dec_outputs) # dec_logits : [batch_size,tgt_len, tgt_vocab_size]# 这里dec_logits进行view操作主要是为了适应后面的CrossEntropyLoss API的参数要求return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns
2.2 模型训练
model = Transformer()criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=0.001)enc_inputs, dec_inputs, target_batch = make_batch(sentences)for epoch in range(10):optimizer.zero_grad()outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)# output:[batch_size x tgt_len,tgt_vocab_size]# 就这份代码而言,这里其实可以不写.contiguous(),因为target_batch这个tensor是连续的loss = criterion(outputs, target_batch.contiguous().view(-1))print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))loss.backward()optimizer.step()
3. 编码器(Encoder)
3.1 编码器
编码器由N个编码层堆叠而成。
流程
输入文本的索引tensor,经过词嵌入层得到词嵌入,然后和位置编码线性相加作为输入层的最终输出;
随后,每一层的输出最为下一层编码块的输入,在每个编码块里进行注意力计算、前馈神经网络、残差连接、层归一化等操作;
最终返回编码器最后一层的输出和每一层的注意力权重矩阵。
class Encoder(nn.Module):def __init__(self):super(Encoder, self).__init__()# 这个其实就是去定义生成一个词嵌入矩阵,大小是 src_vocab_size * d_modelself.src_emb = nn.Embedding(src_vocab_size, d_model)# 位置编码,这里是固定的正余弦函数,也可以使用类似词向量的nn.Embedding获得一个可以更新学习的位置编码self.pos_emb = PositionalEncoding(d_model) # 使用ModuleList对多个encoder进行堆叠,因为后续的encoder并没有使用词向量和位置编码,所以抽离出来;self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)]) def forward(self, enc_inputs):# enc_inputs形状是:[batch_size,src_len]# 下面这个代码通过src_emb,进行索引定位,enc_outputs输出形状是[batch_size, src_len, d_model]enc_outputs = self.src_emb(enc_inputs)# 位置编码和词嵌入相加,具体实现在PositionalEncoding里,enc_outputs:[batch_size,src_len,d_model]enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1)# get_attn_pad_mask是为了得到句子中pad的位置信息,以便在计算注意力时忽略pad符号enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)enc_self_attns = []for layer in self.layers:# 每一层的输出作为下一层的输入,enc_outputs:[batch_size,src_len,d_model]enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask)# 把每一层得到的注意力权重矩阵添加到列表里最后返回,enc_self_attn:[batch_size,src_len,src_len]enc_self_attns.append(enc_self_attn)return enc_outputs, enc_self_attns
3.2 单个编码层
class EncoderLayer(nn.Module):def __init__(self):super(EncoderLayer, self).__init__()self.enc_self_attn = MultiHeadAttention()self.pos_ffn = PoswiseFeedForwardNet()def forward(self, enc_inputs, enc_self_attn_mask):#enc_inputs形状是[batch_size x seq_len_q x d_model],注意,最初始的QKV矩阵是等同于这个输入的enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) # enc_inputs to same Q,K,Venc_outputs = self.pos_ffn(enc_outputs) # enc_outputs: [batch_size x len_q x d_model]return enc_outputs, attn
3.3 Padding Mask
在后面注意力机制的部分,在计算出 Q ∗ KT Q*K^TQ∗KT除以根号之后、softmax之前得到的矩阵大小为[len_input * len_input],代表每个单词对所有(包含自己)单词的影响力。这个函数用于获取一个同等大小形状的矩阵,标记哪个位置是PAD符号,之后在计算softmax之前会把这些地方置为无穷小,避免Query去关注这些无意义的PAD符号。
注意,本函数得到的矩阵形状是[batch_size x len_q x len_k],是对K中的pad符号进行标识,并没有对Q中的做标识,因为没必要。
seq_q和seq_k不一定一致,例如,在交互注意力,q来自解码端,k来自编码端,所以告诉模型编码这边pad符号信息就可以,解码端的pad信息在交互注意力层是没有用到的。
def get_attn_pad_mask(seq_q, seq_k):batch_size, len_q = seq_q.size()batch_size, len_k = seq_k.size()# eq(zero) is PAD tokenpad_attn_mask = seq_k.data.eq(0).unsqueeze(1)# batch_size x 1 x len_k, one is masking# 最终得到的应该是一个最后n列为1的矩阵,即K的最后n个token为PAD。return pad_attn_mask.expand(batch_size, len_q, len_k)# batch_size x len_q x len_k
4. 解码器(Decoder)
4.1 解码器
解码器由N个解码层堆叠而成。
解码器和编码器类似,区别在于每一层解码器有两个多头注意力机制。第一个多头注意力机制中,需要对未来词进行掩码屏蔽;第二个多头注意力机制是交叉注意力机制,需要将Encoder的输出作为K、V,将Decoder上一部分的输出作为Q。其意义就在于解码时能够利用到编码得到的上下文信息。
class DecoderLayer(nn.Module):def __init__(self):super(DecoderLayer, self).__init__()self.dec_self_attn = MultiHeadAttention()self.dec_enc_attn = MultiHeadAttention()self.pos_ffn = PoswiseFeedForwardNet()def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)# 这里用dec_outputs作为Q,enc_outputs作为K和V,实现交叉注意力机制dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask)dec_outputs = self.pos_ffn(dec_outputs)return dec_outputs, dec_self_attn, dec_enc_attn
4.2 单个解码层
class Decoder(nn.Module):def __init__(self):super(Decoder, self).__init__()self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)self.pos_emb = PositionalEncoding(d_model)self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])# dec_inputs : [batch_size x target_len]def forward(self, dec_inputs, enc_inputs, enc_outputs):dec_outputs = self.tgt_emb(dec_inputs)# [batch_size, tgt_len, d_model]dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1) # [batch_size, tgt_len, d_model]# 获取自注意力层pad的mask矩阵,1表示被maskdec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs)# 获取look ahead的mask矩阵,即让注意力机制看不到未来的单词,获得到一个上三角为1的矩阵,1表示被maskdec_self_attn_subsequent_mask = get_attn_subsequent_mask(dec_inputs)# 两个mask矩阵相加,大于0的为1,不大于0的为0,既屏蔽了pad的信息,也屏蔽了未来时刻的信息,为1的在之后就会被fill到无限小# 使用gt()函数,因为可能会有在两个mask都被屏蔽的情况,1+1=2dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequent_mask), 0)# 获取交互注意力机制中的mask矩阵,decoder的输入是q,encoder的输入是k,需要知道k里面哪些是pad符号,# 注意,q肯定也是有pad符号,但是没有必要将其屏蔽# 这里不用再把q的未来词再mask了,因为前面已经mask过一次,输出向量对应部分的值应该都是0了dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)dec_self_attns, dec_enc_attns = [], []for layer in self.layers:dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask)dec_self_attns.append(dec_self_attn)dec_enc_attns.append(dec_enc_attn)return dec_outputs, dec_self_attns, dec_enc_attns
4.3 Sequence Mask
遮蔽未来词,让当前词看不到未来词。这个函数就是用来表示Decoder的输入中哪些是未来词,显然,这个Mask矩阵应该是一个上三角矩阵。
def get_attn_subsequent_mask(seq):"""seq: [batch_size, tgt_len]"""attn_shape = [seq.size(0), seq.size(1), seq.size(1)]# attn_shape: [batch_size, tgt_len, tgt_len]# np.triu()返回一个上三角矩阵,自对角线k以下元素全部置为0,k=0为主对角线subsequence_mask = np.triu(np.ones(attn_shape), k=1)# 生成一个上三角矩阵# 如果没转成byte,这里默认是Double(float64),占据的内存空间大,浪费,用byte就够了subsequence_mask = torch.from_numpy(subsequence_mask).byte()return subsequence_mask# [batch_size, tgt_len, tgt_len]
5. 位置编码
e−2id m o d e l ∗log10000 = 1 1000 02 i d model e^{\frac{-2i}{d_{model}}*log10000}=\frac{1}{10000^{\frac{2i}{d_{model}}}} edmodel−2i∗log10000=10000dmodel2i1
位置编码的实现直接对照着公式写就行,下面这个代码只是其中一种实现方式;
需要注意的是偶数和奇数在公式上有一个共同部分,这里使用指数函数e和log函数(以e为底)把次方拿下来,方便计算;
pos代表的是单词在句子中的绝对索引位置,例如max_len是128,那么索引就是从0,1,2,…,127,假设d_model是512,即用一个512维tensor来编码一个索引位置,那么0<=2i<512,则0<=i<=255,那么2i对应取值就是0,2,4…510,即偶数位置;2i+1的取值是1,3,5…511,即奇数位置。
最后的文本嵌入表征是词嵌入和位置编码相加得到。
class PositionalEncoding(nn.Module):def __init__(self, d_model, dropout=0.1, max_len=5000):super(PositionalEncoding, self).__init__()self.dropout = nn.Dropout(p=dropout)# 生成一个形状为[max_len,d_model]的全为0的tensorpe = torch.zeros(max_len, d_model)# position:[max_len,1],即[5000,1],这里插入一个维度是为了后面能够进行广播机制然后和div_term直接相乘# 注意,要理解一下这里position的维度。每个pos都需要512个编码。position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)# 共有项,利用指数函数e和对数函数log取下来,方便计算div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))# 这里position * div_term有广播机制,因为div_term的形状为[d_model/2],即[256],符合广播条件,广播后两个tensor经过复制,形状都会变成[5000,256],*表示两个tensor对应位置处的两个元素相乘# 这里需要注意的是pe[:, 0::2]这个用法,就是从0开始到最后面,补长为2,其实代表的就是偶数位置赋值给pepe[:, 0::2] = torch.sin(position * div_term)# 同理,这里是奇数位置pe[:, 1::2] = torch.cos(position * div_term)# 上面代码获取之后得到的pe:[max_len*d_model]# 下面这个代码之后,我们得到的pe形状是:[max_len*1*d_model]pe = pe.unsqueeze(0).transpose(0, 1)# 定一个缓冲区,其实简单理解为这个参数不更新就可以,但是参数仍然作为模型的参数保存self.register_buffer('pe', pe)def forward(self, x):"""x: [seq_len, batch_size, d_model]"""# 这里的self.pe是从缓冲区里拿的# 切片操作,把pe第一维的前seq_len个tensor和x相加,其他维度不变# 这里其实也有广播机制,pe:[max_len,1,d_model],第二维大小为1,会自动扩张到batch_size大小。# 实现词嵌入和位置编码的线性相加x = x + self.pe[:x.size(0), :]return self.dropout(x)
6. 多头注意力机制(Muti-Head Attention)
6.1 多头注意力机制
这里实际代码可能跟一些原理的讲解略有出入。
先映射,后分头。即拿到输入以后,映射到d_k * n_heads维度上,然后再通过转置分成n_heads个头,这样就不用写n_heads次参数矩阵,也不用进行拼接操作。
多头注意力机制完成之后还要经过一个线性层。
class MultiHeadAttention(nn.Module):def __init__(self):super(MultiHeadAttention, self).__init__()# Wq,Wk,Wv其实就是一个线性层,用来将输入映射为K、Q、V# 这里输出是d_k * n_heads,因为是先映射,后分头。self.W_Q = nn.Linear(d_model, d_k * n_heads)self.W_K = nn.Linear(d_model, d_k * n_heads)self.W_V = nn.Linear(d_model, d_v * n_heads)self.linear = nn.Linear(n_heads * d_v, d_model)self.layer_norm = nn.LayerNorm(d_model)def forward(self, Q, K, V, attn_mask):# attn_mask:[batch_size,len_q,len_k]# 输入的数据形状: Q: [batch_size x len_q x d_model], K: [batch_size x len_k x d_model], # V: [batch_size x len_k x d_model]residual, batch_size = Q, Q.size(0)# (B, S, D) -proj-> (B, S, D) -split-> (B, S, H, W) -trans-> (B, H, S, W)# 分头;一定要注意的是q和k分头之后维度是一致的,所以一看这里都是d_k# q_s: [batch_size x n_heads x len_q x d_k]q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)# k_s: [batch_size x n_heads x len_k x d_k]k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2)# v_s: [batch_size x n_heads x len_k x d_v]v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2)# attn_mask:[batch_size x len_q x len_k] --> [batch_size x n_heads x len_q x len_k]# 就是把pad信息复制n份,重复到n个头上以便计算多头注意力机制attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)# 计算ScaledDotProductAttention# 得到的结果有两个:context: [batch_size x n_heads x len_q x d_v],# attn: [batch_size x n_heads x len_q x len_k]context, attn = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)# 这里实际上在拼接n个头,把n个头的加权注意力输出拼接,然后过一个线性层,context变成# [batch_size,len_q,n_heads*d_v]。这里context需要进行contiguous,因为transpose后源tensor变成不连续的# 了,view操作需要连续的tensor。context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v)output = self.linear(context)# 过残差、LN,输出output: [batch_size x len_q x d_model]和这一层的加权注意力表征向量return self.layer_norm(output + residual), attn
6.2 点积缩放的注意力机制(ScaledDotProductAttention)
这里有个关键细节就是对Mask矩阵和注意力权重矩阵的联合操作,要把注意力权重矩阵的PAD部分变成无限小,以对Query起到屏蔽作用。
class ScaledDotProductAttention(nn.Module):def __init__(self):super(ScaledDotProductAttention, self).__init__()def forward(self, Q, K, V, attn_mask):# 输入进来的维度分别是Q:[batch_size x n_heads x len_q x d_k]K:[batch_size x n_heads x len_k x d_k]V:[batch_size x n_heads x len_k x d_v]# matmul操作即矩阵相乘# [batch_size x n_heads x len_q x d_k] matmul [batch_size x n_heads x d_k x len_k] -> [batch_size x n_heads x len_q x len_k]scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)# masked_fill_(mask,value)这个函数,用value填充源向量中与mask中值为1位置相对应的元素,# 要求mask和要填充的源向量形状需一致# 把被mask的地方置为无穷小,softmax之后会趋近于0,Q会忽视这部分的权重scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is one.attn = nn.Softmax(dim=-1)(scores)context = torch.matmul(attn, V)# context:[batch_size,n_heads,len_q,d_k]# attn:[batch_size,n_heads,len_q,len_k]return context, attn
7. 前馈神经网络(Poswise-FeedForward)
FeedForward实际上就是一个两层的线性层,对输入进行线性转换。Position-wise意为对每个点独立做,即对序列中的每个token独立过同一个MLP,即作用在输入的最后一个维度上。
MLP有两种实现方式,一种是通过卷积的方式实现,一种是通过线性层实现。二者的区别除了原理上,还有代码细节上,例如,Conv1d要求输入必须是[batch_size,channel,length],必须是三维tensor,而Linear要求输入是[batch_size,*,d_model],可以有多个维度。
7.1 实现方式1:Conv1d
class PoswiseFeedForwardNet(nn.Module):def __init__(self):super(PoswiseFeedForwardNet, self).__init__()self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)self.layer_norm = nn.LayerNorm(d_model)def forward(self, inputs):residual = inputs # inputs : [batch_size, len_q, d_model]# Conv1d的输入为[batch, channel, length],作用于第二个维度channel,所以这里要转置output = nn.ReLU()(self.conv1(inputs.transpose(1, 2)))output = self.conv2(output).transpose(1, 2)return self.layer_norm(output + residual)
7.2 实现方式2:Linear
class PoswiseFeedForwardNet(nn.Module):def __init__(self):super(PoswiseFeedForwardNet, self).__init__()self.fc = nn.Sequential(nn.Linear(d_model, d_ff, bias=False),nn.ReLU(),nn.Linear(d_ff, d_model, bias=False))def forward(self, inputs):# inputs: [batch_size, seq_len, d_model]residual = inputsoutput = self.fc(inputs)return nn.LayerNorm(d_model).(output + residual) # [batch_size, seq_len, d_model]