Transformer序列复制任务效果不佳，求优化方向指引

问题描述 #

我在学习过程中尝试实现Transformer模型完成输入序列到输出序列的复制任务，原本以为实现相对简单，但结果却不尽人意——损失值高于预期，导致序列中的数字复制不准确。我已经尝试了大量超参数调优，但结果没有实质性改善，希望能得到优化方向指引：我的模型架构是否存在问题？有没有遗漏的算法BUG？

运行示例 #

Sample run
Epoch 0 Loss 1.9481123001017469
Epoch 1 Loss 1.4541001472067325
Epoch 2 Loss 1.2569004525529577
Epoch 3 Loss 1.158278153297749
Epoch 4 Loss 1.1283172952367904
Test Case 0
Expected: [2, 3, 4, 2, 5, 9, 6, 7]
Actual: [2, 6, 4, 9, 3, 2, 7, 5]

Test Case 1
Expected: [8, 2, 5, 6, 7, 9, 4, 8]
Actual: [8, 2, 9, 4, 6, 5, 8, 7]

Test Case 2
Expected: [8, 5, 6, 3, 4, 5, 6, 8]
Actual: [8, 6, 5, 3, 4, 6, 8, 5]

Test Case 3
Expected: [5, 5, 7, 5, 2, 2, 9, 2]
Actual: [2, 5, 2, 5, 9, 7, 2]

Test Case 4
Expected: [4, 5, 9, 5, 3, 5, 7, 8]
Actual: [5, 9, 5, 3, 8, 4, 7]

模型实现代码 #

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import math

from torch.utils.data import Dataset, DataLoader

# Set the random seed for reproducibility
torch.manual_seed(0)

# Define the device to run the model on
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Model hyperparameters
input_size = 10  # Size of the input vocabulary (number of unique tokens)
num_layers = 6  # Number of transformer layers
d_model = 512  # The number of expected features in the encoder/decoder inputs
nhead = 8  # The number of heads in the multiheadattention model
dim_feedforward = 2048  # Dimension of the feedforward network model in nn.TransformerEncoder
dropout = 0.01  # Dropout rate
max_seq_length = 10  # Maximum sequence length

# Training hyperparameters
batch_size = 64
num_samples = 3000
epochs = 5
lr = 0.0001  # Learning rate

SOS_token = np.array([0])
EOS_token = np.array([1])


class TransformerModel(nn.Module):
    def __init__(self, input_size, d_model, nhead, num_layers, dim_feedforward, dropout=0.1):
        super(TransformerModel, self).__init__()
        self.model_type = 'Transformer'
        self.tgt_mask = None
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        self.embedding = nn.Embedding(input_size, d_model)
        self.transformer = nn.Transformer(d_model=d_model, nhead=nhead, num_encoder_layers=num_layers, num_decoder_layers=num_layers,
                                          dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True)
        self.linear = nn.Linear(d_model, input_size)

    def _generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def forward(self, src, tgt):
        tgt_mask = self._generate_square_subsequent_mask(tgt.size(1)).to(device)

        src = self.embedding(src) * math.sqrt(d_model)
        tgt = self.embedding(tgt) * math.sqrt(d_model)
        src = self.pos_encoder(src)
        tgt = self.pos_encoder(tgt)

        output = self.transformer(src, tgt, tgt_mask=tgt_mask)
        output = self.linear(output)
        return output


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=1024):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)




class TrainingDataSet(Dataset):
    def __init__(self, num_tokens, seq_length, size):
        self.num_tokens = num_tokens
        self.seq_length = seq_length
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, index):
        sequence = np.concatenate(
            (SOS_token, torch.randint(2, 10, (self.seq_length,), dtype=torch.long), EOS_token))
        return sequence, sequence  # input and target are the same in the copy task


# Create the dataset and dataloader
dataset = TrainingDataSet(input_size,
                          max_seq_length - 2,  # -2 due to SOS end EOS tokens
                          num_samples)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Initialize the model
model = TransformerModel(input_size, d_model, nhead, num_layers, dim_feedforward, dropout).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)

# Training loop
model.train()
for epoch in range(epochs):
    total_loss = 0.
    for src, tgt in dataloader:
        src = src.to(device)
        tgt = tgt.to(device)

        tgt_input = tgt[:, :-1]
        tgt_expected = tgt[:, 1:]

        optimizer.zero_grad()
        output = model(src, tgt_input)
        # loss = criterion(output.view(-1, input_size), tgt_expected.view(-1))
        # loss = criterion(output.reshape(-1, input_size), tgt_expected.reshape(-1))

        # permute (N, L, C) -> (N, C, L) where N is batch, L is seq len, C is class/vocab_size
        output = output.permute(0, 2, 1)
        # shape of tgt_expected is (N, L) where N is batch, L is seq len
        loss = criterion(output, tgt_expected)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    print(f'Epoch {epoch} Loss {total_loss / len(dataloader)}')

