seq2seq编码器encoder和解码器decoder详解

编码器

在序列到序列模型中，编码器将输入序列（如一个句子）转换为一个隐藏状态序列，供解码器生成输出。编码层通常由嵌入层和RNN（如GRU/LSTM)等组成

• Token:是模型处理文本时的基本单元，可以是词,子词,字符等，每个token都有一个对应的ID。是由原始文本中的词或子词通过分词器（Tokenizer）处理后得到的最小单位，这些 token 会被映射为词汇表中的唯一索引 ID
• 输入：
  • 原始输入序列：通常是一个句子的词汇ID序列。例如 [“Hello”, “world”] 可能会被映射为 [1, 2]，假设“Hello”的ID是1，“world”的ID是2；
  • 嵌入向量： token ID 列表 [1, 2] 会作为模型的输入，每一个token即1和2经过嵌入层(embedding layer)都会转换为’encoder_embed_dim’大小的向量，即两个嵌入向量。嵌入向量最初是由模型在训练过程中学到的，初始时通常是随机的。在训练的过程中，嵌入向量会调整，使得语义相似的词在向量空间中更接近。
• 编码层(通常是RNN,GRU,LSTM,Transformer)：
  • 作用: 处理输入序列的时间依赖性，并生成隐藏状态
  • 输入：嵌入向量序列，形状为(seq_len, embed_dim)
  • 输出: 编码器对每个时间步（即每个token）计算一个隐藏状态，这些隐藏状态组成了一个隐藏状态序列，捕捉了当前 token 以及它的上下文信息，每个时间步的隐藏状态不仅考虑当前词的嵌入向量，还结合了之前所有时间步的信息。对于长度为seq_len的输入序列，隐藏状态序列的形状为 (seq_len, hidden_dim)。对于 LSTM 或 GRU，编码器还会输出最后一个时间步的隐藏状态，供解码器初始化使用
• 输出：
  • 隐藏状态序列：编码器处理整个输入序列后，输出的隐藏状态序列通常被称为 encoder_outputs。其中每个隐藏状态序列对应于输入序列中的一个 token，这个序列的形状是 (seq_len, batch_size, hidden_dim)
  • 最终隐藏状态：编码器的最后一个时间步的隐藏状态通常被用作解码器的初始状态。这被称为 encoder_hiddens，形状为 (num_layers, batch_size, hidden_dim)，在双向RNN中，这个向量可能会有两倍的维度，变为(num_layers, batch_size, 2*hidden_dim)

