自定义Transformer模型仅预测START/END tokens问题求助

我从零实现并训练了一个用于共情对话生成的Transformer模型，当前遇到严重问题：无论输入何种目标token到decoder，模型输出层始终只倾向于预测START和END token。已多次检查并修正MultiHead Attention层、tokenization等模块，但问题仍未解决。

损失计算细节 #

使用F.cross_entropy计算损失，输入为Transformer输出的logits（out[:, :-1:, :]）和数据集目标序列（target[:, 1:]），该移位操作符合Transformer预测下一个token的逻辑。曾尝试移除START和END token计算损失（即out[:, :-2:, :]和target[:, 1:-1]），但无效果。logits和target的形状完全符合PyTorch文档要求：分别为(batch_size, sequence_length, classes)和(batch_size, sequence_length)，target为类别索引（padding index已被忽略）。

训练现象 #

训练过程中损失下降至5-6区间后便不再变化，即使训练到第50个epoch也是如此。训练日志片段如下：

Epoch 0:   1%|          | 1/180 [00:05<16:28,  5.53s/it, loss=11, v_num=3, train_loss=11.00]
Epoch 0:   1%|          | 2/180 [00:25<37:55, 12.78s/it, loss=11, v_num=3, train_loss=11.00]
...
Epoch 5:  90%|█████████ | 162/180 [00:58<00:06,  2.77it/s, loss=5.54, v_num=3, train_loss=5.520]
Epoch 5:  90%|█████████ | 162/180 [00:58<00:06,  2.77it/s, loss=5.53, v_num=3, train_loss=5.430]

对logits做softmax后查看概率张量，发现START和END token的概率始终远高于其他token。用beam search测试生成：给定任意上下文序列，第一个目标token设为[START]，模型预测的下一个token始终是[END]，与beam width无关。

已尝试的排查措施 #

• 加入dropout层，尝试解决权重爆炸问题，无效
• 移除decoder中的情感embedding，问题仍存在
• 更换Adam、AdamW优化器，搭配StepLR调度器，无改善

模型完整代码 #

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size: int, heads: int) -> None:
        super().__init__()

        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = self.embed_size // self.heads

        assert self.head_dim * self.heads == self.embed_size

        self.values = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.keys = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.queries = nn.Linear(self.embed_size, self.embed_size, bias=False)

        self.fc_out = nn.Linear(self.embed_size, self.embed_size, bias=False)

    def forward(
        self,
        keys: torch.Tensor, 
        values: torch.Tensor, 
        queries: torch.Tensor, 
        mask: torch.Tensor
    ) -> torch.Tensor:

        N = queries.shape[0]
        keys_len, values_len, queries_len = keys.shape[1], values.shape[1], queries.shape[1]

        values = self.values(values).reshape(N, values_len, self.heads, self.head_dim)
        keys = self.keys(keys).reshape(N, keys_len, self.heads, self.head_dim)
        queries = self.queries(queries).reshape(N, queries_len, self.heads, self.head_dim)

        scores = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])

        # Apply mask to attention scores if specified
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-1e20"))

        # Normalise with respect to all keys
        attention = F.softmax(scores / (self.embed_size ** 0.5), dim=-1)

        out = torch.einsum("nhqk,nvhd->nqhd", [attention, values])
        out = self.fc_out(out.reshape(N, queries_len, self.embed_size))

        return out


class TransformerBlock(nn.Module):
    def __init__(
        self,
        embed_size: int, 
        heads: int, 
        dropout: float, 
        forward_expansion: int
    ) -> None:

        super().__init__()

        self.attention = MultiHeadAttention(embed_size, heads)

        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)

        self.dropout = nn.Dropout(dropout)

        self.ff = nn.Sequential(
            nn.Linear(embed_size, embed_size * forward_expansion),
            nn.ReLU(),
            nn.Linear(embed_size * forward_expansion, embed_size)
        )

    def forward(
        self,
        keys: torch.Tensor, 
        values: torch.Tensor, 
        queries: torch.Tensor, 
        mask: torch.Tensor
    ) -> torch.Tensor:

        attention = self.attention(keys, values, queries, mask)

        contextualised = self.dropout(self.norm1(attention + queries))
        forward = self.ff(contextualised)
        out = self.dropout(self.norm2(forward + contextualised))

        return out

class Encoder(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        padding_idx: int,
        num_layers: int,
        embed_size: int,
        heads: int,
        dropout: float, 
        forward_expansion: int,
        max_seq_len: int,
        num_of_emo_labels: int
    ) -> None:

        super().__init__()

        self.word_embeddings = nn.Embedding(
            vocab_size + 1, embed_size, padding_idx=padding_idx)
        self.pos_embeddings = nn.Embedding(max_seq_len, embed_size)
        self.ds_embeddings = nn.Embedding(2 + 1, embed_size, padding_idx=0)

        self.layers = nn.ModuleList(
            [TransformerBlock(embed_size, heads, dropout, forward_expansion)
             for _ in range(num_layers)]
        )

        self.dropout = nn.Dropout(dropout)
    
    def forward(
        self, 
        context: torch.Tensor, 
        context_ds_state: torch.Tensor,
        mask: torch.Tensor,
        emotion_label: torch.Tensor
    ) -> torch.Tensor:

        N, seq_len = context.shape
        positions = torch.arange(0, seq_len, device=context.device).expand(N, seq_len)

        word_embeddings = self.word_embeddings(context)
        pos_embeddings = self.pos_embeddings(positions)
        ds_embeddings = self.ds_embeddings(context_ds_state)

        out = self.dropout(word_embeddings + pos_embeddings + ds_embeddings)

        for layer in self.layers:
            out = layer(out, out, out, mask)
        
        return out

