Timexer+MKAN+PGA：时间序列预测优化的PyTorch实践

去年我接了一个设备温度场预测项目，数据是感器阵列的时序读数，外生变量不少，天气、负载、开关状态都有。基线我用的是 Timexer，公开 benchmark 上表现确实能打，尤其在多变量、多步预测场景下。但跑下来有个别扭的地方：模型对信号里的局部拐点拟合不够干脆，外生变量进来之后，网络总像是“知道该涨，但涨得拖泥带水”。调参、改层数、换 dropout 都试过，收益有限。

后来我把注意力放到两件事上：一是把 Timexer 内部原生的 MLP 替换成 MKAN（Multi-Kernel KAN），提升非线性拟合能力；二是引入 PGA（Progressive Gradient Aggregation）训练策略，解决多尺度、多损失叠加后的训练震荡。组合之后效果提升明显，而且收敛速度比预期还快了一点。这篇文章就把完整的 PyTorch 复现过程写清楚，从动机、原理到代码骨架，再到踩坑细节，尽量做到能直接照着抄。

适合谁看？如果你正在做时间序列预测，用过或者听说过 Timexer、PatchTST 这类模型，同时对 KAN 方向好奇，又苦于模型训不稳定，那这篇内容会正好对你的胃口。如果你只是想把某个指标的 MSE 往下压一点，里面也有可以直接用的调参套路。

1. 这个组合要解决的实际痛点：Timexer 虽强，但非线性拟合与训练韧性仍有短板

先不急着上代码，得把“为什么是这三样拼在一起”讲明白。只有把动机理清楚，后面调参的时候你才不至于瞎试。

1.1 Timexer 的看家本领和它“不擅长”的地方

Timexer 的核心设计思路，是同时建模内生变量（目标序列自身的历史）和外生变量（影响目标的其他序列）。它把每个变量当作 token 处理，内生变量采用 patch 化嵌入，外生变量单独做 embedding，再通过多头自注意力、内-外变量交叉注意力把信息聚合起来。这样做的优势很明显：对不同变量之间的依赖关系建模得比较充分，而且 patch 化之后计算量可控，长序列也能跑。

但它的主体结构仍然是 Transformer 那一套，每个 Transformer Block 里前馈网络用的是标准 MLP。MLP 拟合非线性靠的是固定激活函数加宽隐藏层，总体来说是一种“堆容量换表达”的思路。信号里有比较尖锐的局部模式、非平稳趋势突变时，MLP 往往需要很宽的层才能拟合好，参数量涨得很快，但效果提升越来越钝。

我当时的直观感受是：模型整体趋势预测没问题，但细节处“手抖”。尤其温度序列在某个采样点突然被设备开关状态影响发生转折时，预测曲线总是显得太圆润，像低通滤波后的结果。

1.2 KAN 系方法补的是哪块短板

KAN 的原理不复杂：把传统的线性加权加固定激活函数，改成“可学习的基函数加权组合”。你可以把 MLP 的神经元想象成固定形状的积木，只能堆数量；而 KAN 的神经元像是可塑形的粘土，每个连接上都有一组可学习的基函数权重，能根据数据自动调整“形状”。

MKAN 是 KAN 的进阶版本，思路是同时用多个不同类型的核函数（比如 Bessel、Chebyshev、RBF 等）并行计算，再学一个融合权重。好处在于：不同核函数对信号的局部变化敏感度不一样，核函数组合起来，等于给模型更丰富的“描述语言”，更贴近真实序列的复杂性。

所以把 Timexer 里的 MLP 替换成 MKAN，逻辑上很顺：保留 Transformer 的全局交互结构，但把局部非线性拟合部分换成表达能力更强的组件。

1.3 PGA 解决的是“训练韧性问题”

组合模型之后，参数变多了，梯度流也更复杂。MKAN 这种多核结构在初期很容易出现某一支路梯度特别大、其他支路梯度特别小的情况，直接导致 loss 震荡。叠加多步预测的辅助损失后，问题更明显。

PGA（Progressive Gradient Aggregation）是我在实践里逐渐摸索出来的训练策略，思路是“渐进式梯度聚合”：训练初期只激活性质量最高的主损失，让模型先学趋势；随着训练进行，逐步放大辅助损失的权重，把细节信息渐进式注入训练过程。相当于让模型先在“粗粒度”上站稳，再慢慢“细化”。

这不是学习率 warmup 能替代的。warmup 只调节步长，PGA 调节的是“梯度信号组合的比例”，两者解决的问题完全不在一个维度上。

2. Transformer 之外的隐藏功臣：MKAN 多核基函数增强的前向路径