'''定义了一个用于自然语言处理的编码器类 RNNEncoder，该类继承自 FairseqEncoder，
并实现了一个双向 GRU（门控循环单元，Gated Recurrent Unit）来对输入的文本进行编码'''
class RNNEncoder(FairseqEncoder):
    def __init__(self, args, dictionary, embed_tokens):
        super().__init__(dictionary)
        self.embed_tokens = embed_tokens   # 嵌入层，用于将 token 索引转换为嵌入向量。
        
        self.embed_dim = args.encoder_embed_dim      # 嵌入维度
        self.hidden_dim = args.encoder_ffn_embed_dim # 隐藏层维度
        self.num_layers = args.encoder_layers        # GRU层数
        
        self.dropout_in_module = nn.Dropout(args.dropout)
        # 双向GRU层，用于处理输入序列
        self.rnn = nn.GRU(
            self.embed_dim, 
            self.hidden_dim, 
            self.num_layers, 
            dropout=args.dropout, 
            batch_first=False, 
            bidirectional=True
        )
        self.dropout_out_module = nn.Dropout(args.dropout)
        
        self.padding_idx = dictionary.pad() # 填充索引，用于处理可变长度的输入序列。
        
    def combine_bidir(self, outs, bsz: int):  
        # outs: 双向RNN的输出，[seq_len, batch_size, hidden_dim * 2],序列长度、批次大小和双向 RNN 的隐藏状态维度（2 倍的 hidden_dim）
        # bsz: 当前batch的大小
        # view: [self.num_layers, 2, bsz, -1]， 2表示RNN双向的两个方向，-1 表示自动计算的隐藏状态维度hidden_dim
        # transpose: [self.num_layers, bsz, 2, -1]
        # 调用 contiguous() 来确保张量在内存中的布局是连续的
        out = outs.view(self.num_layers, 2, bsz, -1).transpose(1, 2).contiguous()
        # 将 out 重新调整为形状 [self.num_layers, bsz, hidden_dim * 2]，即将双向的两个隐藏状态拼接在一起，成为一个新的隐藏状态张量。
        # 这里 -1 表示自动计算合并后的隐藏状态维度，等于 hidden_dim * 2。
        return out.view(self.num_layers, bsz, -1)

    # 执行编码器的前向传播，处理输入的 token 序列并生成输出
    def forward(self, src_tokens, **unused):
        bsz, seqlen = src_tokens.size()
        
        # get embedding 获取输入token的嵌入向量，并进行dropout操作
        x = self.embed_tokens(src_tokens)
        x = self.dropout_in_module(x)

        # [batch_size, sequence_length, hidden_dim] -> [sequence_length,batch_size,hidden_dim]
        # B x T x C -> T x B x C
        x = x.transpose(0, 1)
        
        # pass thru bidirectional RNN
        # 初始化GRU的隐藏状态h0[2*num_layers,batch_size,hidden_dim]
        h0 = x.new_zeros(2 * self.num_layers, bsz, self.hidden_dim)
        x, final_hiddens = self.rnn(x, h0)
        outputs = self.dropout_out_module(x)
        # outputs = [sequence len, batch size, hid dim * directions]
        # hidden =  [num_layers * directions, batch size  , hid dim]
        
        # Since Encoder is bidirectional, we need to concatenate the hidden states of two directions
        final_hiddens = self.combine_bidir(final_hiddens, bsz)
        # hidden =  [num_layers , batch , num_directions*hidden]
        
        encoder_padding_mask = src_tokens.eq(self.padding_idx).t()
        return tuple(
            (
                outputs,  # seq_len , batch , hidden
                final_hiddens,  # num_layers , batch , num_directions*hidden
                encoder_padding_mask,  # seq_len , batch
            )
        )
    
    def reorder_encoder_out(self, encoder_out, new_order):
        # This is used by fairseq's beam search. How and why is not particularly important here.
        return tuple(
            (
                encoder_out[0].index_select(1, new_order),   # outputs
                encoder_out[1].index_select(1, new_order),   # final_hiddens
                encoder_out[2].index_select(1, new_order),   # encoder_padding_mask
            )
        )

解码器

根据编码器的输出生成目标序列，分为训练阶段和推理阶段，略有不同

1. 训练阶段(teaching forcing)
   在训练阶段，解码器知道整个目标序列，它使用前一个正确的 token（即目标序列的上一个 token）作为当前时间步的输入。这种方式称为 Teacher Forcing

• 输入：

  • 初始输入：在序列开始时，解码器通常会接收到一个特殊的开始标记（如 ，表示 “Beginning of Sequence”）作为输入，是目标序列中上一个时间步的实际token ID,形状是(target_seq_len, batch_size)。可以稳定训练，加速收敛，因为训练早期模型生成的token可能不准确，通过使用实际的目标 token 作为输入，可以让模型在训练时保持在正确的轨道上，学习更稳定。能更快地学会生成目标序列的模式，训练过程更快收敛。

• 嵌入层：
  解码器的每个输入 token（包括 和前一个时间步的输出 token）都会通过嵌入层转换成嵌入向量(target_seq_len, batch_size, embed_dim)。embed_dim 是嵌入向量的维度。

• 序列模型：
  在每个时间步接收嵌入向量和隐藏状态，生成当前时间步的输出和更新的隐藏状态。

  • 隐藏状态：解码器会在每个时间步更新它的隐藏状态，这个隐藏状态将在下一个时间步作为输入的一部分。解码器的初始隐藏状态通常是由编码器的最终隐藏状态传递过来的。在双向 RNN 结构中，这个隐藏状态可以是编码器的最后一层前向和后向隐藏状态的拼接。
    • RNN、LSTM、GRU：
      • 形状：(num_layers, batch_size, hidden_dim)
      • num_layers 是 RNN 层的数量。
      • hidden_dim 是隐藏状态的维度。
    • Transformer:
      • (seq_len, batch_size, embed_dim)