# Save the model
torch.save(model.state_dict(), 'transformer_copy_task.pth')

def predict(model, input_sequence):
    """
    Method from "A detailed guide to Pytorch's nn.Transformer() module.", by
    Daniel Melchor: https://medium.com/@danielmelchor/a-detailed-guide-to-pytorchs-nn-transformer-module-c80afbc9ffb1
    """
    model.eval()

    y_input = torch.tensor([SOS_token], dtype=torch.long, device=device)

    for _ in range(max_seq_length+2):
        pred = model(input_sequence, y_input)

        next_item = pred.topk(1)[1].view(-1)[-1].item()  # num with highest probability
        next_item = torch.tensor([[next_item]], device=device)

        # Concatenate previous input with predicted best word
        y_input = torch.cat((y_input, next_item), dim=1)

        # Stop if model predicts end of sentence
        if next_item.view(-1).item() == EOS_token:
            break

    return y_input.view(-1).tolist()


# Test the model
model.eval()

# Here we test some examples to observe how the model predicts
examples = [
    torch.tensor([[0, 2, 3, 4, 2, 5, 9, 6, 7, 1]], dtype=torch.long, device=device),
    torch.tensor([[0, 8, 2, 5, 6, 7, 9, 4, 8, 1]], dtype=torch.long, device=device),
    torch.tensor([[0, 8, 5, 6, 3, 4, 5, 6, 8, 1]], dtype=torch.long, device=device),
    torch.tensor([[0, 5, 5, 7, 5, 2, 2, 9, 2, 1]], dtype=torch.long, device=device),
    torch.tensor([[0, 4, 5, 9, 5, 3, 5, 7, 8, 1]], dtype=torch.long, device=device),
]

for idx, example in enumerate(examples):
    result = predict(model, example)
    print(f"Test Case {idx}")
    print(f"Expected: {example.view(-1).tolist()[1:-1]}")
    print(f"Actual: {result[1:-1]}")
    print()

优化方向与BUG修复 #

1. 位置编码维度匹配错误 #

你的PositionalEncoding类中，由于模型设置了batch_first=True，输入x的形状是[batch_size, seq_len, d_model]，但self.pe的形状是[max_len, 1, d_model]，导致维度无法正确相加。

修改PositionalEncoding的初始化和前向传播：

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=1024):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # 调整为 [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

2. 训练轮数严重不足 #

当前仅训练5轮，对于Transformer模型来说完全不够。将epochs调整为50-100轮，观察损失是否持续下降至接近0。

3. 模型规模过度冗余 #

复制任务属于简单序列任务，你设置的6层Transformer、512维d_model过于庞大，会导致收敛缓慢甚至过拟合。建议缩小模型：

• num_layers改为2-3层
• d_model改为128-256维
• dim_feedforward改为512-1024维

4. 数据类型不统一问题 #

TrainingDataSet中拼接了numpy数组和torch tensor，可能导致类型隐患。统一使用torch tensor：

SOS_token = torch.tensor([0], dtype=torch.long)
EOS_token = torch.tensor([1], dtype=torch.long)

# __getitem__方法修改为
sequence = torch.cat(
    (SOS_token, torch.randint(2, 10, (self.seq_length,), dtype=torch.long), EOS_token)
)

5. 学习率与优化器调整 #

当前lr=0.0001偏低，可尝试提升至0.001，或添加学习率调度器：

optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

# 训练循环中添加
scheduler.step()

6. 预测逻辑优化 #

预测循环次数可调整为max_seq_length，避免多余迭代；同时确保输入形状一致性：

def predict(model, input_sequence):
    model.eval()
    y_input = torch.tensor([[0]], dtype=torch.long, device=device)  # 直接用tensor，避免numpy转换

    for _ in range(max_seq_length):
        pred = model(input_sequence, y_input)
        next_item = pred.topk(1)[1].view(-1)[-1].item()
        next_item = torch.tensor([[next_item]], device=device)
        y_input = torch.cat((y_input, next_item), dim=1)
        
        if next_item.item() == 1:  # EOS_token是1，直接判断
            break

    return y_input.view(-1).tolist()

按上述修改后重新训练，Transformer模型应该能快速收敛，实现近乎完美的序列复制。

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