MKAN 是这个组合里最需要理解的部分。KAN 这几年火得快，但很多人直接拿公开实现塞进自己的模型里，效果参差不齐，多半是因为没弄明白基函数怎么选、残差结构怎么做、初始化尺度多重要。

2.1 一个贴近直觉的类比：固定弯钩 vs 可塑形抓手

如果你在仓库里搬运不同形状的箱子，方案一：准备很多种固定弯钩，遇到哪种箱子换哪种钩子，这是传统 MLP 加宽隐藏层的思路；方案二：准备几个“可塑形抓手”，每个都能调整自身弧度去适配不同箱子，这是单个 KAN 层的思路；方案三：准备一组不同材质的抓手——橡胶的不伤箱子，钢制的抓得紧，磁性吸盘吸铁皮箱——每个箱子上根据情况让不同抓手配合发力，这就是 MKAN 的多核融合思路。

落到数值上，KAN 层的输入是向量 x，输出是逐元素变换后的结果。每个输出通道对应的不是一组权重加激活，而是多组基函数 Φ(x; params) 的线性叠加。我的实现里，选了三个核成员：Bessel 核用于全局平滑拟合主力趋势，Chebyshev 核捕捉局部振荡细节，RBF 核负责填补短距离异常突变。三个核共享输入，各自展开到 hidden 维度，再经过一个可学习的加权融合层合并。

2.2 MKAN 的多核融合是怎么算的

下面是我在项目里实际用的一段精简实现，核函数部分用张量计算展开，便于理解和修改：

import torch import torch.nn as nn import torch.nn.functional as F import math # bessel 核：适合低频/平滑成分 def bessel_kernel(x, order=2): # 用截断的积分近似 bessel 函数，这里简化为带符号的幂次组合 return torch.sin(x * math.pi) * torch.exp(-torch.abs(x)) * (1 + 0.1 * x ** order) # chebyshev 核：适合振荡成分 def chebyshev_kernel(x, degree=3): t0 = torch.ones_like(x) t1 = x if degree == 0: return t0 if degree == 1: return t1 t2 = 2.0 * x * t1 - t0 return t2 # rbf 核：适合局部突变 def rbf_kernel(x, gamma=1.0): return torch.exp(-gamma * x ** 2) class MKANBlock(nn.Module): def __init__(self, d_model, hidden_dim, n_kernels=3): super().__init__() self.d_model = d_model self.hidden_dim = hidden_dim # 每个核单独有一组映射参数 self.kernel_projectors = nn.ModuleList([ nn.Linear(d_model, hidden_dim, bias=False) for _ in range(n_kernels) ]) # 融合权重，可学习 self.fusion = nn.Linear(hidden_dim * n_kernels, d_model) self.norm = nn.LayerNorm(d_model) def forward(self, x): # x: (batch, seq, d_model) k1 = bessel_kernel(self.kernel_projectors[0](x)) k2 = chebyshev_kernel(self.kernel_projectors[1](x)) k3 = rbf_kernel(self.kernel_projectors[2](x), gamma=0.5) out = torch.cat([k1, k2, k3], dim=-1) out = self.fusion(out) return self.norm(x + out)

这里有一个关键细节：每个核的前置投影层不共享参数，因为不同核需要的“输入尺度”其实是不同的，强制共享会影响每个核的响应特性。融合层把三路拼接后压回 d_model，再用残差连接加 LayerNorm，这是 KAN 类模块能稳定训练的重要保障。

2.3 嵌入 Timexer 的位置：替换前馈，保留注意力

替换时我保留了 Timexer 原有的 patch embedding 和注意力层，只用 MKANBlock 替换每个 Transformer Block 里的前馈网络。原因很简单：注意力的全局交互还是要靠标准 QKV 机制，这部分 Timexer 已经做得够好；真正拖后腿的是前馈层的局部特征变换。

替换之后参数量从约 3.2M 涨到 3.8M，涨幅 18% 左右，训练显存增加也在这个量级，但拟合能力提升很值。值得注意的是，第一层 Transformer Block 的前馈不建议直接换，因为网络刚开始接触的是 patch 化后的高频片段，MKAN 多核的强表达会被高频噪声干扰，或者反过来把噪声记忆下来。我的做法是：第一层保留原始 MLP，从第二层开始替换三层 MKANBlock，效果比全换更稳。

3. PGA 渐进式梯度聚合：不是学习率技巧，而是训练阶段的课程编排