• 注意力机制（可选）：如果使用注意力机制，解码器还会基于编码器的输出和当前的隐藏状态计算注意力权重，以对编码器的隐藏状态进行加权求和。这有助于生成时更好地关注输入序列的相关部分。

• 输出：

  • 生成的token:
    • 预测的token概率分布：解码器的最后一层通常是一个全连接层，用于将隐藏状态映射到词汇表中的每个词的概率分布，(target_seq_len, batch_size, vocab_size)，vocab_size 是词汇表的大小
    • 最终生成的 token： 是根据这个概率分布选取的。解码器会在每个时间步生成一个 token，直到生成一个结束标记（如 ，表示 “End of Sequence”）或者达到最大长度，这个输出与实际目标序列的 token 进行比较，以计算损失，在训练时，目标序列通常包括 ，以帮助模型学习生成结束标记
  • 隐藏状态

• 步骤：

  1. 初始化：用编码器的最终隐藏状态初始化解码器的隐藏状态，并输入 作为第一个 token。

  2. 每个时间步：

     • 输入目标序列的上一个 token 以及当前隐藏状态到解码器。
     • 解码器输出当前时间步的预测 token。
     • 计算损失：将解码器的输出与实际目标序列的当前 token 进行比较，并计算损失。

  3. 更新：使用损失反向传播更新模型参数。

2. 推理阶段
   在推理阶段，解码器并不知道目标序列。它使用自己上一步生成的 token 作为当前时间步的输入，逐步生成整个序列。解码器通常是一个 token 一个 token 地进行输入和输出的

• 步骤：

  1. 初始化：与训练阶段相同，解码器的隐藏状态用编码器的最终隐藏状态初始化，并输入 作为第一个 token。
  2. 每个时间步：
     • 使用解码器在前一个时间步生成的 token 作为当前时间步的输入。
     • 解码器输出当前时间步的预测 token。
     • 将预测 token 作为下一个时间步的输入。
     • 如果生成了 ，则终止解码；否则继续。
  3. 输出：最终解码器生成的 token 序列作为输出序列。

• 输入：

  • 在推理的开始阶段，解码器的输入通常是一个特殊的起始标记（），表示序列的开始。
  • 形状：(1, batch_size)，其中 1 是时间步的数量（在初始阶段只有一个 token），batch_size 是批处理的大小。

• 嵌入向量：(1, batch_size, embed_dim)

• 生成token概率分布：
  解码器生成一个 token 的概率分布，这个概率分布表示当前时间步每个词汇的概率。(1, batch_size, vocab_size)

• 更新输入，将生成的token作为下一个时间步的输入，经过嵌入曾，再次生成token分布，一直充分生成token，直到生成 或达到最大长度