class DecoderBlock(nn.Module):
    def __init__(
        self,
        embed_size: int,
        heads: int,
        dropout: float,
        forward_expansion: int
    ) -> None:

        super().__init__()

        self.attention = MultiHeadAttention(embed_size, heads)
        self.norm = nn.LayerNorm(embed_size)
        self.transformer_block = TransformerBlock(
            embed_size,
            heads, 
            dropout, 
            forward_expansion
        )
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        keys: torch.Tensor,
        values: torch.Tensor,
        target_mask: torch.Tensor,
        input_mask: torch.Tensor
    ) -> torch.Tensor:
        
        attention = self.attention(x, x, x, target_mask)
        queries = self.dropout(self.norm(attention + x))
        out = self.transformer_block(keys, values, queries, input_mask)

        return out

class Decoder(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        padding_idx: int,
        num_layers: int,
        embed_size: int,
        heads: int,
        dropout: float, 
        forward_expansion: int,
        max_seq_len: int,
        num_of_emo_labels: int
    ) -> None:

        super().__init__()

        self.word_embeddings = nn.Embedding(
            vocab_size + 1, embed_size, padding_idx=padding_idx)
        self.pos_embeddings = nn.Embedding(max_seq_len, embed_size)
        self.ds_embeddings = nn.Embedding(2 + 1, embed_size, padding_idx=0)
        self.emotion_embedding = nn.Embedding(num_of_emo_labels, embed_size)

        self.layers = nn.ModuleList(
            [DecoderBlock(embed_size, heads, dropout, forward_expansion)
             for _ in range(num_layers)]
        )

        self.dropout = nn.Dropout(dropout)

        self.fc_out = nn.Linear(embed_size, vocab_size)

    def forward(
        self,
        target: torch.Tensor,
        target_ds_state: torch.Tensor,
        encoder_out: torch.Tensor,
        target_mask: torch.Tensor,
        input_mask: torch.Tensor,
        emotion_label: torch.Tensor
    ) -> torch.Tensor:

        N, seq_len = target.shape
        positions = torch.arange(0, seq_len, device=target.device).expand(N, seq_len)

        word_embeddings = self.word_embeddings(target)
        pos_embeddings = self.pos_embeddings(positions)
        ds_embeddings = self.ds_embeddings(target_ds_state)

        out = self.dropout(word_embeddings + pos_embeddings + ds_embeddings)
        
        for layer in self.layers:
            out = layer(out, encoder_out, encoder_out, target_mask, input_mask)
        
        emotion_embedding = self.emotion_embedding(
            emotion_label).unsqueeze(1).expand(-1, seq_len, -1)
        
        out = self.fc_out(out + emotion_embedding)

        return out

class Transformer(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        num_of_emo_labels: int,
        max_seq_len: int,
        padding_idx: int,
        num_layers: int = 6,
        embed_size: int = 256,
        heads: int = 8,
        dropout: float = 0.5, 
        forward_expansion: int = 4
    ) -> None:

        super().__init__()

        self.padding_idx = padding_idx
        self.encoder = Encoder(
            vocab_size,
            padding_idx,
            num_layers, 
            embed_size, 
            heads,
            dropout, 
            forward_expansion, 
            max_seq_len,
            num_of_emo_labels
        )

        self.decoder = Decoder(
            vocab_size,
            padding_idx,
            num_layers, 
            embed_size, 
            heads,
            dropout, 
            forward_expansion, 
            max_seq_len,
            num_of_emo_labels
        )

    def create_padding_mask(self, batch_seq):
        N = batch_seq.size(dim=0)
        padding_mask = (batch_seq != self.padding_idx).unsqueeze(1).unsqueeze(2)
        return padding_mask
    
    def create_lookahead_mask(self, batch_seq):
        N, seq_len = batch_seq.shape
        lookahead_mask = torch.tril(torch.ones(
            N, 1, seq_len, seq_len, device=batch_seq.device))
        return lookahead_mask
    
    def forward(
        self,
        context: torch.Tensor,
        target: torch.Tensor,
        context_ds_state: torch.Tensor,
        target_ds_state: torch.Tensor,
        emotion_label: torch.Tensor
    ) -> None:

        input_mask = self.create_padding_mask(context)
        target_mask = torch.minimum(
            self.create_lookahead_mask(target), 
            self.create_padding_mask(target)
        )

        encoder_out = self.encoder(
            context, 
            context_ds_state, 
            input_mask,
            emotion_label
        )
        out = self.decoder(
            target, 
            target_ds_state,
            encoder_out, 
            target_mask, 
            input_mask, 
            emotion_label
        )

        return out

希望得到技术层面的解决方案，感谢！

内容的提问来源于stack exchange，提问作者Tarush