为何我的递归模型无法在等差数列简单任务上收敛？

我想用递归模型解决最简单的序列问题——等差数列：以a为首项、d为公差，序列形式为：

a, a+d, a+2d, a+3d, a+4d, ...

我设定隐藏状态为h，模型需要学习一个2×2矩阵（实际要实现h₁ = t₀的逻辑），因此设计了一个带2×2全连接层的模型，但模型始终无法收敛，怀疑问题出在训练循环里。特意将batch size设为1，后续再处理填充问题，理论上单样本也应该能完成学习。

完整代码 #

数据集与模型定义 #

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import numpy as np

class CustomDataset(Dataset):
    def __init__(self, size):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, index):
        a0 = (np.random.rand() - 0.5) * 200
        d = (np.random.rand() - 0.5) * 40
        # 注意：MAX_Length_sequence未定义，需提前声明
        length = np.random.randint(2, MAX_Length_sequence + 1)

        sequence = np.arange(length) * d + a0
        next_number = sequence[-1] + d

        return length, torch.tensor(sequence, dtype=torch.float32), torch.tensor(next_number, dtype=torch.float32)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(2, 2, bias=False)

    def forward(self, x):
        x = self.fc1(x)
        return x

# 超参数
EPOCHS = 10
BATCH_SIZE = 1
LEARNING_RATE = 0.001
DATASET_SIZE = 10000
criterion = nn.MSELoss()

# 模型与优化器
model = Model()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

训练循环 #

for epoch in range(EPOCHS):
    dataset = CustomDataset(DATASET_SIZE)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE)
    model.train()
    total_loss = 0

    for length, sequence, next_number in dataloader:
        optimizer.zero_grad()
        loss = 0
        h = torch.zeros(BATCH_SIZE)

        for i in range(length):
            x = torch.cat([h, sequence[0, i].unsqueeze(0)])
            y = sequence[0, i + 1] if i != length - 1 else next_number[0]

            output = model(x)
            h, y_hat = output[0].unsqueeze(0), output[1]

            loss += criterion(y_hat, y)

        loss.backward()
        optimizer.step()
        total_loss += loss.item() 
    
    print(f'Epoch {epoch+1}, Loss: {total_loss/len(dataloader)}')

问题排查与修复建议 #

1. 数据集重复创建 #

每次epoch重新创建CustomDataset会导致数据完全随机，模型难以稳定学习。应将数据集创建移到epoch循环外，同时开启shuffle增强泛化：

# 提前定义缺失常量
MAX_Length_sequence = 10
# 数据集仅创建一次
dataset = CustomDataset(DATASET_SIZE)

for epoch in range(EPOCHS):
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
    # 后续训练代码

2. 张量维度处理错误 #

length是张量类型，直接用range(length)会报错，需转换为整数：length = length_tensor.item()

3. 损失累积导致梯度不稳定 #

当前损失是累加每个时间步的结果，长序列会产生更大的损失值，导致梯度震荡。应对损失取平均：

loss = loss / length

4. 模型逻辑与权重初始化 #

等差数列的下一项满足y_hat = 2*t_i - h_prev（其中h_prev是上一个序列值，t_i是当前值），对应的2×2权重矩阵应为：

[[0, 1],
 [-1, 2]]

可手动初始化权重帮助模型快速收敛：

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(2, 2, bias=False)
        # 初始化权重接近预期值
        with torch.no_grad():
            self.fc1.weight.copy_(torch.tensor([[0.0, 1.0], [-1.0, 2.0]]))

    def forward(self, x):
        x = self.fc1(x)
        return x

5. 学习率调整 #

Adam默认0.001的学习率对简单任务可能偏大，可尝试调整为0.01或0.0005，提升收敛速度。

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