class RNNDecoder(FairseqIncrementalDecoder):
    def __init__(self, args, dictionary, embed_tokens):
        super().__init__(dictionary)
        self.embed_tokens = embed_tokens
        
        # 解码器和编码器的层数必须相同
        assert args.decoder_layers == args.encoder_layers, f"""seq2seq rnn requires that encoder 
        and decoder have same layers of rnn. got: {args.encoder_layers, args.decoder_layers}"""
        # 解码器的隐藏层维度必须是编码器隐藏层维度的两倍，因为在许多的seq2seq模型中，编码器的输出可能是双向的(双向GRU或LSTM)
        assert args.decoder_ffn_embed_dim == args.encoder_ffn_embed_dim*2, f"""seq2seq-rnn requires 
        that decoder hidden to be 2*encoder hidden dim. got: {args.decoder_ffn_embed_dim, args.encoder_ffn_embed_dim*2}"""
        
        self.embed_dim = args.decoder_embed_dim      # 解码器的嵌入维度
        self.hidden_dim = args.decoder_ffn_embed_dim # 解码器RNN的hidden layers维度
        self.num_layers = args.decoder_layers        # 解码器RNN的层数
        
        
        self.dropout_in_module = nn.Dropout(args.dropout)
        self.rnn = nn.GRU(
            self.embed_dim, 
            self.hidden_dim, 
            self.num_layers, 
            dropout=args.dropout, 
            batch_first=False, 
            bidirectional=False
        )
        self.attention = AttentionLayer(
            self.embed_dim, self.hidden_dim, self.embed_dim, bias=False
        ) 
        # self.attention = None
        self.dropout_out_module = nn.Dropout(args.dropout)
        
        if self.hidden_dim != self.embed_dim:
            self.project_out_dim = nn.Linear(self.hidden_dim, self.embed_dim)
        else:
            self.project_out_dim = None
        
        if args.share_decoder_input_output_embed:
            self.output_projection = nn.Linear(
                self.embed_tokens.weight.shape[1],
                self.embed_tokens.weight.shape[0],
                bias=False,
            )
            self.output_projection.weight = self.embed_tokens.weight
        else:
            self.output_projection = nn.Linear(
                self.output_embed_dim, len(dictionary), bias=False
            )
            nn.init.normal_(
                self.output_projection.weight, mean=0, std=self.output_embed_dim ** -0.5
            )
        
    def forward(self, prev_output_tokens, encoder_out, incremental_state=None, **unused):
        # extract the outputs from encoder
        encoder_outputs, encoder_hiddens, encoder_padding_mask = encoder_out
        # outputs:          seq_len x batch x num_directions*hidden
        # encoder_hiddens:  num_layers x batch x num_directions*encoder_hidden
        # padding_mask:     seq_len x batch
        
        if incremental_state is not None and len(incremental_state) > 0:
            # if the information from last timestep is retained, we can continue from there instead of starting from bos
            prev_output_tokens = prev_output_tokens[:, -1:]
            cache_state = self.get_incremental_state(incremental_state, "cached_state")
            prev_hiddens = cache_state["prev_hiddens"]
        else:
            # incremental state does not exist, either this is training time, or the first timestep of test time
            # prepare for seq2seq: pass the encoder_hidden to the decoder hidden states
            prev_hiddens = encoder_hiddens
        
        bsz, seqlen = prev_output_tokens.size()
        
        # embed tokens
        x = self.embed_tokens(prev_output_tokens)
        x = self.dropout_in_module(x)

        # B x T x C -> T x B x C
        x = x.transpose(0, 1)
                
        # decoder-to-encoder attention
        if self.attention is not None:
            x, attn = self.attention(x, encoder_outputs, encoder_padding_mask)
                        
        # pass thru unidirectional RNN
        x, final_hiddens = self.rnn(x, prev_hiddens)
        # outputs = [sequence len, batch size, hid dim]
        # hidden =  [num_layers * directions, batch size  , hid dim]
        x = self.dropout_out_module(x)
                
        # project to embedding size (if hidden differs from embed size, and share_embedding is True, 
        # we need to do an extra projection)
        if self.project_out_dim != None:
            x = self.project_out_dim(x)
        
        # project to vocab size
        x = self.output_projection(x)
        
        # T x B x C -> B x T x C
        x = x.transpose(1, 0)
        
        # if incremental, record the hidden states of current timestep, which will be restored in the next timestep
        cache_state = {
            "prev_hiddens": final_hiddens,
        }
        self.set_incremental_state(incremental_state, "cached_state", cache_state)
        
        return x, None
    
    def reorder_incremental_state(
        self,
        incremental_state,
        new_order,
    ):
        # This is used by fairseq's beam search. How and why is not particularly important here.
        cache_state = self.get_incremental_state(incremental_state, "cached_state")
        prev_hiddens = cache_state["prev_hiddens"]
        prev_hiddens = [p.index_select(0, new_order) for p in prev_hiddens]
        cache_state = {
            "prev_hiddens": torch.stack(prev_hiddens),
        }
        self.set_incremental_state(incremental_state, "cached_state", cache_state)
        return