LSTM编解码器释义生成模型损失不下降、同轮生成重复问题求助

已知问题现象 #

• 训练损失全程无下降，模型完全不收敛
• 单个epoch内所有样本生成结果完全一致，仅在epoch间权重更新后输出会整体变化，同轮次输出无样本区分度

现有基础信息 #

模型结构 #

StackedResidualLSTM(
  (encoder): RecurrentEncoder(
    (embed_tokens): Embedding(30522, 256)
    (dropout): Dropout(p=0.5, inplace=False)
    (rnn): LSTM(256, 256, num_layers=2, batch_first=True, dropout=0.5)
  )
  (decoder): RecurrentDecoder(
    (embed_tokens): Embedding(30522, 128)
    (dropout_in_module): Dropout(p=0.5, inplace=False)
    (dropout_out_module): Dropout(p=0.1, inplace=False)
    (layers): ModuleList(
      (0): LSTMCell(384, 256)
      (1): LSTMCell(256, 256)
    )
    (fc_out): Linear(in_features=256, out_features=30522, bias=True)
  )
)

训练样本样例 #

Source: [CLS] where can i get quality services in brisbane for plaster and drywall repair? [SEP] [PAD]*N

Decoder Input: [CLS] [CLS] where can i get quality services for plaster and drywall repairs in brisbane? [SEP] [PAD]*N

Preds: [CLS] the? [SEP]? [SEP]? [SEP]? [SEP]? [SEP]? [SEP]? [SEP]? [SEP]? [SEP]

Target: [CLS] where can i get quality services for plaster and drywall repairs in brisbane? [SEP] [PAD]*N

现有核心代码 #

损失计算逻辑：

loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))

Torch-Ignite训练循环逻辑：

model.train()
accumulation_steps = cfg.train.get("accumulation_steps", 1)

if batch["labels"].device != device:
      batch = {k: v.to(device) for (k, v) in batch.items()}

      input_ids = batch["input_ids"]
      labels = batch["labels"]

      loss = model(input_ids=input_ids, labels=labels)[0]
      loss /= accumulation_steps

      scaler.scale(loss).backward()
      if engine.state.iteration % accumulation_steps == 0:
          scaler.step(optimizer)
          scaler.update()
          optimizer.zero_grad()

      return {"batch loss": loss.item()}

排查修复方向（按优先级从高到低排序） #

1. 优先排查训练循环逻辑错误 #

当前训练循环把设备迁移、前向传播、反向传播全写在了if batch["labels"].device != device判断块内部。第一个batch完成设备迁移后，后续所有batch的labels设备已经和目标设备一致，永远不会进入判断块，等于整个epoch除了第一个batch，剩下的所有步都不执行前向、反向、参数更新。
这个bug和描述的现象完全匹配：单epoch内只有第一个batch更新一次权重，后续所有样本推理都用同一套冻结的权重，输出自然完全一致；下一个epoch第一个batch再次更新权重，整轮输出整体变化一次，之后又保持固定。

修复方式是把设备迁移逻辑和训练逻辑拆开：

model.train()
accumulation_steps = cfg.train.get("accumulation_steps", 1)
# 仅执行设备迁移
if batch["labels"].device != device:
    batch = {k: v.to(device) for (k, v) in batch.items()}
# 前向、反向、参数更新逻辑移到判断块外，每个batch都执行
input_ids = batch["input_ids"]
labels = batch["labels"]
loss = model(input_ids=input_ids, labels=labels)[0]
loss /= accumulation_steps

scaler.scale(loss).backward()
if engine.state.iteration % accumulation_steps == 0:
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

return {"batch loss": loss.item()}

验证方式：在原判断块内部加打印，会发现只有每个epoch的第一个batch会进入代码块。

2. 排查Decoder输入拼接逻辑 #

第一层LSTMCell输入维度为384，正常逻辑是128维token embedding拼接256维编码器上下文向量，维度刚好匹配。需要确认两点：

• 编码器输出是否做了错误聚合，比如对所有时间步的encoder output取平均后，固定为同一个向量传给所有时间步、所有batch样本，导致解码器拿不到不同样本的差异化输入
• 训练时teacher forcing逻辑是否异常，有没有出现不管输入什么样本，解码器第一步之后的输入都没有正确传入对应位置token embedding的问题

3. 排查残差连接与初始状态逻辑 #

带残差的堆叠LSTM结构很容易出现连接错误：

• 确认两层LSTMCell的初始hidden、cell状态是否正确用编码器最终输出生成，有没有出现初始状态全0、或者所有样本初始状态完全一致的问题
• 确认残差连接是否加对位置，有没有把某一层输出错误覆盖为固定值，导致梯度完全传不回编码器

4. 排查正则项与学习率配置 #

当前编码器embedding后、LSTM层间dropout为0.5，解码器输入dropout也为0.5，比例过高很容易直接淹没输入信号：

• 先把所有dropout概率调到0.1-0.2区间，用小批量数据测试能不能正常过拟合小样本集
• 确认优化器学习率是否在合理区间，LSTM序列生成模型初始学习率一般在1e-3到1e-4区间，过高会导致梯度爆炸损失不收敛，过低会导致参数几乎不更新

5. 排查损失计算逻辑 #

• 确认CrossEntropyLoss输入维度顺序是否正确，logits最后一维是否为词表维度，有没有维度转置错误导致损失计算位置完全不对
• 确认label的padding mask是否正确替换为-100，有没有把非padding的有效token也错误替换成-100，导致损失计算没有有效监督信号

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