T5是Google提出的Seq2Seq结构的预训练语言模型，一经提出便登上了GLUE、SuperGLUE等各大NLP榜单第一，而它的升级版本mT5因为用了多国语言语料，在中文任务上可以开箱即用。
HuggingFace的Transformers包里的MT5ForConditionalGeneration，就是MT5生成模型的实现，我们只需调用它的model.generate()函数，就能用mT5模型进行推理和生成，简单易用。model.generate()函数背后的逻辑和内部的实现细节又是什么样的呢，本文带你一窥究竟。
Step 1. 关闭梯度
def generate(self)有装饰器@torch.no_grad()，不计算梯度和反向传播，节约计算资源。
@torch.no_grad()
def generate(self, ...)
Step 2. Encoder推理
这里用 get_encoder 函数调用 self.encoder，self.encoder的参数都存储在 named_modules()[‘encoder’] 中。
encoder = self.get_encoder()
ModelOutput = encoder(input_ids, return_dict=True, **encoder_kwargs)
self.encoder是T5Stack的一个实例，这个实例的 forward 函数由字嵌入、数值变化、逐层编码、LN&dropout、存储Dict 共5个部分组成。接下来展开介绍：
Part 1. 字嵌入
这里把input_ids传入 self.embed_tokens 函数，调用Embedding(50000, 512)。Embedding(50000, 512) 的参数存储在 named_modules()[‘embed_tokens’]中。它是torch.nn.functional.embedding的一个实例，以lookup table方式快速取得input_ids对应的inputs_embeds。inputs_embeds维度是(16, 104, 512)，也就是(bc, seq_len, emb_dim)。
inputs_embeds = self.embed_tokens(input_ids)
Part 2. 数值变化
这部分包括key_value、attention mask、head mask、hidden states的变化。
- 初始化past_key_value为[None] * 8，也就是[None] * self.encoder层数
# initialize past_key_values with `None` if past does not exist if past_key_values is None: past_key_values = [None] * len(self.block) - 把attention_mask传入get_extend_attention_mask函数，这个函数把attention_mask中值为1的元素变为0，把值为0的元素变为-10000.0，用于忽略encoder中的pad-masked元素。代码中的[:, None, None, :]用于把attention_mask的维度(bc, seq_len)更改为(bc, 1, 1, seq_len)，例如(16, 1, 1, 104)。
# Provided a padding mask of dimensions [batch_size, seq_length] # - if the model is a decoder, apply a causal mask in addition to the padding mask # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length] if not self.config.is_decoder: extended_attention_mask = attention_mask[:, None, None, :] # Since attention_mask is 1.0 for positions we want to attend and 0.0 for # masked positions, this operation will create a tensor which is 0.0 for # positions we want to attend and -10000.0 for masked positions. # Since we are adding it to the raw scores before the softmax, this is # effectively the same as removing these entirely. extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0 - 调用get_head_mask函数，初始化head_mask和encoder_head_mask为[None] * self.encoder层数
if head_mask is not None: head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) if is_attention_chunked is True: head_mask = head_mask.unsqueeze(-1) else: head_mask = [None] * num_hidden_layers - 把在Part 1. 字嵌入得到的 inputs_embeds 传入self.dropout，self.dropout是torch.nn.functional.dropout的一个实例，以p=0.1的概率丢弃神经元。
hidden_states = F.dropout(inputs_embeds, self.p, self.training, self.inplace)
Part 3. 逐层编码
这部分进入Encoder的各个层，用for循环逐个访问self.encoder的8个层。我们把每层称为一个block。
for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)):
layer_outputs = layer_module(hidden_states,
attention_mask=extended_attention_mask,
use_cache=False,
output_attentions=False,
...)
这8个block的内部实现几乎一样，唯一差别是block_0(从0开始计数)比其他block多了一个relative_position_bias。进入到block内部，由Self-Attention Layer和Feed Forward layer两个模块组成。
# Apply Self-Attention layer
hidden_states = self.layer[0](hidden_states, ...)[0]
# Apply Feed Forward layer
hidden_states = self.layer[-1](hidden_states)
下面展开介绍Self-Attention Layer和Feed Forward layer。
Part 3.1 Self-Attention Layer
Self-Attention Layer 由 layernorm、SelfAttention、残差连接共三步组成。
- 第一步，用self.layernorm调用T5LayerNorm，对传入hidden_states进行层归一化。
注意，mT5的layer_norm和常规layernorm不同，计算时省去了均值。normed_hidden_states = self.layer_norm(hidden_states)def forward(self, hidden_states): variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states - 第二步，用self.SelfAttention调用T5Attention
T5Attention的内部计算和常规的self-attention计算类似，也就是q、k、v分别做线性变换(分别和参数矩阵w_q, w_k, w_v相乘)，得到query_states、key_states、value_states，并且维度统一为(bc, n_head, seq_len, head_dim)，然后query_states和key_states相乘得到scoreattention_output = self.SelfAttention(normed_hidden_states, mask=attention_mask, ...)
不同之处在于，mT5采用相对位置编码，所以会在self-attention计算时，把relative_position_bias和extended_attention_mask一同加在score上。query_states = shape(self.q(hidden_states)) # (batch_size, n_heads, seq_length, dim_per_head) key_states = project(hidden_states, self.k, key_value_states, past_key_value[0] if past_key_value is not None else None) value_states = project(hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None) # compute scores scores = torch.matmul(query_states, key_states.transpose(3, 2))
extended_attention_mask在本文Part 2.数值变化部分讲解过；position_bias是如何计算的呢？这个计算过程是比较特别的。在计算当前token和目标token的attention值时，记录两个token的距离的绝对值，我们不直接使用这个距离值，而是根据距离值的大小进行一定程度的缩小，距离值越大缩小倍数越大，距离值越小缩小倍数越小。具体实现时，采用一种bucket的方法，bucket方法的实现在self._relative_position_bucket函数，代码比较长，这里就不贴了。position_bias = self.compute_bias(real_seq_length, key_length) if extended_attention_mask is not None: position_bias = position_bias + extended_attention_mask # (batch_size, n_heads, seq_length, key_length) scores += position_bias
接下来，和常规self-attention一样，经过softmax得到归一化后的score，再dropout一下，再和value_states相乘，然后维度恢复为(bc, seq_len, n_head*head_dim)，例如(16, 104, 384)，最后做线性变换(和w_o相乘)。此外有一个细节，计算attention score时不除以dk。def compute_bias(self, query_length, key_length): """ Compute binned relative position bias """ context_position = torch.arange(query_length, dtype=torch.long)[:, None] memory_position = torch.arange(key_length, dtype=torch.long)[None, :] relative_position = memory_position - context_position # shape (query_length, key_length) relative_position_bucket = self._relative_position_bucket( relative_position, # shape (query_length, key_length) bidirectional=(not self.is_decoder), num_buckets=self.relative_attention_num_buckets, ) relative_position_bucket = relative_position_bucket.to(self.relative_attention_bias.weight.device) values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) return valuesattn_weights = F.softmax(scores.float(), dim=-1).type_as(scores) # (batch_size, n_heads, seq_length, key_length) attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training) # (batch_size, n_heads, seq_length, key_length) attn_output = unshape(torch.matmul(attn_weights, value_states)) # (batch_size, seq_length, dim) attn_output = self.o(attn_output) - 第三步，残差连接，把self-attention的输出传入self.dropout，再把dropout的输出与hidden_states直接相加。
hidden_states = hidden_states + self.dropout(attention_output[0])
Part 3.2 Feed Forward Layer
Feed Forward Layer调用T5LayerFF，它由layernorm、DenseReluDense、残差连接共三步组成。
- 第一步，layernorm，和Part 3.1一样，这里不再赘述。
forwarded_states = self.layer_norm(hidden_states) - 第二步，DenseReluDense，它是T5DenseGatedGeluDense的一个实例，它的forward函数由以下3部分组成：
- gelu_act(wi_0( ~ )) * wi_1( ~ )，"~"表示输入，wi_0和wi_1是两个线性变换层，gelu_act的计算公式是 g e l u _ a c t ( x ) = 1 2 ⋅ x ⋅ ( 1 + tanh ( 2 π ⋅ x + 0.44715 ∗ x 3 ) ) gelu\_act(x)=\frac{1}{2}·x·(1+\tanh{(\sqrt{\frac{2}{\pi}}·x+0.44715*x^3)}) gelu_act(x)=21⋅x⋅(1+tanh(π2⋅x+0.44715∗x3))
- dropout层，以p=0.1丢弃神经元，输出维度(16, 104, 1024)，也就是(bc, seq_len, fc_dim)
- wo线性变换层，输出维度(16, 104, 512)，也就是(bc, seq_len, model_dim)
def gelu_act(x): return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) hidden_gelu = self.gelu_act(self.wi_0(hidden_states)) hidden_linear = self.wi_1(hidden_states) hidden_states = hidden_gelu * hidden_linear hidden_states = self.dropout(hidden_states) hidden_states = self.wo(hidden_states) - 第三步，残差连接，和Part 3.1一样，这里不再赘述。
hidden_states = hidden_states + self.dropout(forwarded_states)
Part 4. LN&dropout
这部分由layernorm和dropout两部分组成。layernorm和dropout细节在前文介绍过，这里不再重复。
hidden_states = self.final_layer_norm(hidden_states)
hidden_states = self.dropout(hidden_states)
Part 5. 存储Dict
这里存储的是BaseModelOutputWithPastAndCrossAttentions的一个实例，是一个OrderedDict，包含encoder的输出。
return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states,
past_key_values=present_key_value_states,
hidden_states=all_hidden_states,
attentions=all_attentions,
cross_attentions=all_cross_attentions,)
Step 3. Decoder推理
Decoder的input_ids是什么呢？分两种情况：
- 默认情况下，Decoder的input_ids是[101]*batch_size，101是’[bos]'的token_id；
- 喂给Decoder一个prompt，例如“上文提到了”，Decoder会接着这个输入进行生成。Decoder的输入是就是这5个字的token_id；
# set input_ids as decoder_input_ids if "decoder_input_ids" in model_kwargs: input_ids = model_kwargs.pop("decoder_input_ids") else: input_ids = self._prepare_decoder_input_ids_for_generation(input_ids, decoder_start_token_id=decoder_start_token_id, bos_token_id=bos_token_id)
除了准备decoder的input_ids，还需要准备一个logits_processor，这里采用MinLengthLogitsProcessor(min_len, eos_id)，其中min_len是0，eos_id是102，控制生成序列长度不得小于min_len。
# enforcing a min-length by setting EOS probability to 0.
class MinLengthLogitsProcessor(LogitsProcessor):
def __call__(self, input_ids, scores):
cur_len = input_ids.shape[-1]
if cur_len < self.min_length:
scores[:, self.eos_token_id] = -float("inf")
return scores
准备完毕后，进入生成过程，本文采用greedy_search，greedy_search每个时间步只选择条件概率值最高的生成序列。self.greedy_search内部用while循环，一个字一个字地生成，直到生成的字是’<\s>'或者达到预设定的max_length值，例如40，便会停止生成。所以max_length值越大，推理耗时往往越长。while循环的内部，分为两部分：
第一部分，调用BaseModelOutput，获取Encoder的输出，前文讲到encoder的输出存储在Dict中，所以读字典取值即可。
if encoder_outputs is not None and return_dict and not isinstance(encoder_outputs, BaseModelOutput):
encoder_outputs = BaseModelOutput(last_hidden_state=encoder_outputs[0],
hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,)
第二部分，正式进入Decoder解码阶段，调用 self.decoder，self.decoder的参数存储在 named_modules()[‘decoder’] 中。self.decoder是T5Stack的一个实例，这个实例的 forward 函数由字嵌入、数值变化、逐层解码、LN&dropout、存储Dict、获取生成字 共6个部分组成，和self.encoder很相似，但在部分细节处有所不同。下面展开来看：
Part 1. 字嵌入
这里把decoder_input_ids传入 self.embed_tokens 函数，调用Embedding(50000, 512)。Embedding(50000, 512) 的参数存储在 named_modules()[‘embed_tokens’]中。它是torch.nn.functional.embedding的一个实例，以lookup table方式快速取得decoder_input_ids对应的inputs_embeds。初始状态下，inputs_embeds维度是(16, 1, 512)，也就是(bc, decoder_seq_len, emb_dim)。
inputs_embeds = self.embed_tokens(input_ids)
此外与encoder不同的是，config.use_cache设置为True，这个下文会用到
Part 2. 数值变化
这部分包括key_value、attention mask、head mask、hidden states的变化。
- 其中key_value、head mask、hidden states这三个部分和encoder完全相同（详见Step 2. Encoder推理 - Part 2. 数值变化），此处不再重复介绍
- attention mask的计算
- 与encoder相同之处：这里把decoder_attention_mask传入get_extend_attention_mask函数，这个函数把decoder_attention_mask中值为1的元素变为0，把值为0的元素变为-10000.0，用于decoder中的pad元素的mask；
- 与encoder不同之处：
- 这里额外计算causal_mask，并为decoder_attention_mask乘上causal_mask，用于decoder中的future元素的mask。
if self.config.is_decoder: batch_size, seq_length = input_shape seq_ids = torch.arange(seq_length, device=device) causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None] extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :] - 还调用self.invert_attention_mask函数，计算encoder_extended_attention_mask，把encoder_attention_mask中值为1的元素变为0，把值为0的元素变为-1e9。
encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :] encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -1e9
- 这里额外计算causal_mask，并为decoder_attention_mask乘上causal_mask，用于decoder中的future元素的mask。
Part 3. 逐层解码
这部分进入Decoder的各个层，用for循环逐个访问self.decoder的8个层，我们把每层称为一个block，这8个block的内部完全一样（都没有relative_position_bias）。
for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)):
layer_outputs = layer_module(hidden_states,
attention_mask=extended_attention_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_extended_attention_mask,
use_cache=True,
output_attentions=False,
...)
以decoder的block_0为例，进入到block内部，由Self-Attention Layer、Cross_attention Layer和Feed Forward layer 共三个模块组成。
# Apply Self-Attention layer
hidden_states = self.layer[0](hidden_states, ...)[0]
# Apply Cross-Attention layer
cross_attention_outputs = self.layer[1](hidden_states, ...)[0]
# Apply Feed Forward layer
hidden_states = self.layer[-1](hidden_states)
下面展开介绍Self-Attention Layer、Cross_attention Layer和Feed Forward layer。
Part 3.1 Self-Attention Layer
这部分和encoder完全一致（详见Step 2. Encoder推理 - Part 3.1 Self-Attention Layer），不再重复介绍。
Part 3.2 Cross-Attention Layer
cross-attention与self-attention不同在于，qkv不是相同的序列，q是decoder当前时刻的self-Attention Layer的输出，k和v是encoder的输出。
这里是调用T5LayerCrossAttention的forward函数，输入是decoder_input_emb, encoder_attention_mask, 以及key_value_states，这里的key_value_states就是encoder的输出，计算过程分为3部分：
- 第一步，self.layernorm进行层归一化
normed_hidden_states = self.layer_norm(hidden_states) - 第二步，self.EncDecAttention，它是T5Attention的一个实例，与encoder的SelfAttention计算几乎完全相同，不再重复罗列计算过程，不同之处在于：
- query_states对应于decoder的输入，维度是(bc, n_head, decoder_input_len, head_dim)，例如(16, 6, 1, 64)；key_states和value_states对应于encoder的输出，维度是(bc, n_head, encoder_input_len, head_dim)，例如(16, 6, 104, 64)
- 由于decoder没有relative_position_bias，这里的positon_bias元素值全为零
- 由于decoder的use_cache为True，这里会存储present_key_value_states，也就是(key_states和value_states)
- 第三步，残差连接。
layer_output = hidden_states + self.dropout(attention_output[0])
Part 3.3 Feed Forward Layer
这部分和encoder完全一致（详见Step 2. Encoder推理 - Part 3.2 Feed Forward Layer），不再重复介绍。
Part 4. LN&dropout
这部分和encoder完全一致（详见Step 2. Encoder推理 - Part 4. LN&dropout），不再重复介绍。
Part 5. 存储Dict
这部分和encoder完全一致（详见Step 2. Encoder推理 - Part 5. 存储Dict），不再重复介绍。
Part 6. 获取生成字
第一步，把decoder的输出送入self.lm_head，它是nn.modules.Linear的一个实例，用于得到生成字的logits。logits的维度是(16, 1, 50000)，也就是(bc, 1, word_dict_size)。
lm_logits = self.lm_head(sequence_output)
第二步，用[:, -1, :]取当前生成字的logits，然后调用logits_processor预处理这个logits，logits_processor就是前文提到的MinLengthLogitsProcessor，控制生成序列长度不得小于min_len。输出维度为(16, 50000)，也就是(bc, word_dict_size)
next_token_logits = outputs.logits[:, -1, :]
next_tokens_scores = logits_processor(input_ids, next_token_logits)
第三步，torch.argmax，得到当前概率值最大的token的id，维度为(16,)，也就是(bc)，
next_tokens = torch.argmax(next_tokens_scores, dim=-1)
第四步，把这个token_id拼接到input_ids上，用于下一个时刻的while循环。如果这个token_id是’<\s>'或者while循环结束，就会停止继续生成下一个字。
至此，greedy search就结束了，这里的input_ids就是最终的输出。
2662
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
