做时间序列预测做到这两年大家应该都有一种很真实的纠结Transformer 架构在长序列上的能力强但开销也大想用混合专家MoE来摊薄成本吧又常常被离散路由的 token 丢弃、负载不均和训练不稳定搞到心态爆炸。最近我在整理实验笔记时又翻到 TIME-MADEIRA 这个模型全称是 Time-Series Forecasting via Time-SoftMoE一个把块状稀疏注意力和时间维度上的软混合专家融合进同一个 Transformer 块的工作。名字里的 Madeira 是陈年加强葡萄酒作者用这个名字想表达的意思我理解下来大概是时间序列的每一段局部都值得被更细腻、更专业地加工。这篇文章我准备顺着这条线从头拆到底这个模型到底在解决什么问题、两个核心组件各自的工作原理是什么、复现时有哪些关键实现细节以及我在实际调参中踩过的坑。适合正在做长序列预测、想尝试在时序任务里引入 MoE、或者被传统 MoE 各种毛病折磨过的同学阅读。1. 项目解读TIME-MADEIRA 到底在解决什么问题1.1 长序列预测绕不开的两个瓶颈先说第一个瓶颈注意力机制的计算量。标准自注意力中长度为 n 的序列要计算 n×n 的注意力矩阵复杂度是 O(n²)。在文本任务里人们可以用稀疏注意力、线性注意力这些手段来缓解但时间序列有自己的脾气——序列不是一堆无关 token 的集合而是高度自相关的连续观测盲目剪掉注意力边又会破坏长距离依赖。PatchTST 的思路是把连续观测切成若干 patch用 patch 作为 token一下子把 n 缩小了一个量级。但 patch 只是在输入端减少了 token 数Transformer 内部每个 token 依然走同一套前馈网络换句话说模型对凌晨 3 点的用电负荷和中午 12 点的用电负荷用的是同一组参数来建模。第二个瓶颈就藏在这里时序数据天然具有局部模式差异。同一个预测任务里不同时间段往往对应不同的形态——工作日早高峰、节假日波动、季节性周期、突发漂移。一个稠密 Transformer 只有一个共享的 FFN它必须把所有模式都压缩进同一组权重里而一个通常意义上的 MoE 又只在 token 级别做离散路由选中的专家处理、未选中的 token 直接被丢掉。对 NLP 来说丢掉一些冗余 token 可能还能接受对时间序列来说丢掉任何一个观测点都可能意味着丢掉一个拐点、一个异常峰值甚至一个周期相位。这就是时序任务里直接套用传统 MoE 最容易翻车的地方。所以在读 TIME-MADEIRA 的时候我特别注意到它的两个设计取舍一是把稀疏从 token 选择层面换到了路径层面二是把 MoE 从特征维度搬到了时间维度。下面分开拆。1.2 从硬 MoE 到软 MoE时序场景为什么更难传统 MoE比如 Switch Transformer、GShard在每层都放一个 router它对每个 token 做 top-k 选择k 一般取 1 或 2。这种硬路由有四个公认的问题token 丢弃没被选中的 token 不进入任何专家前馈信息直接蒸发、负载不均衡router 容易把所有 token 都塞给同一个专家、专家坍缩少数专家学会偷懒梯度集中在几个专家上、训练不稳定离散选择和后续层梯度之间缺乏平滑的映射。时序任务还会放大其中两个问题。第一相邻 token 高度相关top-k 硬路由经常把相邻的一段时间片断分配给同一个专家表面上看负载均衡 loss 正常了实际上专家成了分段的二传手根本没有学到有区分度的局部模式。第二时间序列的分布不是静态的训练集里工作日和周末的比例可能差异很大硬路由会放大这种分布偏差让某些专家在某个时间段内完全失活。TIME-MADEIRA 把 Soft MoE 引入时间维度正是冲着这两点去的路由不再做离散取舍而是给每个 token 在所有专家槽位上分配一个连续权重所有 token 的信息都以加权形式进入专家、再以同样的加权方式汇合回来。信息不会被丢弃梯度也能平滑回传负载由软分配天然摊开。这一节的要点是在做模型选型时时序任务比 NLP 更需要软的 MoE 设计。如果只把 NLP 的 MoE 代码搬到时间序列上大概率会在 20 步以内见到 loss 抖动和专家退化而不是稳定下降。2. 核心机制拆解Time-SoftMoE 与块状稀疏注意力2.1 软路由怎么做到既稀疏又稳定Soft MoE 的核心是槽位slot。假设我们有 E 个专家每个专家配置 s 个槽位那么这一层总共有 E×s 个槽位。前向过程可以分成三步第一步对每个输入 token 学出一个长度为 E×s 的分配权重用 softmax 在槽位维度归一第二步把每个槽位当作一个容器把所有 token 按权重加权求和到这个容器里得到每个槽位的输入向量注意这里是加权和而不是选 top-k所以每个槽位都收到了全部 token 的信息只是侧重不同第三步每个专家只处理属于自己那 s 个槽位的向量专家的输出再按同样的分配权重加权回传给所有 token。对比一下硬路由软路由最大的变化在于没有离散选择只有连续加权。你可以想象一个自助餐厅传统 MoE 是窗口打饭服务员只给你面前这一排选好的菜Soft MoE 是每个槽位都像一口锅所有食材按不同比例倒进每一口锅每一口锅炖完后又按原来的比例把味道分回每一个人的盘子里。因为整个过程完全可微梯度可以顺着权重一路传回 router训练时就不需要辅助负载均衡 loss 了run 起来会省心不少。在 TIME-MADEIRA 里这个 Soft MoE 被特意作用在时间维度token 序列被划分成若干块路由权重是在块内甚至是块间学出来的分配关系。设计意图很清楚——让模型自动学会哪些时间段的 token 内容应该交给哪些更懂这种模式的专家处理。专家容量不取决于输入序列长度而取决于槽位总数所以当槽位总数明显小于输入 token 数时FFN 的计算量可以被显著压低这也是它在参数增加的同时还能控制 FLOPs 的原因。2.2 块状稀疏注意力注意力开销从平方降到接近线性光有软路由还不够注意力开销依然是时序 Transformer 的大头。TIME-MADEIRA 给注意力做的处理叫块状稀疏注意力Blockwise Sparse Attention。做法不复杂把长度为 n 的序列平均分成 B 个块每块有 L 个 token注意力不再允许每个 token 看全序列而是只允许它看到自己所在块内的 token外加少数几个全局 token一般是可学习的全局向量或者序列开头的一个汇总 token。这样复杂度就变成了 B×L² 加上对全局 token 的 n×g其中 g 是全局 token 数。拿一个具体数字感受一下输入长度 1024块长 32全局 token 数取 2注意力计算量大约是 32×32² 1024×2 ≈ 34816 次关联计算而全量自注意力要算 1024×1024 ≈ 1048576。差了约 30 倍。更关键的是B×L² 本质上等于 n×L当块长 L 固定、序列长度 n 变大时计算量是线性增长的这就是它可以平稳外推到更长序列的原因。这里要特别澄清一个常见误解块状稀疏注意力并不是把所有 token 锁死在局部视野里。因为每个块都还保留了对全局 token 的注意力全局 token 相当于一个信息汇聚中心序列的长期趋势、全局周期信息是通过这个中心在各块之间流动的。你可以把它理解成每个小区都保有一个通向城市主干道的匝道局部信息在小区内交换全局信息通过匝道汇总。实际使用时全局 token 的个数不需要多1 到 4 个基本够了太多反而会稀释局部注意力。2.3 两个组件如何配合出局部专家把两块拼到一起才看得出 TIME-MADEIRA 真正的设计心法。注意力负责的是信息视野每个 token 该跟谁交换信息是空间维度的关系建模Soft MoE 负责的是处理策略这些信息该由哪组参数来加工是能力维度的专业化分工。块状边界同时服务于两者——注意力在块内收敛MoE 也以块为基本粒度去学习路由分配。这个设计的直接收益是模型参数量可以做得比较大但实际计算量不会跟着线性膨胀。比如你可以把专家数从 4 加到 16每个专家的 FFN 参数也成倍增加但只要槽位总数固定前馈计算量基本维持不变。传统稠密 Transformer 想做同样的事只能硬着头皮加大 hidden size计算量和参数量一起涨。论文在 ETTh、ETTm、Electricity、Traffic 等常见基准的多数预测尺度上用这样的结构比 PatchTST、TiDE 等基线拿到了更好的 MSE/MAE本质上就是用更聪明的参数布局换精度而不是单纯堆参数。我在实际理解时喜欢用科室分工来类比一般 Transformer 像是一个全科医生接诊所有病人PatchTST 像是把病人先统一分诊到内科但内科只有一位医生TIME-MADEIRA 则是分诊之后内科里还分了呼吸、心血管、神经几个亚专科每个专科医生只看自己最擅长的病例而分诊台router由模型自己动态学习不动声色。3. 从论文到复现关键实现细节与超参经验3.1 数据与任务设定先把基准跑对如果想验证 TIME-MADEIRA 的真实水平第一步是把实验和论文对齐。常见的公开数据集有四个系列ETT 系列电力变压器温度包含 ETTh1/ETTh2、ETTm1/ETTm2分别按小时和 15 分钟采样、Weather气象观测站数据10 分钟采样、Electricity321 个用户的用电负荷15 分钟采样、Traffic道路占用率小时采样。绝大部分论文都会在这四类数据上做长序列预测预测长度取 96、192、336、720输入长度一般取 96 或者 336。数据划分和归一化是两个最容易埋雷的地方。ETT 系列内部有官方划分前 12 个月训练、4 个月验证、4 个月测试其他数据集的常见做法也是按时间顺序切一定不能随机打乱——时序数据一旦随机切验证集就把未来信息泄漏进训练集了。归一化我建议按通道做独立统计每个 channel 用自己的均值方差做标准化预测完再反归一化回原始量纲如果为了省事用全局统计遇上通道之间量纲差异大的数据集比如 Electricity 和 Traffic误差会明显放大。数据集采样频率特征通道数常见预测长度ETTh1 / ETTh2小时级796 / 192 / 336 / 720ETTm1 / ETTm215 分钟级796 / 192 / 336 / 720Weather10 分钟级2196 / 192 / 336 / 720Electricity15 分钟级32196 / 192 / 336 / 720Traffic小时级86296 / 192 / 336 / 720需要提醒一句ETT 系列里面 ETTh 和 ETTm 虽然都只有 7 个通道但采样频率差 4 倍同样的块长在不同数据上对应的真实时间跨度完全不同后面调 block size 时要格外小心。3.2 TST 块实现的关键点TIME-MADEIRA 的完整代码在官方仓库里可以找到但如果你是自己复现核心的 Soft MoE 前向其实可以浓缩成下面这段逻辑import torch import torch.nn.functional as F def soft_moe_forward(x, expert_ffn, slot_projector, expert_proj, num_experts, slots_per_expert): # x: (batch, n, d) B, n, d x.shape total_slots num_experts * slots_per_expert # 1. 为每个 token 生成分配到所有槽位的权重 dispatch_logits slot_projector(x) # (B, n, total_slots) dispatch_weights F.softmax(dispatch_logits, dim-1) # 槽位维度归一全局无丢弃 # 2. 把所有 token 按权重加权汇入每个槽位 slot_input torch.einsum(bns,bsd-bsd, dispatch_weights, x) # 3. 每个专家只处理自己那 slots_per_expert 个槽位 slot_input slot_input.view(B, num_experts, slots_per_expert, d) expert_out torch.empty_like(slot_input) for e in range(num_experts): expert_out[:, e] expert_ffn[e](slot_input[:, e]) # 每个专家一个小 FFN # 4. 输出按相同权重加权回每个 token expert_out expert_out.view(B, total_slots, d) out torch.einsum(bns,bsd-bnd, dispatch_weights, expert_out) return out四行核心逻辑分别对应路由、汇聚、专家处理、回传。实际工程里 expert_ffn 通常是一个共享参数但分组的 FFN可以一次性完成 E 个专家的前向计算不需要循环上面的写法只是为了把语义讲清楚。要注意的是 dispatch_weights 在 slot 维度做了 softmax但如果总槽位很多softmax 输入很大数值上容易推不动可以在 softmax 之前对 logits 做一次缩放除以 sqrt(d) 之类这是我从实际训练里学到的稳定化技巧。块状稀疏注意力的实现更依赖框架习惯。我用的方案是先把序列按块长 reshape 成 (B, num_blocks, block_len, d) 的形状在块内做普通自注意力同时把全局 token 的 KV 拼到块内注意力后面。PyTorch 自带 F.scaled_dot_product_attention 支持 mask可以把它当成一个小型 meta-batch 去做如果数据通道多还可以考虑把通道维度并进 batch 维实现 channel independence 的同时减少显存碎片。3.3 训练配置与显存优化复现时序 Transformer 的训练配置有一个通用起点我测过几次基本都能收敛优化器用 Adam初始学习率 1e-4 到 5e-4前 5 到 10 个 epoch 线性 warmup后面用 cosine 或者阶段性下降batch size 32 到 128 之间看显存和数据集调整训练轮数不用太多早停在验证集上基本 30 到 60 epoch 就有结论了。MSE 作为主损失函数评估指标用 MSE 和 MAE。三个训练细节值得单独说。第一混合精度建议优先用 bf16 而不是 fp16因为路由权重和 softmax 对精度更敏感fp16 在大专家数量下容易出现梯度溢出。第二梯度累积节省不了实验时间但能稳定训练如果你的 batch size 因为显存限制压得特别小先用梯度累积把等效 batch 撑到 64 以上再去看模型结构问题。第三Soft MoE 层不强制加辅助负载均衡损失如果你想加也可以加但加了之后路由权重的熵会明显上升反而可能抑制专家的专业化我实测下来不如不加。显存方面最大的两个头是注意力矩阵和路由矩阵。块注意力把注意力矩阵从 O(n²) 降到了 O(n·L)这个头基本解决了路由矩阵是 n×E×s如果专家数到了 16 且输入长度 720它也是一笔不小的开销。建议槽位数先从 1 开始验证集有提升再加到 2全局 token 数量保持在 2 到 4实在不够就开 gradient checkpointing只对 TST 块做 checkpoint其他层照常。3.4 与 PatchTST 公平对比的注意事项如果你是想跟 PatchTST 做对比或者复现论文表格有两点容易被人忽略。第一是 patch 长度的对齐。PatchTST 的 patch size 通常取 16而 MADEIRA 的块长本质是一个类似的概念如果你用块长 32 的 MADEIRA 去比 patch size 16 的 PatchTST比较的是两种 token 粒度结论没意义。建议先在同一输入长度下扫一遍块长16、24、32各自取最优值再比较。第二是参数量的口径。MoE 模型的总参数量包含所有专家的参数但每次前向只激活一部分直接拿总参数量对比稠密模型会显得不公平主流的做法是同时汇报参数量和 FLOPs单看 FLOPs 更有实际价值。最后多 seed 平均。时序模型在不同随机种子下的结果波动比我预想的大尤其路由初始化的位置对最终收敛影响不小。我一般跑 3 个 seed报均值和方差如果时间紧张至少跑两个 seed 确认趋势不是随机种子带来的。4. 常见问题与排查实录4.1 效果反而不如 PatchTST我在复现 MADEIRA 的第一周就遇到过这个情况块长 32、专家数 8、槽位 1号称改进了所有基线结果在 ETTh1 的 720 步预测上就是追不上 PatchTST。后来逐项排查问题出在块长和数据周期的错配上。ETTh1 是小时级数据日周期是 24 个点周周期是 168 个点我选的块长 32 既不能完整包含一个 24 点日周期又和 168 不整除路由学到的大多是一些破碎的边界。把块长改成 24 后正好对齐一个自然日周期效果立刻上来了。第二个常见问题是专家数量设置不敏感。如果槽位总数远小于块数路由退化成每个专家平均分摊所有块没有分工如果槽位总数接近 token 数Soft MoE 又退化成普通的稠密 FFN。我的经验是先把槽位总数控制在块数的四分之一到一半这个量级再看路由矩阵的熵是否处于中间状态——熵太大说明分工不明显熵太小说明路由太自信容易脆。第三个原因在输入长度上。PatchTST 在短输入96时表现很好因为 patch 天然做了局部压缩MADEIRA 的块注意力需要更长的上下文才能体现优势。如果你在预测长度较大336/720但输入长度只有 96 的配置下测试模型会有点吃亏建议对比实验统一用输入长度 336。4.2 训练不稳定或显存不够训练过程中最容易遇到的是 loss 曲线突然跳到 NaN。优先查这三处是否用了 fp16 混合精度换成 bf16路由 logits 是否没做缩放加上类似 sqrt(d) 的缩放batch size 是否过小导致梯度异常。还有一个不那么常见但真的会发生的数据集里存在 NaN 或无穷值特别是 Weather 这种传感器数据预处理时最好先做一次缺失值检查。显存不够的优化顺序我建议是先减全局 token 数从 4 降到 2→ 再减槽位数从 2 降到 1→ 再减专家数从 16 降到 8→ 还不行才用梯度 checkpointing。为什么不建议一上来就开 checkpointing因为它会显著拖慢训练等于用时间换显存排查问题时节奏会变得很慢。设置上还有一个容易被忽略的点PyTorch 默认会为每个优化器参数保有 momentum 和二阶矩MoE 层参数量大优化器状态可能比模型本身还占显存必要时对路由层用 Adafactor 或者 SGD 试试虽然收敛速度会变但显存明显缓解。4.3 推理阶段的效率与部署心得模型跑通之后如果还要上线或者做大批量回测推理效率是另一个重点。块状稀疏注意力在推理阶段的优势是 KV cache 小——每个 token 只需要缓存本块的 KV 和全局 token 的 KV不需要缓存全序列这在线性增长的序列长度下非常友好。实际部署时可以把每个专家 FFN 单独导出或量化按输入动态加载存储和内存压力都能摊薄。一个小坑是批量推演时的 padding。序列长度不是块长的整数倍时直接 reshape 会出问题我的做法是在序列末尾做零填充并把填充位置加到注意力 mask 里防止填充 token 污染全局 token 的聚合。如果你要对比推理速度务必统一 batch 大小、输入长度和输出长度只改变模型结构否则测出来的差异更多是框架层面的。最后说一点我自己实验中的体会TIME-MADEIRA 最值得借鉴的其实不是某一个组件而是把稀疏性从选择层面挪到权重层面这个思路。我在很多任务里把硬路由换成软路由之后即使模型结构没怎么改训练的稳定性和最终精度都有改善这个收益比单纯调专家数大得多。另外块长与数据周期的对齐几乎是免费的增益动手之前先去做一下周期分析比闷头扫参数高效得多。如果你打算在自己数据上试建议从块长 24、专家 8、槽位 1 这个配置出发先跑通一个短预测长度的 baseline再逐步加长预测距离。这个模型后续往多变量通道交互、非平稳序列分解这些方向扩展的空间还很大值得持续关注。 TIME-MADEIRA模型解析：软MoE与块稀疏注意力融合的时序预测
做时间序列预测做到这两年大家应该都有一种很真实的纠结Transformer 架构在长序列上的能力强但开销也大想用混合专家MoE来摊薄成本吧又常常被离散路由的 token 丢弃、负载不均和训练不稳定搞到心态爆炸。最近我在整理实验笔记时又翻到 TIME-MADEIRA 这个模型全称是 Time-Series Forecasting via Time-SoftMoE一个把块状稀疏注意力和时间维度上的软混合专家融合进同一个 Transformer 块的工作。名字里的 Madeira 是陈年加强葡萄酒作者用这个名字想表达的意思我理解下来大概是时间序列的每一段局部都值得被更细腻、更专业地加工。这篇文章我准备顺着这条线从头拆到底这个模型到底在解决什么问题、两个核心组件各自的工作原理是什么、复现时有哪些关键实现细节以及我在实际调参中踩过的坑。适合正在做长序列预测、想尝试在时序任务里引入 MoE、或者被传统 MoE 各种毛病折磨过的同学阅读。1. 项目解读TIME-MADEIRA 到底在解决什么问题1.1 长序列预测绕不开的两个瓶颈先说第一个瓶颈注意力机制的计算量。标准自注意力中长度为 n 的序列要计算 n×n 的注意力矩阵复杂度是 O(n²)。在文本任务里人们可以用稀疏注意力、线性注意力这些手段来缓解但时间序列有自己的脾气——序列不是一堆无关 token 的集合而是高度自相关的连续观测盲目剪掉注意力边又会破坏长距离依赖。PatchTST 的思路是把连续观测切成若干 patch用 patch 作为 token一下子把 n 缩小了一个量级。但 patch 只是在输入端减少了 token 数Transformer 内部每个 token 依然走同一套前馈网络换句话说模型对凌晨 3 点的用电负荷和中午 12 点的用电负荷用的是同一组参数来建模。第二个瓶颈就藏在这里时序数据天然具有局部模式差异。同一个预测任务里不同时间段往往对应不同的形态——工作日早高峰、节假日波动、季节性周期、突发漂移。一个稠密 Transformer 只有一个共享的 FFN它必须把所有模式都压缩进同一组权重里而一个通常意义上的 MoE 又只在 token 级别做离散路由选中的专家处理、未选中的 token 直接被丢掉。对 NLP 来说丢掉一些冗余 token 可能还能接受对时间序列来说丢掉任何一个观测点都可能意味着丢掉一个拐点、一个异常峰值甚至一个周期相位。这就是时序任务里直接套用传统 MoE 最容易翻车的地方。所以在读 TIME-MADEIRA 的时候我特别注意到它的两个设计取舍一是把稀疏从 token 选择层面换到了路径层面二是把 MoE 从特征维度搬到了时间维度。下面分开拆。1.2 从硬 MoE 到软 MoE时序场景为什么更难传统 MoE比如 Switch Transformer、GShard在每层都放一个 router它对每个 token 做 top-k 选择k 一般取 1 或 2。这种硬路由有四个公认的问题token 丢弃没被选中的 token 不进入任何专家前馈信息直接蒸发、负载不均衡router 容易把所有 token 都塞给同一个专家、专家坍缩少数专家学会偷懒梯度集中在几个专家上、训练不稳定离散选择和后续层梯度之间缺乏平滑的映射。时序任务还会放大其中两个问题。第一相邻 token 高度相关top-k 硬路由经常把相邻的一段时间片断分配给同一个专家表面上看负载均衡 loss 正常了实际上专家成了分段的二传手根本没有学到有区分度的局部模式。第二时间序列的分布不是静态的训练集里工作日和周末的比例可能差异很大硬路由会放大这种分布偏差让某些专家在某个时间段内完全失活。TIME-MADEIRA 把 Soft MoE 引入时间维度正是冲着这两点去的路由不再做离散取舍而是给每个 token 在所有专家槽位上分配一个连续权重所有 token 的信息都以加权形式进入专家、再以同样的加权方式汇合回来。信息不会被丢弃梯度也能平滑回传负载由软分配天然摊开。这一节的要点是在做模型选型时时序任务比 NLP 更需要软的 MoE 设计。如果只把 NLP 的 MoE 代码搬到时间序列上大概率会在 20 步以内见到 loss 抖动和专家退化而不是稳定下降。2. 核心机制拆解Time-SoftMoE 与块状稀疏注意力2.1 软路由怎么做到既稀疏又稳定Soft MoE 的核心是槽位slot。假设我们有 E 个专家每个专家配置 s 个槽位那么这一层总共有 E×s 个槽位。前向过程可以分成三步第一步对每个输入 token 学出一个长度为 E×s 的分配权重用 softmax 在槽位维度归一第二步把每个槽位当作一个容器把所有 token 按权重加权求和到这个容器里得到每个槽位的输入向量注意这里是加权和而不是选 top-k所以每个槽位都收到了全部 token 的信息只是侧重不同第三步每个专家只处理属于自己那 s 个槽位的向量专家的输出再按同样的分配权重加权回传给所有 token。对比一下硬路由软路由最大的变化在于没有离散选择只有连续加权。你可以想象一个自助餐厅传统 MoE 是窗口打饭服务员只给你面前这一排选好的菜Soft MoE 是每个槽位都像一口锅所有食材按不同比例倒进每一口锅每一口锅炖完后又按原来的比例把味道分回每一个人的盘子里。因为整个过程完全可微梯度可以顺着权重一路传回 router训练时就不需要辅助负载均衡 loss 了run 起来会省心不少。在 TIME-MADEIRA 里这个 Soft MoE 被特意作用在时间维度token 序列被划分成若干块路由权重是在块内甚至是块间学出来的分配关系。设计意图很清楚——让模型自动学会哪些时间段的 token 内容应该交给哪些更懂这种模式的专家处理。专家容量不取决于输入序列长度而取决于槽位总数所以当槽位总数明显小于输入 token 数时FFN 的计算量可以被显著压低这也是它在参数增加的同时还能控制 FLOPs 的原因。2.2 块状稀疏注意力注意力开销从平方降到接近线性光有软路由还不够注意力开销依然是时序 Transformer 的大头。TIME-MADEIRA 给注意力做的处理叫块状稀疏注意力Blockwise Sparse Attention。做法不复杂把长度为 n 的序列平均分成 B 个块每块有 L 个 token注意力不再允许每个 token 看全序列而是只允许它看到自己所在块内的 token外加少数几个全局 token一般是可学习的全局向量或者序列开头的一个汇总 token。这样复杂度就变成了 B×L² 加上对全局 token 的 n×g其中 g 是全局 token 数。拿一个具体数字感受一下输入长度 1024块长 32全局 token 数取 2注意力计算量大约是 32×32² 1024×2 ≈ 34816 次关联计算而全量自注意力要算 1024×1024 ≈ 1048576。差了约 30 倍。更关键的是B×L² 本质上等于 n×L当块长 L 固定、序列长度 n 变大时计算量是线性增长的这就是它可以平稳外推到更长序列的原因。这里要特别澄清一个常见误解块状稀疏注意力并不是把所有 token 锁死在局部视野里。因为每个块都还保留了对全局 token 的注意力全局 token 相当于一个信息汇聚中心序列的长期趋势、全局周期信息是通过这个中心在各块之间流动的。你可以把它理解成每个小区都保有一个通向城市主干道的匝道局部信息在小区内交换全局信息通过匝道汇总。实际使用时全局 token 的个数不需要多1 到 4 个基本够了太多反而会稀释局部注意力。2.3 两个组件如何配合出局部专家把两块拼到一起才看得出 TIME-MADEIRA 真正的设计心法。注意力负责的是信息视野每个 token 该跟谁交换信息是空间维度的关系建模Soft MoE 负责的是处理策略这些信息该由哪组参数来加工是能力维度的专业化分工。块状边界同时服务于两者——注意力在块内收敛MoE 也以块为基本粒度去学习路由分配。这个设计的直接收益是模型参数量可以做得比较大但实际计算量不会跟着线性膨胀。比如你可以把专家数从 4 加到 16每个专家的 FFN 参数也成倍增加但只要槽位总数固定前馈计算量基本维持不变。传统稠密 Transformer 想做同样的事只能硬着头皮加大 hidden size计算量和参数量一起涨。论文在 ETTh、ETTm、Electricity、Traffic 等常见基准的多数预测尺度上用这样的结构比 PatchTST、TiDE 等基线拿到了更好的 MSE/MAE本质上就是用更聪明的参数布局换精度而不是单纯堆参数。我在实际理解时喜欢用科室分工来类比一般 Transformer 像是一个全科医生接诊所有病人PatchTST 像是把病人先统一分诊到内科但内科只有一位医生TIME-MADEIRA 则是分诊之后内科里还分了呼吸、心血管、神经几个亚专科每个专科医生只看自己最擅长的病例而分诊台router由模型自己动态学习不动声色。3. 从论文到复现关键实现细节与超参经验3.1 数据与任务设定先把基准跑对如果想验证 TIME-MADEIRA 的真实水平第一步是把实验和论文对齐。常见的公开数据集有四个系列ETT 系列电力变压器温度包含 ETTh1/ETTh2、ETTm1/ETTm2分别按小时和 15 分钟采样、Weather气象观测站数据10 分钟采样、Electricity321 个用户的用电负荷15 分钟采样、Traffic道路占用率小时采样。绝大部分论文都会在这四类数据上做长序列预测预测长度取 96、192、336、720输入长度一般取 96 或者 336。数据划分和归一化是两个最容易埋雷的地方。ETT 系列内部有官方划分前 12 个月训练、4 个月验证、4 个月测试其他数据集的常见做法也是按时间顺序切一定不能随机打乱——时序数据一旦随机切验证集就把未来信息泄漏进训练集了。归一化我建议按通道做独立统计每个 channel 用自己的均值方差做标准化预测完再反归一化回原始量纲如果为了省事用全局统计遇上通道之间量纲差异大的数据集比如 Electricity 和 Traffic误差会明显放大。数据集采样频率特征通道数常见预测长度ETTh1 / ETTh2小时级796 / 192 / 336 / 720ETTm1 / ETTm215 分钟级796 / 192 / 336 / 720Weather10 分钟级2196 / 192 / 336 / 720Electricity15 分钟级32196 / 192 / 336 / 720Traffic小时级86296 / 192 / 336 / 720需要提醒一句ETT 系列里面 ETTh 和 ETTm 虽然都只有 7 个通道但采样频率差 4 倍同样的块长在不同数据上对应的真实时间跨度完全不同后面调 block size 时要格外小心。3.2 TST 块实现的关键点TIME-MADEIRA 的完整代码在官方仓库里可以找到但如果你是自己复现核心的 Soft MoE 前向其实可以浓缩成下面这段逻辑import torch import torch.nn.functional as F def soft_moe_forward(x, expert_ffn, slot_projector, expert_proj, num_experts, slots_per_expert): # x: (batch, n, d) B, n, d x.shape total_slots num_experts * slots_per_expert # 1. 为每个 token 生成分配到所有槽位的权重 dispatch_logits slot_projector(x) # (B, n, total_slots) dispatch_weights F.softmax(dispatch_logits, dim-1) # 槽位维度归一全局无丢弃 # 2. 把所有 token 按权重加权汇入每个槽位 slot_input torch.einsum(bns,bsd-bsd, dispatch_weights, x) # 3. 每个专家只处理自己那 slots_per_expert 个槽位 slot_input slot_input.view(B, num_experts, slots_per_expert, d) expert_out torch.empty_like(slot_input) for e in range(num_experts): expert_out[:, e] expert_ffn[e](slot_input[:, e]) # 每个专家一个小 FFN # 4. 输出按相同权重加权回每个 token expert_out expert_out.view(B, total_slots, d) out torch.einsum(bns,bsd-bnd, dispatch_weights, expert_out) return out四行核心逻辑分别对应路由、汇聚、专家处理、回传。实际工程里 expert_ffn 通常是一个共享参数但分组的 FFN可以一次性完成 E 个专家的前向计算不需要循环上面的写法只是为了把语义讲清楚。要注意的是 dispatch_weights 在 slot 维度做了 softmax但如果总槽位很多softmax 输入很大数值上容易推不动可以在 softmax 之前对 logits 做一次缩放除以 sqrt(d) 之类这是我从实际训练里学到的稳定化技巧。块状稀疏注意力的实现更依赖框架习惯。我用的方案是先把序列按块长 reshape 成 (B, num_blocks, block_len, d) 的形状在块内做普通自注意力同时把全局 token 的 KV 拼到块内注意力后面。PyTorch 自带 F.scaled_dot_product_attention 支持 mask可以把它当成一个小型 meta-batch 去做如果数据通道多还可以考虑把通道维度并进 batch 维实现 channel independence 的同时减少显存碎片。3.3 训练配置与显存优化复现时序 Transformer 的训练配置有一个通用起点我测过几次基本都能收敛优化器用 Adam初始学习率 1e-4 到 5e-4前 5 到 10 个 epoch 线性 warmup后面用 cosine 或者阶段性下降batch size 32 到 128 之间看显存和数据集调整训练轮数不用太多早停在验证集上基本 30 到 60 epoch 就有结论了。MSE 作为主损失函数评估指标用 MSE 和 MAE。三个训练细节值得单独说。第一混合精度建议优先用 bf16 而不是 fp16因为路由权重和 softmax 对精度更敏感fp16 在大专家数量下容易出现梯度溢出。第二梯度累积节省不了实验时间但能稳定训练如果你的 batch size 因为显存限制压得特别小先用梯度累积把等效 batch 撑到 64 以上再去看模型结构问题。第三Soft MoE 层不强制加辅助负载均衡损失如果你想加也可以加但加了之后路由权重的熵会明显上升反而可能抑制专家的专业化我实测下来不如不加。显存方面最大的两个头是注意力矩阵和路由矩阵。块注意力把注意力矩阵从 O(n²) 降到了 O(n·L)这个头基本解决了路由矩阵是 n×E×s如果专家数到了 16 且输入长度 720它也是一笔不小的开销。建议槽位数先从 1 开始验证集有提升再加到 2全局 token 数量保持在 2 到 4实在不够就开 gradient checkpointing只对 TST 块做 checkpoint其他层照常。3.4 与 PatchTST 公平对比的注意事项如果你是想跟 PatchTST 做对比或者复现论文表格有两点容易被人忽略。第一是 patch 长度的对齐。PatchTST 的 patch size 通常取 16而 MADEIRA 的块长本质是一个类似的概念如果你用块长 32 的 MADEIRA 去比 patch size 16 的 PatchTST比较的是两种 token 粒度结论没意义。建议先在同一输入长度下扫一遍块长16、24、32各自取最优值再比较。第二是参数量的口径。MoE 模型的总参数量包含所有专家的参数但每次前向只激活一部分直接拿总参数量对比稠密模型会显得不公平主流的做法是同时汇报参数量和 FLOPs单看 FLOPs 更有实际价值。最后多 seed 平均。时序模型在不同随机种子下的结果波动比我预想的大尤其路由初始化的位置对最终收敛影响不小。我一般跑 3 个 seed报均值和方差如果时间紧张至少跑两个 seed 确认趋势不是随机种子带来的。4. 常见问题与排查实录4.1 效果反而不如 PatchTST我在复现 MADEIRA 的第一周就遇到过这个情况块长 32、专家数 8、槽位 1号称改进了所有基线结果在 ETTh1 的 720 步预测上就是追不上 PatchTST。后来逐项排查问题出在块长和数据周期的错配上。ETTh1 是小时级数据日周期是 24 个点周周期是 168 个点我选的块长 32 既不能完整包含一个 24 点日周期又和 168 不整除路由学到的大多是一些破碎的边界。把块长改成 24 后正好对齐一个自然日周期效果立刻上来了。第二个常见问题是专家数量设置不敏感。如果槽位总数远小于块数路由退化成每个专家平均分摊所有块没有分工如果槽位总数接近 token 数Soft MoE 又退化成普通的稠密 FFN。我的经验是先把槽位总数控制在块数的四分之一到一半这个量级再看路由矩阵的熵是否处于中间状态——熵太大说明分工不明显熵太小说明路由太自信容易脆。第三个原因在输入长度上。PatchTST 在短输入96时表现很好因为 patch 天然做了局部压缩MADEIRA 的块注意力需要更长的上下文才能体现优势。如果你在预测长度较大336/720但输入长度只有 96 的配置下测试模型会有点吃亏建议对比实验统一用输入长度 336。4.2 训练不稳定或显存不够训练过程中最容易遇到的是 loss 曲线突然跳到 NaN。优先查这三处是否用了 fp16 混合精度换成 bf16路由 logits 是否没做缩放加上类似 sqrt(d) 的缩放batch size 是否过小导致梯度异常。还有一个不那么常见但真的会发生的数据集里存在 NaN 或无穷值特别是 Weather 这种传感器数据预处理时最好先做一次缺失值检查。显存不够的优化顺序我建议是先减全局 token 数从 4 降到 2→ 再减槽位数从 2 降到 1→ 再减专家数从 16 降到 8→ 还不行才用梯度 checkpointing。为什么不建议一上来就开 checkpointing因为它会显著拖慢训练等于用时间换显存排查问题时节奏会变得很慢。设置上还有一个容易被忽略的点PyTorch 默认会为每个优化器参数保有 momentum 和二阶矩MoE 层参数量大优化器状态可能比模型本身还占显存必要时对路由层用 Adafactor 或者 SGD 试试虽然收敛速度会变但显存明显缓解。4.3 推理阶段的效率与部署心得模型跑通之后如果还要上线或者做大批量回测推理效率是另一个重点。块状稀疏注意力在推理阶段的优势是 KV cache 小——每个 token 只需要缓存本块的 KV 和全局 token 的 KV不需要缓存全序列这在线性增长的序列长度下非常友好。实际部署时可以把每个专家 FFN 单独导出或量化按输入动态加载存储和内存压力都能摊薄。一个小坑是批量推演时的 padding。序列长度不是块长的整数倍时直接 reshape 会出问题我的做法是在序列末尾做零填充并把填充位置加到注意力 mask 里防止填充 token 污染全局 token 的聚合。如果你要对比推理速度务必统一 batch 大小、输入长度和输出长度只改变模型结构否则测出来的差异更多是框架层面的。最后说一点我自己实验中的体会TIME-MADEIRA 最值得借鉴的其实不是某一个组件而是把稀疏性从选择层面挪到权重层面这个思路。我在很多任务里把硬路由换成软路由之后即使模型结构没怎么改训练的稳定性和最终精度都有改善这个收益比单纯调专家数大得多。另外块长与数据周期的对齐几乎是免费的增益动手之前先去做一下周期分析比闷头扫参数高效得多。如果你打算在自己数据上试建议从块长 24、专家 8、槽位 1 这个配置出发先跑通一个短预测长度的 baseline再逐步加长预测距离。这个模型后续往多变量通道交互、非平稳序列分解这些方向扩展的空间还很大值得持续关注。