为何Transformer模型无法训练实现加1任务？

你的Transformer模型无法收敛的核心原因存在几个关键实现错误和选型问题，下面逐一拆解并给出修正方案：

1. 固定非可训练的目标序列（tgt）是核心问题 #

你在__init__中定义的self.tgt是普通Tensor，没有被注册为模型参数，训练过程中不会被optimizer更新，始终是初始化时的随机值。Decoder基于这个固定随机序列计算输出，完全无法建立输入（X）和目标（X+1）的映射关系，自然不可能收敛。

2. 输入数值范围过大导致训练不稳定 #

你的输入是0到10^{7的整数，数值量级极大，会让MSELoss的初始值达到10}14级别，梯度更新极易出现爆炸或消失，即使是简单的FFN也需要更长时间收敛，更不用说结构更复杂的Transformer。

3. Transformer结构选型不符合任务需求 #

加1是单输入单输出的静态映射任务，不需要完整的Encoder-Decoder架构（该架构更适合翻译、文本生成等序列到序列任务），只用TransformerEncoder直接输出预测值即可，或者调整Decoder的输入逻辑。

修正后的代码示例 #

下面是调整后的代码，解决了上述所有问题：

import torch
from torch import nn
import torch.optim as optim

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
batch = 1024
# 缩小输入范围，避免数值爆炸
upper = 100

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 改用TransformerEncoder，适配单输入单输出的映射任务
        encoder_layer = nn.TransformerEncoderLayer(d_model=1, nhead=1, dim_feedforward=4)
        self.trans_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)
        # 加一层线性层映射到输出（可选，d_model=1时可直接输出，但增加一层更稳定）
        self.fc = nn.Linear(1, 1)

    def forward(self, x):
        # Encoder输出后直接映射
        x = self.trans_encoder(x)
        x = self.fc(x)
        return x

model = MyModel()
model.to(device)

optimizer = optim.Adam(model.parameters(), lr=0.001)
# 降低学习率，配合缩小后的输入范围
criterion = nn.MSELoss()

running_loss = 0.0
for epoch in range(1000):
    # 生成0-99的随机数，转换为浮点型并调整形状为(S,B,E)=(1, batch, 1)
    training = torch.randint(low=0, high=upper, size=(batch,)).float().to(device)
    training = training.unsqueeze(0).unsqueeze(-1)
    labels = training + 1

    optimizer.zero_grad()
    outputs = model(training)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    running_loss += loss.item()
    if (epoch + 1) % 100 == 0:
        print(f'[{epoch + 1}] loss: {running_loss / 100:.6f}')
        running_loss = 0.0

print('Finished Training')
# 测试模型效果
test_input = torch.tensor([[5.0]]).unsqueeze(0).to(device)
print(f"输入5，预测输出：{model(test_input).item():.2f}")

额外优化建议 #

• 如果一定要用完整的Encoder-Decoder结构，需要将tgt设置为可训练参数（用nn.Parameter包裹），或者让tgt与输入关联（比如用输入作为Decoder的初始输入），但这对于单步映射任务完全没必要。
• 当d_model=1时，建议调小学习率（比如0.001），并适当增加dim_feedforward的维度（比如4或8），提升模型的拟合能力和训练稳定性。
• 对输入做归一化处理（比如除以upper缩放到0-1区间），可以进一步提升训练效率。

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