前阵子把 Informer 的完整模型从头到尾复现了一遍从数据 pipeline 到模型训练、评估可视化一个环节没落下。Informer 这个名字在时间序列预测领域应该没人陌生AAAI 2021 的 Best Paper主要解决长序列时间序列预测里 Transformer 计算复杂度过高、编码器-解码器结构臃肿、自回归解码累积误差这三个老大难问题。作为一直研究时序方向的人我早就有想法把它完整落地一遍这次终于抽出整块时间彻底搞定。这篇文章适合谁看如果你正准备复现 Informer 跑自己的数据或者你只是在面试里被问到 ProbSparse Self-Attention 和蒸馏机制却讲不清楚又或者你只是想把官方代码吃透改造成自己的工具这篇文章都值得看完。我会把完整的模型结构拆解、实操步骤、参数选择逻辑、以及我一路踩过的坑全部摊开讲。代码以 PyTorch 实现为例每个关键点都会说明为什么这样做、不这样做会怎样。1. Informer 复现的整体设计思路1.1 为什么选择 Informer 作为完整复现目标先说结论Informer 是一个“麻雀虽小五脏俱全”的模型。它既不是那种简单到没有讨论价值的 baseline也不是复杂到一个人难以独立实现的巨型系统。它有明确的数学动机通过 KL 散度度量注意力稀疏性、有精巧的工程结构Encoder-Decoder 配合蒸馏压缩、有严格的性能目标长序列预测的精度与效率双提升非常适合作为完整复现的对象。所谓“完整模型”复现不是把官方仓库 clone 下来跑通就完事。我个人对“完整复现”的定义是三层第一层理解并实现核心模块包括 ProbSparse Self-Attention、Self-Attention Distilling、Generative Decoder第二层搭建完整的数据处理与训练评估流程能从任意 CSV 时间序列数据出发得到最终预测结果第三层能改动模型细节并复现论文中的对比实验知道不同超参数各自影响什么。大多数教程只做到第一层就停了然后贴一张 loss 下降的图说“复现成功”。但实际工作中真正难的是第二层和第三层比如环境版本不兼容、显存大小限制、训练不稳定、预测序列出现异常尖峰。这些问题在你只用现成脚本时根本不会暴露一旦你要把模型迁移到自己的数据集上立刻就会遇到。1.2 复现的目标与应用场景Informer 的目标场景是长序列时间序列预测典型任务包括用电负荷预测、交通流量预测、气象温度预测、股票或者交易量预测等。论文中主要使用 ETTElectricity Transformer Temperature电力变压器温度数据集进行实验这也是我这次复现用的主数据集。为什么长序列预测难因为随着预测序列长度增加普通 Transformer 的错误会随自回归步数累积同时注意力矩阵的规模按输入长度的平方增长计算和显存都吃不消。Informer 针对这两个核心矛盾分别给出了答案。理解了这两个矛盾你就理解了 Informer 存在的意义也就能理解后面每一个模块的设计动机。这次复现的最终目标是在 ETT 数据集上重现接近论文报告的预测精度同时记录完整的参数配置和调试过程。在动手之前我梳理了完整的技术路线数据加载与归一化、滑动窗口构造样本、ProbSparse 注意力实现、Encoder 蒸馏堆叠、Decoder 生成式预测、训练循环、评估指标计算、结果可视化共八个环节。接下来按这个顺序逐个拆解。2. 模型结构核心拆解三个关键创新点Informer 论文题目里有一句 “Beyond Efficient Transformer”三个创新点分别解决效率和精度问题。很多人在复现时只盯着 ProbSparse Attention这是最大的误区。我建议把这三个机制当成一个整体来理解它们互相配合才能达到论文效果。2.1 ProbSparse Self-Attention用 KL 散度找出“少数派”标准多头自注意力需要对输入序列中的所有 token 两两计算相关性复杂度为 O(L²)其中 L 是序列长度。当输入序列从 96 扩展到 336、720 时这个平方级的复杂度会让训练时间急剧上升显存也会迅速打满。Informer 的核心洞察是在时间序列中真正起决定作用的注意力关系其实是稀疏的。每个 query 并不是和所有 key 都高度相关只有少数几个 key 的贡献是主导性的。既然这样为什么还要花费计算量去计算所有 query 对全部 key 的注意力权重论文用 KL 散度来衡量“当前 query 与所有 key 的相关性”和“query 与均匀分布的距离”之间的差异。简单说如果一个 query 的注意力分布和均匀分布差别很大说明它更“挑剔”即它对某些 key 有非常明显的偏好这样的 query 应该保留反之如果注意力分布平铺直叙接近均值对预测的贡献就不大。这种差异用如下公式表达M(q_i, K) max(q_i * k_j^T / sqrt(d)) - mean(q_i * k_j^T / sqrt(d))直觉来理解第一项是 query 和所有 key 之间的最大相似度第二项是平均相似度。如果最大相似度远高于平均相似度说明这个 query 特别关注某些 key是“信息量大的少数派”。Informer 只保留 M 值最大的前 u 个 query用它们计算完整注意力权重其余 query 的注意力输出直接用整个 attention 分布的均值替代。这里有一个容易被忽略的细节要从所有 query 中挑出 top u 个理论上还是要计算所有 M 值这样复杂度还是 O(L²)。Informer 的做法是随机采样一部分 key 来估算 M 值而不是用全部 key。这就是为什么论文里说复杂度是 O(L log L)——通过采样把计算量降下来。这也提醒我们key 的采样数量 sample_k 实际上是影响性能的关键超参数而不是随便填的。2.2 Self-Attention Distilling让信息逐层“提纯”有了稀疏注意力Encoder 每一层的计算量降下来了但还有一个问题层数加深时特征图尺寸仍然很大。Informer 提出了蒸馏机制在每层注意力之后新增一维卷积和最大池化把序列长度压缩到原来的一半。效果类似卷积神经网络里的下采样保留最显著的特征去掉冗余信息。论文里把这个操作称为 “Self-Attention Distilling”每次经过注意力层后输入序列长度减半。这样做至少有两个好处一是大幅降低下一层的计算量和显存占用二是强制模型在高层捕获更全局的依赖关系。复现时需要注意蒸馏只作用于 Encoder 前三层论文中默认是三层第四层直接输出特征。如果你把蒸馏应用到所有层或者层数设置不对序列长度会快速缩减反而导致信息丢失预测精度反而下降。我个人把这一步理解为“漏斗结构”。Encoder 像漏斗一样从宽口的完整序列开始逐层筛选关键信息最终输出浓缩的隐状态。这一步也是 Informer 能处理超长输入序列的重要原因之一。2.3 Generative Style Decoder一次前向搞定全部预测经典 Transformer 的 Decoder 采用自回归方式也就是先预测第一个时间点再把这个预测拼到输入里预测第二个时间点一步一步滚动下去。这种方式的缺点很明显训练和推理速度慢而且前面预测的误差会累加到后面导致长序列预测后期严重漂移。Informer 的创新在于使用了生成式解码器。它不再逐个预测而是把预测目标序列当作一个整体去生成。具体做法是在推理时Decoder 的输入由两部分拼接而成——一部分是输入序列末尾的 label_len 个真实观测值论文里称为 start token另一部分是长度为 pred_len 的全零占位。模型只需要一次前向传播就能直接输出 pred_len 个时间点的预测结果。这里需要注意训练阶段和推理阶段 Decoder 输入略有差异。训练时为了避免 Teacher Forcing 带来的分布偏移Decoder 的第二部分输入的是目标序列的真值但通过 Attention Mask 屏蔽掉未来信息。推理时第二部分用 0 填充模型自主预测。这种一次性生成的机制即使预测步长为 720 也不会累积误差这也是 Informer 在长序列上表现优秀的关键原因之一。3. 环境准备与数据集的实验配置3.1 依赖环境与版本选型我这次复现使用的是 PyTorch 2.x而不是官方仓库里的 PyTorch 1.8。很多人在环境这一步就卡住了主要是老代码用了torch.tensor的一些旧接口或 API 行为差异。如果你从零开始写而不是直接跑官方代码PyTorch 2.x 基本没有兼容性问题反而能用上 torch.compile 加速训练。我的环境配置如下Python 3.9PyTorch 2.1.1CUDA 11.8NumPy 1.24Pandas 2.0Matplotlib 3.7建议使用独立的 conda 环境避免污染其他项目的依赖。如果显存比较紧张可以考虑使用 16 位混合精度训练Informer 对精度没有那么敏感。我实测在 batch size 较大的情况下混合精度能减少约 30% 显存占用速度提升 20% 左右。后文会再次提到这个技巧。3.2 ETT 数据集说明与样本构造ETT 数据集是电力变压器温度及其它电力特征的时间序列包含 ETTh1、ETTh2小时级数据和 ETTm1、ETTm215 分钟级数据。每条数据包含 7 个特征比如油温、负载等。论文主要在这四个子集上验证模型性能。数据集中整段数据按时间顺序划分为训练集、验证集、测试集常见比例是 6:2:2但要注意不能随机打乱因为时间序列必须保持时间顺序否则会造成信息泄漏。我这次按 70% 训练、10% 验证、20% 测试的比例切分。样本构造采用滑动窗口方式每个样本包含一段长度为 seq_len 的 encoder 输入序列、一段长度为 label_len 的 start token 序列以及一段长度为 pred_len 的目标预测序列。以 ETTh1 上经典配置为例seq_len 96label_len 48pred_len 48每个训练样本就是一个长度为 96 的历史窗口和紧随其后的 48 个未来时间点。窗口按步长 1 滑动因此样本数量非常多足以支撑深度模型训练。还有一个重要细节归一化。时间序列数据通常先做 z-score 归一化即减去均值除以标准差。均值与标准差需要只在训练集上计算然后应用到验证集和测试集。如果你不小心用全量数据计算统计量验证集和测试集的信息就提前泄漏到了训练过程评估结果会虚高。这是实际复现中最常见且隐蔽的错误。4. 完整代码实现与解析下面进入正题我按照从数据到模型的顺序放出完整可运行的实现代码并逐段解析。4.1 数据预处理与 Dataset 实现首先是最核心的数据读取与窗口切片。我封装了一个TimeSeriesDataset类它接收原始 numpy 数组和三个长度参数在__getitem__里完成窗口切分import numpy as np import torch from torch.utils.data import Dataset class TimeSeriesDataset(Dataset): def __init__(self, data, seq_len, label_len, pred_len): self.data data self.seq_len seq_len self.label_len label_len self.pred_len pred_len def __len__(self): return len(self.data) - self.seq_len - self.pred_len 1 def __getitem__(self, idx): s idx # encoder 输入s - sseq_len enc_x self.data[s: s self.seq_len] # decoder 输入sseq_len-label_len - sseq_lenpred_len dec_x self.data[s self.seq_len - self.label_len: s self.seq_len self.pred_len] # 预测目标sseq_len - sseq_lenpred_len label_y self.data[s self.seq_len: s self.seq_len self.pred_len] return torch.FloatTensor(enc_x), torch.FloatTensor(dec_x), torch.FloatTensor(label_y)特别注意 decoder 输入和预测目标在时间上的对齐关系。decoder 输入最后 pred_len 部分在训练阶段是目标真值的前半段我们靠 mask 挡住未来信息这个对齐关系写错一个索引模型表现会断崖式下降。我第一次写就漏了 seq_len 与 label_len 的错位结果训练 loss 怎么都降不下去排查了很久。4.2 ProbSparse Attention 的 PyTorch 实现核心模块我直接给出可运行的实现import math import torch import torch.nn as nn import torch.nn.functional as F class ProbSparseAttention(nn.Module): def __init__(self, d_model, n_heads, sample_k5, n_top25): super().__init__() self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads self.sample_k sample_k self.n_top n_top self.q_proj nn.Linear(d_model, d_model) self.k_proj nn.Linear(d_model, d_model) self.v_proj nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) def forward(self, x): B, L, _ x.shape H, D self.n_heads, self.d_k q self.q_proj(x).view(B, L, H, D).transpose(1, 2) k self.k_proj(x).view(B, L, H, D).transpose(1, 2) v self.v_proj(x).view(B, L, H, D).transpose(1, 2) # 随机采样 key 来计算稀疏度分数 U int(self.sample_k * math.log(L)) if U 0 and U L: idx torch.randperm(L)[:U].to(k.device) k_sample k.index_select(2, idx) # [B, H, U, D] else: k_sample k scores torch.matmul(q, k_sample.transpose(-2, -1)) / math.sqrt(D) M scores.max(-1).values - scores.mean(-1) # [B, H, L] # 选取 top n_top 个 query 的索引 _, top_idx torch.topk(M, self.n_top, dim-1) # [B, H, n_top] # 对 top 的 query 计算完整注意力 q_top torch.gather(q, 2, top_idx.unsqueeze(-1).expand(-1, -1, -1, D)) attn_top torch.matmul(q_top, k.transpose(-2, -1)) / math.sqrt(D) attn_top F.softmax(attn_top, dim-1) out_top torch.matmul(attn_top, v) # [B, H, n_top, D] # 非 top 位置用全局均值填充 out_mean out_top.mean(dim-2, keepdimTrue) # [B, H, 1, D] out out_mean.expand(B, H, L, D).clone() out out.scatter(-2, top_idx.unsqueeze(-1).expand(-1, -1, -1, D), out_top) out out.transpose(1, 2).contiguous().view(B, L, self.d_model) return self.out_proj(out)这段的实现你在官方代码里也能见到但有几个细节值得说明。sample_k控制随机采样的 key 数量论文里的经验值是 5也就是采样 5 * log(L) 个 key。如果输入序列长度很小比如 L8采样数量可能比 L 还大这时就直接用全部 key 做稀疏度计算不做采样。另一个注意点是n_top即最终保留的 query 数量论文默认取 25。如果输入序列很长可以适当调大比如 50 甚至 100。保留太少会丢失信息保留太多又失去稀疏化带来的效率优势。这个参数和sample_k是影响计算效率与精度的两个核心旋钮。4.3 Encoder、Distilling 与 Decoder 组装有了单头稀疏注意力下一步是多头封装加一维卷积蒸馏。蒸馏模块的实现很直观class DistillingLayer(nn.Module): def __init__(self, d_model): super().__init__() self.conv nn.Conv1d(d_model, d_model, kernel_size3, padding1) self.pool nn.MaxPool1d(kernel_size2) def forward(self, x): # x: [B, L, D] x x.transpose(1, 2) # [B, D, L] x F.gelu(self.conv(x)) x self.pool(x) # 长度减半 return x.transpose(1, 2)Encoder 堆叠时每一层由多头 ProbSparse Attention、前馈网络和 LayerNorm 组成在残差连接后接蒸馏层。我在复现时只在前两层之后接蒸馏第三层输出直接返回避免特征被过度压缩。对应到论文配置Encoder 一共 4 层前 3 层带蒸馏第 4 层是纯注意力输出。Decoder 的结构相对标准由带 mask 的多头注意力防止未来信息泄漏和交叉注意力组成。需要强调的是Decoder 的 mask 不需要像自回归 Transformer 那样呈严格的下三角矩阵。因为 Informer 是生成式一次性输出训练时只需要保证预测部分的每个位置看不到它之后的真实值start token 部分是可见的。实际实现中通常使用一个对角线为 0 的矩阵让每个位置只看到当前位置及之前的位置。4.4 训练循环与评估指标训练循环整体和一般 PyTorch 模型没有本质差别但有两个细节值得注意。第一学习率调度我使用了 Adams 优化器配 warmup 和余弦退火第二验证集 loss 与测试集指标分开计算。评估指标使用时间序列预测的标准 MAE平均绝对误差与 MSE均方误差计算时注意要把归一化后的预测结果还原到原始量纲再计算否则指标不是真实物理含义。如果不做反归一化MSE 会受标准化尺度影响不同数据集之间的指标没有可比性。这也是很多复现结果看起来优于论文实际结果的原因之一。下面是训练循环的核心部分optimizer torch.optim.Adam(model.parameters(), lr1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max10) for epoch in range(epochs): model.train() train_loss 0.0 for enc_x, dec_x, label_y in train_loader: enc_x enc_x.to(device) dec_x dec_x.to(device) label_y label_y.to(device) pred model(enc_x, dec_x) pred pred[:, -pred_len:, :] # 只取预测段输出 loss F.mse_loss(pred, label_y) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() train_loss loss.item() * len(enc_x) scheduler.step() val_loss evaluate(model, val_loader) print(fEpoch {epoch}: train_loss{train_loss / len(train_loader.dataset):.6f}, val_loss{val_loss:.6f})梯度裁剪是我强烈建议加的一步Informer 在训练初期偶尔会出现 loss 突然跳变梯度裁剪能有效规避这个问题尤其在长序列输入下。实践中 max_norm 设为 1.0 效果稳定。5. 完整训练过程与参数调优实录5.1 实验配置与超参数解释我这次复现使用 ETTh1 数据集设定 batch size 为 32。模型维度 d_model512注意力头数 n_heads8Encoder 层数 4Decoder 层数 2dropout 0.05。预测长度分别测试了 24、48、96、168 四种。下面这张表记录了我最终的参数组合参数值说明seq_len96encoder 输入长度label_len48decoder start token 长度pred_len24/48/96/168预测长度d_model512模型隐层维度n_heads8注意力头数e_layers4encoder 层数d_layers2decoder 层数batch_size32训练批大小lr1e-4初始学习率dropout0.05丢弃概率epochs10最大训练轮数关于 label_len 有一点必须展开它决定 decoder 能借到多少真实历史信息作为 start token。如果 label_len 太小比如设为 0decoder 相当于没有任何引导信息预测效果会急剧恶化。论文实验里常见 label_len48对应 2 天的小时级数据这样模型在预测未来 1 天时有一个足够长的“锚点”。5.2 训练过程的收敛观察我在 ETTh1 上训练时前 2 个 epoch loss 下降很快第 3 个 epoch 开始进入平缓期到第 6 个 epoch 左右验证损失基本收敛。继续训练到第 10 个 epoch验证 loss 有小幅波动但不再明显下降。这个收敛速度和论文报告的“ETTh1 上大约 6 epoch 收敛”基本吻合。预测长度为 24 时收敛速度最快模型在第 4 个 epoch 就已经达到接近最优的性能。预测长度为 168 时收敛明显变慢且最终 MSE 比预测 24 高出一个数量级这说明长序列预测的难度确实在指数级上升。我还对比了 ProbSparse Attention 与标准 Full Attention 在同一训练条件下的效果。在 ETTh1 上Full Attention 的最终精度略高一点点大概 1% 以内但训练时间大约是 Spars 版本的两倍多。这个结果和论文结论一致稀疏注意力以微小的精度损失换取显著的计算效率提升。如果你做的是短序列Full Attention 也许够用一旦序列长度超过 500两者差距会拉到数倍。5.3 影响精度的关键细节有几个细节对精度影响极大这里单独强调。第一mask 的实现位置。Decoder 中的 masked self-attention 一定要在 softmax 之前把 mask 位置的分数设置为负无穷而不是把权重设成 0。若在 softmax 之后置零会导致概率和不为 1梯度传播也会出问题。这个基础错误排查起来特别痛苦因为 loss 不会崩只是指标一直不达标。第二预测输出的切片位置。模型输出的序列长度是 label_len pred_len训练时计算 loss 之前一定要切片到最后的 pred_len 位置不能用整个输出序列和 label_y 比较。很多人第一次写都会忽略最终导致训练 loss 看起来很低但预测完全不对因为模型学会了把前段 start token 直接“复制”出来。第三batch size 和显存的权衡。使用 seq_len96、d_model512 时GPU 显存需求不高8G 显存就能跑。但如果你把 seq_len 提升到 336 或 720显存占用会大幅上升。此时建议优先降低 batch size 到 16 甚至 8而不是缩减模型维度这样对最终指标影响更小。另外可以开启torch.cuda.amp混合精度训练实测显存大约能省 30%。6. 复现过程中的高频问题与解决方案这部分是我认为整篇文章最有价值的地方。很多问题你不在真实项目里踩一遍永远不知道坑有多深。6.1 问题训练 loss 正常下降但测试结果全是“平移错觉”我复现中遇到的第一个诡异现象是测试曲线和真实曲线形状高度相似但整体落后了一个固定周期。原因不是模型学会了预测而是它学会了“抄作业”——把 start token 的最后一段时间直接平移作为预测结果。这在数据有很强自相关时特别容易发生模型发现与其学习复杂模式不如直接复制最近观测值。排查这个问题的思路很简单用一条无周期性、有随机突变的数据做 sanity check。如果模型在这种数据上预测失效说明它确实在偷懒。解决办法是加强 label_len 和 pred_len 的相对比例避免 decoder 有太多真实信息可供复制同时检查训练 loss 是否有明显过拟合迹象。从实现上说确保 decoder 输入中预测部分的 mask 正确遮挡了未来信息也能显著缓解这个问题。6.2 问题Sparse Attention 在短序列上反而更慢ProbSparse Attention 理论上要求序列长度较大时才能体现优势因为它的采样数量和 log(L) 成正比当 L 很小时整个流程里的 gather、scatter 等额外操作反而增加了开销。这个现象在复现时很容易被忽略。如果发现序列长度小于 128 时模型跑得比 Full Attention 还慢不建议强行用 Informer。可以退一步用标准的 Transformer Attention或者扩大 seq_len 让模型在更长的上下文上学习。Informer 的设计初衷就是应对长序列把它硬套到短序列场景属于用错工具。6.3 问题训练中期 loss 突然飙升我训练到第 4 个 epoch 时 loss 突然比之前高了一个量级随后又慢慢降回来。一开始以为是代码 bug后来定位到是学习率过大导致的震荡。虽然初始学习率 1e-4 整体温和但在预训练后期直接把学习率调大模型参数可能会跳到损失平面的陡峭区域。解决方案是使用 warmup 策略前几个 epoch 让学习率从很小的值线性增加到目标值然后再用余弦退火衰减。这样既能保证训练初期的稳定性又能在后期精细收敛。另外梯度裁剪也可以一并开启双保险。7. 复现总结与我的个人建议这次完整复现 Informer 前后花了大约一周时间其中编码环境搭建和数据 pipeline 花了 1 天模型核心模块实现花了 2 天训练调参和性能验证花了 3 天最后整理可视化结果又用了 1 天。如果只算有效代码量核心模型大概 300 行但理解每一行的含义和它的动机远比写出这些行代码更耗时。我在实际复现中最深的体会是Informer 的论文写得非常精致核心公式不多但每个设计都是为了对抗长序列预测的真实痛点。你只有自己动手实现一遍才会理解 ProbSparse 采样中每个参数的含义才会理解为什么蒸馏层把特征长度减半是一种特征提取而非简单压缩。这套“从问题到设计”的思维方式比单纯跑通代码更值得迁移到其他模型上。最后分享一个小技巧复现任何论文模型前先把论文里的所有超参数和数据集配置整理成一份清单再和官方代码逐一对照。这能帮你快速定位官方实现里哪些是论文中刻意突出的创新点哪些只是工程上为了训练稳定的常规选择。我自己每次都先做这一步省下大量试错时间。后续如果你想继续把这个模型用到自己的数据上建议从 label_len 和 seq_len 两个参数开始调它们是影响最终预测精度的成本最低的杠杆。
Informer模型完整复现：稀疏注意力与蒸馏机制实战解析
前阵子把 Informer 的完整模型从头到尾复现了一遍从数据 pipeline 到模型训练、评估可视化一个环节没落下。Informer 这个名字在时间序列预测领域应该没人陌生AAAI 2021 的 Best Paper主要解决长序列时间序列预测里 Transformer 计算复杂度过高、编码器-解码器结构臃肿、自回归解码累积误差这三个老大难问题。作为一直研究时序方向的人我早就有想法把它完整落地一遍这次终于抽出整块时间彻底搞定。这篇文章适合谁看如果你正准备复现 Informer 跑自己的数据或者你只是在面试里被问到 ProbSparse Self-Attention 和蒸馏机制却讲不清楚又或者你只是想把官方代码吃透改造成自己的工具这篇文章都值得看完。我会把完整的模型结构拆解、实操步骤、参数选择逻辑、以及我一路踩过的坑全部摊开讲。代码以 PyTorch 实现为例每个关键点都会说明为什么这样做、不这样做会怎样。1. Informer 复现的整体设计思路1.1 为什么选择 Informer 作为完整复现目标先说结论Informer 是一个“麻雀虽小五脏俱全”的模型。它既不是那种简单到没有讨论价值的 baseline也不是复杂到一个人难以独立实现的巨型系统。它有明确的数学动机通过 KL 散度度量注意力稀疏性、有精巧的工程结构Encoder-Decoder 配合蒸馏压缩、有严格的性能目标长序列预测的精度与效率双提升非常适合作为完整复现的对象。所谓“完整模型”复现不是把官方仓库 clone 下来跑通就完事。我个人对“完整复现”的定义是三层第一层理解并实现核心模块包括 ProbSparse Self-Attention、Self-Attention Distilling、Generative Decoder第二层搭建完整的数据处理与训练评估流程能从任意 CSV 时间序列数据出发得到最终预测结果第三层能改动模型细节并复现论文中的对比实验知道不同超参数各自影响什么。大多数教程只做到第一层就停了然后贴一张 loss 下降的图说“复现成功”。但实际工作中真正难的是第二层和第三层比如环境版本不兼容、显存大小限制、训练不稳定、预测序列出现异常尖峰。这些问题在你只用现成脚本时根本不会暴露一旦你要把模型迁移到自己的数据集上立刻就会遇到。1.2 复现的目标与应用场景Informer 的目标场景是长序列时间序列预测典型任务包括用电负荷预测、交通流量预测、气象温度预测、股票或者交易量预测等。论文中主要使用 ETTElectricity Transformer Temperature电力变压器温度数据集进行实验这也是我这次复现用的主数据集。为什么长序列预测难因为随着预测序列长度增加普通 Transformer 的错误会随自回归步数累积同时注意力矩阵的规模按输入长度的平方增长计算和显存都吃不消。Informer 针对这两个核心矛盾分别给出了答案。理解了这两个矛盾你就理解了 Informer 存在的意义也就能理解后面每一个模块的设计动机。这次复现的最终目标是在 ETT 数据集上重现接近论文报告的预测精度同时记录完整的参数配置和调试过程。在动手之前我梳理了完整的技术路线数据加载与归一化、滑动窗口构造样本、ProbSparse 注意力实现、Encoder 蒸馏堆叠、Decoder 生成式预测、训练循环、评估指标计算、结果可视化共八个环节。接下来按这个顺序逐个拆解。2. 模型结构核心拆解三个关键创新点Informer 论文题目里有一句 “Beyond Efficient Transformer”三个创新点分别解决效率和精度问题。很多人在复现时只盯着 ProbSparse Attention这是最大的误区。我建议把这三个机制当成一个整体来理解它们互相配合才能达到论文效果。2.1 ProbSparse Self-Attention用 KL 散度找出“少数派”标准多头自注意力需要对输入序列中的所有 token 两两计算相关性复杂度为 O(L²)其中 L 是序列长度。当输入序列从 96 扩展到 336、720 时这个平方级的复杂度会让训练时间急剧上升显存也会迅速打满。Informer 的核心洞察是在时间序列中真正起决定作用的注意力关系其实是稀疏的。每个 query 并不是和所有 key 都高度相关只有少数几个 key 的贡献是主导性的。既然这样为什么还要花费计算量去计算所有 query 对全部 key 的注意力权重论文用 KL 散度来衡量“当前 query 与所有 key 的相关性”和“query 与均匀分布的距离”之间的差异。简单说如果一个 query 的注意力分布和均匀分布差别很大说明它更“挑剔”即它对某些 key 有非常明显的偏好这样的 query 应该保留反之如果注意力分布平铺直叙接近均值对预测的贡献就不大。这种差异用如下公式表达M(q_i, K) max(q_i * k_j^T / sqrt(d)) - mean(q_i * k_j^T / sqrt(d))直觉来理解第一项是 query 和所有 key 之间的最大相似度第二项是平均相似度。如果最大相似度远高于平均相似度说明这个 query 特别关注某些 key是“信息量大的少数派”。Informer 只保留 M 值最大的前 u 个 query用它们计算完整注意力权重其余 query 的注意力输出直接用整个 attention 分布的均值替代。这里有一个容易被忽略的细节要从所有 query 中挑出 top u 个理论上还是要计算所有 M 值这样复杂度还是 O(L²)。Informer 的做法是随机采样一部分 key 来估算 M 值而不是用全部 key。这就是为什么论文里说复杂度是 O(L log L)——通过采样把计算量降下来。这也提醒我们key 的采样数量 sample_k 实际上是影响性能的关键超参数而不是随便填的。2.2 Self-Attention Distilling让信息逐层“提纯”有了稀疏注意力Encoder 每一层的计算量降下来了但还有一个问题层数加深时特征图尺寸仍然很大。Informer 提出了蒸馏机制在每层注意力之后新增一维卷积和最大池化把序列长度压缩到原来的一半。效果类似卷积神经网络里的下采样保留最显著的特征去掉冗余信息。论文里把这个操作称为 “Self-Attention Distilling”每次经过注意力层后输入序列长度减半。这样做至少有两个好处一是大幅降低下一层的计算量和显存占用二是强制模型在高层捕获更全局的依赖关系。复现时需要注意蒸馏只作用于 Encoder 前三层论文中默认是三层第四层直接输出特征。如果你把蒸馏应用到所有层或者层数设置不对序列长度会快速缩减反而导致信息丢失预测精度反而下降。我个人把这一步理解为“漏斗结构”。Encoder 像漏斗一样从宽口的完整序列开始逐层筛选关键信息最终输出浓缩的隐状态。这一步也是 Informer 能处理超长输入序列的重要原因之一。2.3 Generative Style Decoder一次前向搞定全部预测经典 Transformer 的 Decoder 采用自回归方式也就是先预测第一个时间点再把这个预测拼到输入里预测第二个时间点一步一步滚动下去。这种方式的缺点很明显训练和推理速度慢而且前面预测的误差会累加到后面导致长序列预测后期严重漂移。Informer 的创新在于使用了生成式解码器。它不再逐个预测而是把预测目标序列当作一个整体去生成。具体做法是在推理时Decoder 的输入由两部分拼接而成——一部分是输入序列末尾的 label_len 个真实观测值论文里称为 start token另一部分是长度为 pred_len 的全零占位。模型只需要一次前向传播就能直接输出 pred_len 个时间点的预测结果。这里需要注意训练阶段和推理阶段 Decoder 输入略有差异。训练时为了避免 Teacher Forcing 带来的分布偏移Decoder 的第二部分输入的是目标序列的真值但通过 Attention Mask 屏蔽掉未来信息。推理时第二部分用 0 填充模型自主预测。这种一次性生成的机制即使预测步长为 720 也不会累积误差这也是 Informer 在长序列上表现优秀的关键原因之一。3. 环境准备与数据集的实验配置3.1 依赖环境与版本选型我这次复现使用的是 PyTorch 2.x而不是官方仓库里的 PyTorch 1.8。很多人在环境这一步就卡住了主要是老代码用了torch.tensor的一些旧接口或 API 行为差异。如果你从零开始写而不是直接跑官方代码PyTorch 2.x 基本没有兼容性问题反而能用上 torch.compile 加速训练。我的环境配置如下Python 3.9PyTorch 2.1.1CUDA 11.8NumPy 1.24Pandas 2.0Matplotlib 3.7建议使用独立的 conda 环境避免污染其他项目的依赖。如果显存比较紧张可以考虑使用 16 位混合精度训练Informer 对精度没有那么敏感。我实测在 batch size 较大的情况下混合精度能减少约 30% 显存占用速度提升 20% 左右。后文会再次提到这个技巧。3.2 ETT 数据集说明与样本构造ETT 数据集是电力变压器温度及其它电力特征的时间序列包含 ETTh1、ETTh2小时级数据和 ETTm1、ETTm215 分钟级数据。每条数据包含 7 个特征比如油温、负载等。论文主要在这四个子集上验证模型性能。数据集中整段数据按时间顺序划分为训练集、验证集、测试集常见比例是 6:2:2但要注意不能随机打乱因为时间序列必须保持时间顺序否则会造成信息泄漏。我这次按 70% 训练、10% 验证、20% 测试的比例切分。样本构造采用滑动窗口方式每个样本包含一段长度为 seq_len 的 encoder 输入序列、一段长度为 label_len 的 start token 序列以及一段长度为 pred_len 的目标预测序列。以 ETTh1 上经典配置为例seq_len 96label_len 48pred_len 48每个训练样本就是一个长度为 96 的历史窗口和紧随其后的 48 个未来时间点。窗口按步长 1 滑动因此样本数量非常多足以支撑深度模型训练。还有一个重要细节归一化。时间序列数据通常先做 z-score 归一化即减去均值除以标准差。均值与标准差需要只在训练集上计算然后应用到验证集和测试集。如果你不小心用全量数据计算统计量验证集和测试集的信息就提前泄漏到了训练过程评估结果会虚高。这是实际复现中最常见且隐蔽的错误。4. 完整代码实现与解析下面进入正题我按照从数据到模型的顺序放出完整可运行的实现代码并逐段解析。4.1 数据预处理与 Dataset 实现首先是最核心的数据读取与窗口切片。我封装了一个TimeSeriesDataset类它接收原始 numpy 数组和三个长度参数在__getitem__里完成窗口切分import numpy as np import torch from torch.utils.data import Dataset class TimeSeriesDataset(Dataset): def __init__(self, data, seq_len, label_len, pred_len): self.data data self.seq_len seq_len self.label_len label_len self.pred_len pred_len def __len__(self): return len(self.data) - self.seq_len - self.pred_len 1 def __getitem__(self, idx): s idx # encoder 输入s - sseq_len enc_x self.data[s: s self.seq_len] # decoder 输入sseq_len-label_len - sseq_lenpred_len dec_x self.data[s self.seq_len - self.label_len: s self.seq_len self.pred_len] # 预测目标sseq_len - sseq_lenpred_len label_y self.data[s self.seq_len: s self.seq_len self.pred_len] return torch.FloatTensor(enc_x), torch.FloatTensor(dec_x), torch.FloatTensor(label_y)特别注意 decoder 输入和预测目标在时间上的对齐关系。decoder 输入最后 pred_len 部分在训练阶段是目标真值的前半段我们靠 mask 挡住未来信息这个对齐关系写错一个索引模型表现会断崖式下降。我第一次写就漏了 seq_len 与 label_len 的错位结果训练 loss 怎么都降不下去排查了很久。4.2 ProbSparse Attention 的 PyTorch 实现核心模块我直接给出可运行的实现import math import torch import torch.nn as nn import torch.nn.functional as F class ProbSparseAttention(nn.Module): def __init__(self, d_model, n_heads, sample_k5, n_top25): super().__init__() self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads self.sample_k sample_k self.n_top n_top self.q_proj nn.Linear(d_model, d_model) self.k_proj nn.Linear(d_model, d_model) self.v_proj nn.Linear(d_model, d_model) self.out_proj nn.Linear(d_model, d_model) def forward(self, x): B, L, _ x.shape H, D self.n_heads, self.d_k q self.q_proj(x).view(B, L, H, D).transpose(1, 2) k self.k_proj(x).view(B, L, H, D).transpose(1, 2) v self.v_proj(x).view(B, L, H, D).transpose(1, 2) # 随机采样 key 来计算稀疏度分数 U int(self.sample_k * math.log(L)) if U 0 and U L: idx torch.randperm(L)[:U].to(k.device) k_sample k.index_select(2, idx) # [B, H, U, D] else: k_sample k scores torch.matmul(q, k_sample.transpose(-2, -1)) / math.sqrt(D) M scores.max(-1).values - scores.mean(-1) # [B, H, L] # 选取 top n_top 个 query 的索引 _, top_idx torch.topk(M, self.n_top, dim-1) # [B, H, n_top] # 对 top 的 query 计算完整注意力 q_top torch.gather(q, 2, top_idx.unsqueeze(-1).expand(-1, -1, -1, D)) attn_top torch.matmul(q_top, k.transpose(-2, -1)) / math.sqrt(D) attn_top F.softmax(attn_top, dim-1) out_top torch.matmul(attn_top, v) # [B, H, n_top, D] # 非 top 位置用全局均值填充 out_mean out_top.mean(dim-2, keepdimTrue) # [B, H, 1, D] out out_mean.expand(B, H, L, D).clone() out out.scatter(-2, top_idx.unsqueeze(-1).expand(-1, -1, -1, D), out_top) out out.transpose(1, 2).contiguous().view(B, L, self.d_model) return self.out_proj(out)这段的实现你在官方代码里也能见到但有几个细节值得说明。sample_k控制随机采样的 key 数量论文里的经验值是 5也就是采样 5 * log(L) 个 key。如果输入序列长度很小比如 L8采样数量可能比 L 还大这时就直接用全部 key 做稀疏度计算不做采样。另一个注意点是n_top即最终保留的 query 数量论文默认取 25。如果输入序列很长可以适当调大比如 50 甚至 100。保留太少会丢失信息保留太多又失去稀疏化带来的效率优势。这个参数和sample_k是影响计算效率与精度的两个核心旋钮。4.3 Encoder、Distilling 与 Decoder 组装有了单头稀疏注意力下一步是多头封装加一维卷积蒸馏。蒸馏模块的实现很直观class DistillingLayer(nn.Module): def __init__(self, d_model): super().__init__() self.conv nn.Conv1d(d_model, d_model, kernel_size3, padding1) self.pool nn.MaxPool1d(kernel_size2) def forward(self, x): # x: [B, L, D] x x.transpose(1, 2) # [B, D, L] x F.gelu(self.conv(x)) x self.pool(x) # 长度减半 return x.transpose(1, 2)Encoder 堆叠时每一层由多头 ProbSparse Attention、前馈网络和 LayerNorm 组成在残差连接后接蒸馏层。我在复现时只在前两层之后接蒸馏第三层输出直接返回避免特征被过度压缩。对应到论文配置Encoder 一共 4 层前 3 层带蒸馏第 4 层是纯注意力输出。Decoder 的结构相对标准由带 mask 的多头注意力防止未来信息泄漏和交叉注意力组成。需要强调的是Decoder 的 mask 不需要像自回归 Transformer 那样呈严格的下三角矩阵。因为 Informer 是生成式一次性输出训练时只需要保证预测部分的每个位置看不到它之后的真实值start token 部分是可见的。实际实现中通常使用一个对角线为 0 的矩阵让每个位置只看到当前位置及之前的位置。4.4 训练循环与评估指标训练循环整体和一般 PyTorch 模型没有本质差别但有两个细节值得注意。第一学习率调度我使用了 Adams 优化器配 warmup 和余弦退火第二验证集 loss 与测试集指标分开计算。评估指标使用时间序列预测的标准 MAE平均绝对误差与 MSE均方误差计算时注意要把归一化后的预测结果还原到原始量纲再计算否则指标不是真实物理含义。如果不做反归一化MSE 会受标准化尺度影响不同数据集之间的指标没有可比性。这也是很多复现结果看起来优于论文实际结果的原因之一。下面是训练循环的核心部分optimizer torch.optim.Adam(model.parameters(), lr1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max10) for epoch in range(epochs): model.train() train_loss 0.0 for enc_x, dec_x, label_y in train_loader: enc_x enc_x.to(device) dec_x dec_x.to(device) label_y label_y.to(device) pred model(enc_x, dec_x) pred pred[:, -pred_len:, :] # 只取预测段输出 loss F.mse_loss(pred, label_y) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() train_loss loss.item() * len(enc_x) scheduler.step() val_loss evaluate(model, val_loader) print(fEpoch {epoch}: train_loss{train_loss / len(train_loader.dataset):.6f}, val_loss{val_loss:.6f})梯度裁剪是我强烈建议加的一步Informer 在训练初期偶尔会出现 loss 突然跳变梯度裁剪能有效规避这个问题尤其在长序列输入下。实践中 max_norm 设为 1.0 效果稳定。5. 完整训练过程与参数调优实录5.1 实验配置与超参数解释我这次复现使用 ETTh1 数据集设定 batch size 为 32。模型维度 d_model512注意力头数 n_heads8Encoder 层数 4Decoder 层数 2dropout 0.05。预测长度分别测试了 24、48、96、168 四种。下面这张表记录了我最终的参数组合参数值说明seq_len96encoder 输入长度label_len48decoder start token 长度pred_len24/48/96/168预测长度d_model512模型隐层维度n_heads8注意力头数e_layers4encoder 层数d_layers2decoder 层数batch_size32训练批大小lr1e-4初始学习率dropout0.05丢弃概率epochs10最大训练轮数关于 label_len 有一点必须展开它决定 decoder 能借到多少真实历史信息作为 start token。如果 label_len 太小比如设为 0decoder 相当于没有任何引导信息预测效果会急剧恶化。论文实验里常见 label_len48对应 2 天的小时级数据这样模型在预测未来 1 天时有一个足够长的“锚点”。5.2 训练过程的收敛观察我在 ETTh1 上训练时前 2 个 epoch loss 下降很快第 3 个 epoch 开始进入平缓期到第 6 个 epoch 左右验证损失基本收敛。继续训练到第 10 个 epoch验证 loss 有小幅波动但不再明显下降。这个收敛速度和论文报告的“ETTh1 上大约 6 epoch 收敛”基本吻合。预测长度为 24 时收敛速度最快模型在第 4 个 epoch 就已经达到接近最优的性能。预测长度为 168 时收敛明显变慢且最终 MSE 比预测 24 高出一个数量级这说明长序列预测的难度确实在指数级上升。我还对比了 ProbSparse Attention 与标准 Full Attention 在同一训练条件下的效果。在 ETTh1 上Full Attention 的最终精度略高一点点大概 1% 以内但训练时间大约是 Spars 版本的两倍多。这个结果和论文结论一致稀疏注意力以微小的精度损失换取显著的计算效率提升。如果你做的是短序列Full Attention 也许够用一旦序列长度超过 500两者差距会拉到数倍。5.3 影响精度的关键细节有几个细节对精度影响极大这里单独强调。第一mask 的实现位置。Decoder 中的 masked self-attention 一定要在 softmax 之前把 mask 位置的分数设置为负无穷而不是把权重设成 0。若在 softmax 之后置零会导致概率和不为 1梯度传播也会出问题。这个基础错误排查起来特别痛苦因为 loss 不会崩只是指标一直不达标。第二预测输出的切片位置。模型输出的序列长度是 label_len pred_len训练时计算 loss 之前一定要切片到最后的 pred_len 位置不能用整个输出序列和 label_y 比较。很多人第一次写都会忽略最终导致训练 loss 看起来很低但预测完全不对因为模型学会了把前段 start token 直接“复制”出来。第三batch size 和显存的权衡。使用 seq_len96、d_model512 时GPU 显存需求不高8G 显存就能跑。但如果你把 seq_len 提升到 336 或 720显存占用会大幅上升。此时建议优先降低 batch size 到 16 甚至 8而不是缩减模型维度这样对最终指标影响更小。另外可以开启torch.cuda.amp混合精度训练实测显存大约能省 30%。6. 复现过程中的高频问题与解决方案这部分是我认为整篇文章最有价值的地方。很多问题你不在真实项目里踩一遍永远不知道坑有多深。6.1 问题训练 loss 正常下降但测试结果全是“平移错觉”我复现中遇到的第一个诡异现象是测试曲线和真实曲线形状高度相似但整体落后了一个固定周期。原因不是模型学会了预测而是它学会了“抄作业”——把 start token 的最后一段时间直接平移作为预测结果。这在数据有很强自相关时特别容易发生模型发现与其学习复杂模式不如直接复制最近观测值。排查这个问题的思路很简单用一条无周期性、有随机突变的数据做 sanity check。如果模型在这种数据上预测失效说明它确实在偷懒。解决办法是加强 label_len 和 pred_len 的相对比例避免 decoder 有太多真实信息可供复制同时检查训练 loss 是否有明显过拟合迹象。从实现上说确保 decoder 输入中预测部分的 mask 正确遮挡了未来信息也能显著缓解这个问题。6.2 问题Sparse Attention 在短序列上反而更慢ProbSparse Attention 理论上要求序列长度较大时才能体现优势因为它的采样数量和 log(L) 成正比当 L 很小时整个流程里的 gather、scatter 等额外操作反而增加了开销。这个现象在复现时很容易被忽略。如果发现序列长度小于 128 时模型跑得比 Full Attention 还慢不建议强行用 Informer。可以退一步用标准的 Transformer Attention或者扩大 seq_len 让模型在更长的上下文上学习。Informer 的设计初衷就是应对长序列把它硬套到短序列场景属于用错工具。6.3 问题训练中期 loss 突然飙升我训练到第 4 个 epoch 时 loss 突然比之前高了一个量级随后又慢慢降回来。一开始以为是代码 bug后来定位到是学习率过大导致的震荡。虽然初始学习率 1e-4 整体温和但在预训练后期直接把学习率调大模型参数可能会跳到损失平面的陡峭区域。解决方案是使用 warmup 策略前几个 epoch 让学习率从很小的值线性增加到目标值然后再用余弦退火衰减。这样既能保证训练初期的稳定性又能在后期精细收敛。另外梯度裁剪也可以一并开启双保险。7. 复现总结与我的个人建议这次完整复现 Informer 前后花了大约一周时间其中编码环境搭建和数据 pipeline 花了 1 天模型核心模块实现花了 2 天训练调参和性能验证花了 3 天最后整理可视化结果又用了 1 天。如果只算有效代码量核心模型大概 300 行但理解每一行的含义和它的动机远比写出这些行代码更耗时。我在实际复现中最深的体会是Informer 的论文写得非常精致核心公式不多但每个设计都是为了对抗长序列预测的真实痛点。你只有自己动手实现一遍才会理解 ProbSparse 采样中每个参数的含义才会理解为什么蒸馏层把特征长度减半是一种特征提取而非简单压缩。这套“从问题到设计”的思维方式比单纯跑通代码更值得迁移到其他模型上。最后分享一个小技巧复现任何论文模型前先把论文里的所有超参数和数据集配置整理成一份清单再和官方代码逐一对照。这能帮你快速定位官方实现里哪些是论文中刻意突出的创新点哪些只是工程上为了训练稳定的常规选择。我自己每次都先做这一步省下大量试错时间。后续如果你想继续把这个模型用到自己的数据上建议从 label_len 和 seq_len 两个参数开始调它们是影响最终预测精度的成本最低的杠杆。