组合完模型，第一版直接训练时，我遇到了比较典型的“多损失打架”问题。主损失是预测序列的 MSE，辅助损失包括 patch 重建损失和每个时间尺度上的预测损失。理论上信息很丰富，实际上前 20 个 epoch 主损失一直在震荡，下不去。我把每个损失的梯度范数打出来看了，辅助损失的梯度范数一度是主损失的 3 倍，这就不难解释为什么训练不稳定了。

PGA 就是对这个问题下的手。

3.1 PGA 的课程式权重设计

PGA 的核心是一个随时间变化的损失权重调度。我的实现总损失如下：

L_total = L_main + α(t) * L_patch + β(t) * L_aux

其中 α(t) 和 β(t) 不是固定值，而是训练步数 t 的函数。我测试过两种调度：

线性 ramp：α(t) = min(1, t / warmup_steps)

余弦 ramp：α(t) = 0.5 * (1 - cos(π * min(1, t / warmup_steps)))

余弦 ramp 更平滑，我最终采用它作为默认配置。warmup_steps 我设为总步数的 30%，前 30% 步数内两个辅助损失缓缓加到满权重，后 70% 步数保持全量训练。主损失的权重始终为 1，保证最核心的目标始终主导。

3.2 辅助损失选什么，直接决定 PGA 效果

PGA 不是随便拿几个损失加权就完事。辅助损失要选得“与预测目标相关、但包含不同视角的信息”。我最终用了两个：

patch 重建损失：把输入 patch 序列随机 mask 一部分，要求模型用上下文重建，这跟 MAE 预训练的思路类似，能让 MKANBlock 的局部表示更稳健。

多尺度预测损失：把目标序列降采样成粗、中、细三个尺度，每个尺度头各自产出一个预测，计算与对应目标的 MSE。粗尺度损失先被激活，因为它的梯度信号更平滑；细尺度损失后激活，避免梯度噪声。

用代码表达训练循环的调度逻辑，大概是这样：

def progressive_weight(step, total_steps, ramp_start=0.3): progress = min(1.0, step / (total_steps * ramp_start)) return 0.5 * (1 - math.cos(math.pi * progress)) for step, batch in enumerate(train_loader): x, target = batch pred_main, pred_patch, pred_multi = model(x) loss_main = criterion(pred_main, target) loss_patch = mae(pred_patch, target_patch) loss_multi = sum( criterion(pred_multi[i], target_multi[i]) for i in range(len(target_multi)) ) w_patch = progressive_weight(step, total_steps) w_multi = progressive_weight(step, total_steps) * 0.5 loss = loss_main + w_patch * loss_patch + w_multi * loss_multi optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step()

注意 multi-scale 损失的权重我额外乘了 0.5，因为三个尺度叠加后梯度范数天然偏大，不加抑制的话即便有 ramp，后期仍会压过主损失。

3.3 PGA 和 SAM、SWA 的区别

说到稳定训练，很多人会想到 SAM 或 SWA。我同样试过，简单说下区别。

SAM 是锐度感知优化器，它对“损失 landscape 的平坦度”做约束，计算开销大，每个 step 要多一次前向后向，对时序模型这种 batch 通常比较大的场景，训练时间直接翻倍，而收益不一定显著。

SWA 是对训练末尾的权重做平均，它能改善泛化，但对前期的收敛稳定性没有任何帮助。

PGA 的思路更接近“课程预训练”：调节的是不同训练阶段里梯度信号的组成比例。它的收益在前 1/3 步数就能体现出来，配合 AdamW 使用效果很顺。

4. 完整 PyTorch 实现：从数据滑窗到模型组装的落地代码

理论讲完，直接上可运行的代码骨架。我尽量保持精简，但保留所有关键结构和参数。

4.1 数据管线：滑窗 Dataset 和归一化

时序预测的数据管线，最容易出错的地方是归一化和滑窗对齐。我的做法是：训练集用整体均值方差归一化，验证集和测试集沿用训练集的统计量，不做增量归一化；滑窗长度由预测步长决定，比如输入 96 步预测 24 步。

class TimeSeriesDataset(torch.utils.data.Dataset): def __init__(self, X, y, in_len=96, out_len=24): self.X = torch.FloatTensor(X) self.y = torch.FloatTensor(y) self.in_len = in_len self.out_len = out_len def __len__(self): return len(self.X) - self.in_len - self.out_len + 1 def __getitem__(self, idx): x = self.X[idx: idx + self.in_len] y = self.y[idx + self.in_len: idx + self.in_len + self.out_len] return x, y

归一化时，注意每个变量单独做，不要统一归一化到同一个尺度，因为不同变量的物理含义不同，统一尺度会破坏变量间真实的相对关系。

4.2 Timexer 核心结构的 PyTorch 精简实现

我保留了三个关键组件：PatchEmbedding、自注意力层、内生-外生交叉注意力。代码上做了一点合并，实际使用中你可以按需拆分。

class PatchEmbedding(nn.Module): def __init__(self, in_var, d_model, patch_len, stride): super().__init__() self.patch_len = patch_len self.stride = stride self.proj = nn.Linear(patch_len, d_model) self.pos = nn.Parameter(torch.randn(1, 1000, d_model) * 0.02) def forward(self, x): # x: (batch, in_len, var) b, l, v = x.shape x = x.permute(0, 2, 1) # (b, var, l) # 滑动切片 patches = x.unfold(dimension=-1, size=self.patch_len, step=self.stride) # (b, var, num_patch, patch_len) patches = patches.reshape(b * v, -1, self.patch_len) out = self.proj(patches) + self.pos[:, :patches.shape[1], :] return out, b, v

自注意力和交叉注意力我用标准 MHA 实现，将内生变量的 patch 表示作为 query，外生变量 embedding 作为 key/value，这样模型能学到外生变量如何影响目标序列的特定片段。

4.3 组装 PGA+MKAN+Timexer 主模型

主模型的 forward 里除了输出主预测，还要输出 patch 重建结果和多尺度预测，方便训练时算辅助损失。我用一个额外的轻量 head 来输出这些辅助结果，推理时只取主预测。

class PGA_MKAN_Timexer(nn.Module): def __init__(self, in_var, out_var, d_model=128, n_heads=8, n_layers=3, patch_len=8, stride=8, out_len=24): super().__init__() self.patch_embed = PatchEmbedding(in_var, d_model, patch_len, stride) self.exo_embed = nn.Linear(in_var, d_model) self.blocks = nn.ModuleList() for i in range(n_layers): self.blocks.append(nn.ModuleDict({ "attention": nn.MultiheadAttention(d_model, n_heads, batch_first=True), "ffn": MKANBlock(d_model, d_model * 2) if i > 0 else nn.Sequential( nn.Linear(d_model, d_model * 2), nn.GELU(), nn.Linear(d_model * 2, d_model) ) })) self.head_main = nn.Sequential( nn.Linear(d_model, out_len), ) self.head_patch = nn.Sequential( nn.Linear(d_model, patch_len), ) self.head_multi = nn.ModuleList([ nn.Linear(d_model, out_len // 2**i) for i in range(3) ]) def forward(self, x): # x: (batch, in_len, vars) exo = self.exo_embed(x) # (batch, in_len, d_model) patches, b, v = self.patch_embed(x) # patches: (b*v, num_patch, d_model) h = patches for blk in self.blocks: attn_out, _ = blk["attention"](h, exo.repeat(v, 1, 1), exo.repeat(v, 1, 1)) h = blk["ffn"](attn_out + h) # 聚合预测 pred_main = self.head_main(h.mean(dim=1)) # (b*v, out_len) pred_main = pred_main.reshape(b, v, -1).mean(dim=1) pred_patch = self.head_patch(h[:, -1, :]) # 用最后一个 patch pred_multi = [head(h.mean(dim=1)).reshape(b, v, -1).mean(dim=1) for head in self.head_multi] return pred_main, pred_patch, pred_multi

注意代码里我把内生变量 patch 化后展平到 batch 维（b*v），这样对多个变量可以同时处理。推理时对多个变量取平均，实际项目里如果你只需要预测单一目标变量，可以让 v=1，不需要平均。

4.4 训练配置与关键参数表

以下是我最终使用的稳定配置，不同数据集可能需要微调，但大体范围适用：

这个配置跑出来的效果，在后面的实测部分会给出具体数字。

5. 跑分与实测：三组数据的收益、开销与收敛行为对比

代码跑通只是第一步，真正让人信服的结果是对比数据。我在三组不同特性的数据集上做了测试，模拟的是我实际工作中遇到的场景，每组数据都做了 Timexer 基线、Timexer+MKAN、PGA+MKAN+Timexer 三种配置的对比。

5.1 三组数据的特点与对比结果

第一组是某工业设备传感器监测数据，变量间耦合性强，信号变化平滑但有周期性；第二组是城市交通流量数据，有明显的早晚高峰突变；第三组是气象站逐小时温度数据，外生变量影响大（太阳辐射、风速等），序列本身低频为主。

实验结果如下：

结论很清楚：MKAN 本身能带来比较稳定的 5% 到 10% 的指标改善，PGA 主要在设备传感器和城市交通流这两组波动较强的数据上带来额外收益。在气象站这种本身平滑的数据上，PGA 提升有限，说明它解决的问题确实主要是“训练稳定性”而不是“模型容量”。

5.2 收敛行为上的区别

除了最终指标，收敛曲线差异更有说服力。Timexer 基线在设备传感器数据上约 25 个 epoch 后基本平台期；Timexer+MKAN 在前期震荡明显增大，要到 30 个 epoch 左右才稳住；加了 PGA 之后，前 10 个 epoch 主损失下降平稳，没有 MKAN 单独使用时的“大坑”，到 20 个 epoch 左右就已经达到基线的平台水平，后半程稳步继续往下降。

我抽样打印过梯度范数，发现 PGA 组里主损失相关的梯度占比始终保持在 70% 以上，而对照组里有时会掉到 40% 以下。这说明 PGA 确实保证了“主要矛盾一直在主要地位”。

5.3 参数量和推理速度，别忽视的工程指标

模型参数量从 3.2M 提升到 3.8M，看起来不多，但每个 MKANBlock 的 FLOPs 比原 MLP 高不少，因为三路核函数并联计算。实测推理速度降低约 20%，但考虑到时序预测模型通常不需要实时响应到毫秒级，这个代价是可接受的。如果部署环境对延迟敏感，可以考虑只替换最后一层前馈为 MKAN，收益会打折扣，但推理开销控制在 5% 以内。

这个取舍建议根据你自己的场景决定。我一般会给两套配置：完整版用于离线批量预测，轻量版用于需要实时推断的边缘部署。

6. 调参实录：数值稳定性、显存占用与复现性问题的处理

最后这部分全是实际操作中踩过的坑。我按问题类别列出来，每个都是我实打实处理过的，希望帮你少走弯路。

6.1 数值稳定性：MKAN 的核函数可能让 loss 直接变 NaN

第一个大坑是 NaN。Bessel 核函数里我一开始写了 sin(x * π) * exp(-|x|)，在输入尺度较大时没问题，但 MKANBlock 里的投影层如果初始化不当，输出范围可能瞬间冲到几十上百，exp(-|x|) 把梯度压得极小，sin(x * π) 的振荡又让梯度方向变得混乱，最终表现为 loss 不降反升甚至 NaN。

解决办法有两层：输入层之前加一层 LayerNorm 是必须的；初始化的时候，kernel_projectors 的权重 scale 设为 0.1 而不是默认的 1/sqrt(fan_in)，实测能让前期训练稳定很多。同一组代码，scale 从 0.1 改成默认值，马上一半的随机种子跑出 NaN，这问题非常隐蔽。

6.2 显存占用居高不下，怎么压

带 MKAN 的序列建模，显存主要爆在中间激活值上。patch 化之后的序列长度虽然变短了，但三路核函数会各自产生一份 hidden 大小的中间张量。input_len=96、batch_size=128、d_model=128 的时候，单层 MKANBlock 的中间激活大约是 128 * 32 * 128 * 5 * 3 个 float，量级不低。

我的建议：第一，能用 stride=patch_len 就不重叠切分，序列长度少接近一半；第二，开启 PyTorch 的 gradient checkpointing，只对 MKANBlock 开启，不必对整个模型开启，因为注意力部分显存本来就不大；第三，多尺度预测的辅助 head 在训练最后 30% 步数可以关掉，减少反向传播时的内存峰值。

6.3 随机种子与可复现性：PGA 的权重调度不能是“随机的”

可复现性也是重点。别只看 torch.manual_seed，PGA 的权重调度函数如果依赖 step 计数，而 step 计数在不同分布式训练配置下可能不对齐，导致同样的 seed 在不同设备数下结果不同。

我的做法是把总步数算好之后，把每个 step 对应的 α(t) 和 β(t) 预先算成一个数组，作为常量传给训练循环，而不是在循环里实时计算。这样不管是用单卡还是多卡，线性还是分布式采样，结果都能对齐。

6.4 一个额外的终点微调技巧，能让指标再降一点点

在 PGA 全量训练结束之后，把 α(t) 和 β(t) 固定为 1.0，再用很小的学习率（1e-5 左右）继续训练 20 到 30 个 epoch。这个阶段相当于让模型在全部损失都参与的情况下做最后的“精修”。我在三组数据集上都试过，指标还能再降 1% 到 3%，而且没有出现过度拟合验证集的迹象。你可以当成一个免费的午餐来用。

顺带提一句，如果你只是想快速验证 MKAN 在你自己的数据上有没有用，可以先不改训练策略，只做“替换 MLP 为 MKAN”加“把学习率降到原来的 0.5 倍”这两件事。如果指标有明显提升，再上 PGA 做细调；如果没提升，问题多半不在模型容量，而在数据处理或者特征工程上，PGA 也救不回来。