Timestep-Conditioned Transformers for Global Weather Forecasting 这个方向核心是把时间步信息作为条件输入到 Transformer 里用全局注意力做气象场演变预测。它和常见的图像生成、文本翻译类 Transformer 不一样也不只是把气象图当成普通图片去做预测。这里的难点在于大气状态包含多个变量、多个气压层还要在不同预报时效下保持物理合理性。所以下面直接按实操视角拆一拆这类模型到底在解决什么问题输入输出怎么组织时间步条件如何注入训练和评估需要注意什么以及如果要自己复现会踩到哪些坑。如果你正在关注气象AI、时序预测或者想用 Transformer 去做科学计算任务这个方向很值得看。最值得关注的不是模型堆了多少层而是“时间步条件”这个细节怎么设计它直接影响模型能不能区分6小时预报和10天预报。下面按实际落地顺序拆。1. 为什么说时间步条件是天气预报 Transformer 的关键设计1.1 全球天气预报为什么不能只依赖空间建模天气演变是典型的时空耦合过程。同一时刻的气象场确实具有很强的空间结构比如高低压系统、锋面、急流这些结构在卫星云图上能看到。但如果只把当前时刻的气象场当作输入去预测未来模型就需要在内部隐式推断大量物理过程比如大气斜压不稳定、水汽凝结、辐射传输等。让 Transformer 纯靠空间卷积或注意力去“脑补”这些过程很难在多个预报时效之间保持一致。更重要的是不同预报时效对应的预测难度和误差来源完全不同。6小时预报主要受初始条件和局地地形影响24小时预报开始依赖中尺度系统演变5天以后则更多受到大尺度环流和遥相关影响。如果模型对所有预报时效都使用同一套特征表达它在训练时会不断被不同时效的目标拉扯结果往往是短期预报不够精细长期预报不够平稳。时间步条件就是来解决这个问题的它把“当前要预报多远”这一信息显式告诉模型。这样模型可以在不同时效选择不同的特征组合方式。从实践来看这个设计有点像给模型加了一个旋钮让它在同一个参数集合下适配不同任务。1.2 Time-step Conditioning 真正解决的是什么从名称上看Timestep-Conditioned 的核心在于“conditioned”。它不是把时间步作为普通特征拼在输入里而是作为条件向量在模型内部参与特征调制。常见实现方式有很多。最简单的做法是把预报时效的小时数归一化比如除以一个固定值然后通过一个 MLP 映射成条件向量。更精细一些的做法是借鉴 Transformer 位置编码先对时间做 sinusoidal 编码再送入一个小网络。拿到条件向量后可以在 Patch Embedding 之后直接加到每个 token 特征上在每层注意力之后做 Scale-Shift 操作与注意力权重做二次融合。为什么要建议在多层注入因为只在输入层加一次网络深层的特征可能逐渐丢失时效信息。尤其在 Transformer 很深的情况下每一层都会重新混洗特征如果不持续控制模型深层就不知道自己在预测第几天效果会打折扣。每层注入会增加一点计算量但对多时效预测更稳。1.3 和传统数值预报、常规卷积预测的区别传统数值预报基于物理方程用超级计算机做逐步积分。优点是物理一致性强缺点是计算资源消耗大且许多参数化方案存在不确定性。AI 气象预报从再分析数据中学习训练阶段重计算推理阶段一次前向传播就能给出一个或多个时刻的结果速度快很多。但代价是模型学到的规律取决于训练数据分布对历史中没有出现过的极端状态容易失效。常规卷积模型通过卷积核捕获局部空间关系。问题在于局部感受野有限要覆盖数千公里的远程相关需要堆很多层。Transformer 的全局自注意力天然支持远程依赖适合捕捉大气科学里非常关心的“遥相关”也就是相隔几千公里的气象要素之间的统计关系。这也是标题里 Global Weather Forecasting 的核心原因。不过全局注意力也带来更大的计算压力。全球网格在等经纬度下分辨率越高token 数量增长越快。因此后续所有工程优化基本都围绕着“如何让全局注意力在气象网格上跑得动”展开。2. 任务定义从再分析数据到未来变量2.1 输入变量的组织方式讲模型之前先要把任务说清楚。全球天气预报模型通常使用再分析数据作为训练标签。再分析数据是把历史观测和数值模式同化得到的“最佳历史气象场”最常用的是 ERA5。实际训练时一般不会把所有变量都塞进模型因为数据量和存储成本都会爆炸。一组常见设计是选择几个关键气压层上的变量再叠加地表变量。例如气压层变量位势高度、温度、比湿、U 风、V 风气压层次850hPa、700hPa、500hPa、250hPa 等地表变量2 米温度、海平面气压、10 米 U/V 风、地表温度。将这些变量按“通道”组织起来一个样本的输入张量可以写成[历史时间步, 变量通道, 高度, 宽度]。这里的变量通道实际是“变量类型 × 气压层”展平后的结果也可能会把地表变量单独放在一起。为什么要叠这么多变量而不是只用 500hPa 位势高度因为天气系统是三维的低层水汽和风场决定了降水和高低空配置高层位势场对应大尺度环流。模型只有同时看到多层信息才能学习垂直结构。变量类型层次气象意义位势高度500hPa / 850hPa环流形势和槽脊位置温度850hPa / 1000hPa热力分层、冷暖气团比湿850hPa / 700hPa水汽输送和湿度条件U/V 风场250hPa / 850hPa急流和低层风场2 米温度等地表贴近地面的预报结果表里没有给具体层数因为不同实现不一样。原始材料没有明确指定层次落地时先按自己手上的数据源和显存条件来定。2.2 预测目标与时间步条件的关系预测目标可以定义成回归任务给定历史时刻的气象场预测未来某个时刻的气象场。这里的“某个时刻”不是固定的可能是一组不同间隔的时效。假设输入时间窗口是过去的 12 小时也就是两个 6 小时时间步那模型可以输出未来 6 小时、12 小时、18 小时甚至更远。时间步条件在这里的作用就是告诉模型当前这次预测对应的间隔是 6 小时还是 24 小时。如果不加这个条件模型只能通过输入场的演变趋势间接推断出预报时效非常不牢靠。尤其在测试阶段如果模型训练时见过的时间步分布和测试时间步不一致表现会明显变差。有的工作还会把目标时间步作为序列输入让模型按时间顺序逐步输出多帧。这种做法的优点是模型可以共享不同时效之间的特征缺点是误差会随时间步累积。Timestep-Conditioned 的设计在这里更灵活每个输出帧都携带自己的时间条件模型可以针对每个时效单独调整预测。2.3 数据预处理和标准化预处理是整个流程里最容易被忽略但最容易出问题的部分。第一分辨率。ERA5 原始网格可以是 0.25°1440×721对 Transformer 来说直接做全局注意力几乎不可能。常见做法是先降采样到 1.40625° 或 2.8125°让网格变成 128×64 或 64×32这样 token 数量可控。也可以不变分辨率但使用 Patch 化来减少 token 数量。第二标准化。每个变量、每个气压层的物理量纲都不一样必须分别计算均值和标准差把输入输出都归一化到零均值单位方差。如果不做这一步500hPa 位势高度的数值很大2 米温度的数值也很突出模型会优先拟合数值大的变量其他变量被压制。第三气候态去除。气象学里评估 ACC 时会用到距离气候态的异常场。如果训练时没有把气候态从数据里去掉模型很容易学到“直接输出气候态”这种偷懒策略这样异常场不准确。更稳的做法是在标准化的同时保存每个网格点在每个月份和时次的气候态供训练和评估使用。第四数据泄漏。训练、验证、测试的时间段必须严格分开。如果测试期的数据参与过标准化统计或者参与过训练数据采样评估结果会虚高。常见划分是训练 1979—2015验证 2016测试 2017—2018但具体以数据版本为准。3. 模型架构中的关键设计3.1 时间步条件如何注入这一节给出一个可以借鉴的最小模块。时间步条件最核心的计算是把一维时间信息映射成与模型维度匹配的条件向量。伪代码如下import torch import torch.nn as nn import math class TimeStepFourierEmbedding(nn.Module): def __init__(self, dim): super().__init__() half dim // 2 self.freq nn.Parameter(torch.linspace(0.0, 1.0, half), requires_gradFalse) self.net nn.Sequential( nn.Linear(dim, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim) ) def forward(self, hours): # hours: [B, 1]已经归一化到 [0, 1] 或相近范围 t hours * self.freq * 2.0 * math.pi emb torch.cat([torch.sin(t), torch.cos(t)], dim-1) return self.net(emb)拿到条件向量后常用方式是在每个 Transformer Block 里做 Scale-Shiftclass ConditionedBlock(nn.Module): def __init__(self, block, dim): super().__init__() self.block block self.shift nn.Linear(dim, dim) self.scale nn.Linear(dim, dim) def forward(self, x, cond): # x: [B, L, D] scale self.scale(cond).unsqueeze(1) shift self.shift(cond).unsqueeze(1) x self.block(x) * scale shift return x我一般会在每个 Block 内部做一次 Scale-Shift而不是只在输入层加一次。这样做的原因是条件向量要能对每层特征产生作用网络越深时效信息越容易被后续特征覆盖掉。从实际使用来看加入条件注入确实会让模型在训练时更稳定。如果训练时发现短期时效和长期时效的 loss 相互干扰优先检查条件注入是否只在浅层以及条件向量的维度是否过小。3.2 空间 Patch 化与全局注意力Transformer 的 Self-Attention 复杂度是 O(n²)n 是 token 数量。全球网格如果直接用每个网格点作为 tokenn 128×64 8192还能接受如果是 0.25° 分辨率n 1440×721 ≈ 103 万直接做注意力完全不现实。所以常见的做法是“Patch 化”。把空间网格划分成 P×P 的小块例如 4×4 或 8×8每个 patch 内部展平后投影成一个特征向量。这样 token 数量可以降到原来的 1/P²同时保留局部空间相关性。气象场的相邻网格之间本身就高度相关patch 化并不会丢失太多信息。另一类做法是把经纬度两个方向分开做注意力比如先沿纬度再沿经度复杂度可以降到 O(n^{1.5}) 或线性。这类变体也适合气象数据但实现起来更复杂。如果以标题中的 Transformers 为主先按 patch 化理解更合适。Patching 还有一个好处可以把不同分辨率的数据统一到相近的 token 数量。比如 128×64 网格、patch4 时得到 32×16512 个 token64×32、patch2 时也是 32×16512 个 token。这样模型主干可以复用便于做多分辨率训练。3.3 损失函数权重、掩码与多步损失训练阶段需要指定模型优化的目标。最常见的损失是加权均方根误差RMSE因为天气预测的误差通常看空间每个点的偏差。由于等经纬度网格在高纬度的网格面积变小如果直接平均每个网格的误差会高估高纬度区域的影响。所以计算 loss 时通常按纬度的余弦值加权也就是 cos(latitude)。这和高纬地区的实际面积比例一致。另外要处理 mask。地球网格在极点附近或者某些区域可能存在无效值尤其是在重采样后插值边界附近。训练时如果 mask 不设置模型会拼命去拟合无效区域造成梯度异常。建议在 loss 计算时对 mask 区域乘以 0。多步预测时如果模型一次输出多个未来时刻可以对每个时刻分别计算 loss再取平均。这里要注意不同时效的误差尺度不同6 小时误差小5 天误差大直接平均会让模型更关注长时效挤压短期预报的优化空间。可以给每个时效设置权重但这属于超参数需要实验确定。3.4 训练策略与评估协议数据采样建议随机化。不要把连续年份按顺序喂给模型否则模型会学到时间不平移的静态分布。每次训练迭代随机选一个起始时间再截取历史窗口和目标时间可以显著增加样本多样性。训练阶段可以使用 AdamW 优化器配合余弦退火学习率。由于气象场是大尺度空间数据batch size 不宜过小否则梯度的空间统计性不稳定。但显存有限时不要硬撑可以使用梯度累积。评估协议要提前确定。常见做法是从测试期某一天开始以 6 小时为间隔每次用最近 12 小时或 24 小时数据作为输入预测未来 6 小时到 120 小时甚至 240 小时。然后画 RMSE 和 ACC 随预报时效变化的曲线。曲线中应该包含 Persistence 基线和 Climatology 基线用来判断模型是否真的学到东西而不只是跟着气候态走。4. 实操复现环境、训练和评估4.1 环境与依赖复现这类模型建议使用 Linux 系统尤其是训练全球分辨率数据时Linux 在 IO 和分布式方面更顺手。Windows 也能跑通小规模实验但处理 NetCDF 和多卡训练会更折腾。GPU 方面如果只做低分辨率验证显存 11GB 也能起步但需要把 batch size 和模型维度调小。如果目标是在 128×64 或更高分辨率下训练较深 Transformer建议至少 24GB 显存。使用混合精度和 gradient checkpointing 可以缓解显存压力但代码里要提前规划。常见 Python 依赖依赖用途torch模型训练和推理xarray读取气象 NetCDF 数据netCDF4NetCDF 格式支持numpy数值计算einops张量维度重排dask大数据并行加载matplotlib画评估曲线4.2 数据准备流程如果是从零开始不要一上来下载全量 ERA5。数据量太大光读写就会卡住。先按变量和层次筛选子集降到需要的分辨率再按年份切块。一个建议流程在数据服务页面选择要用的变量和层次。把时间范围按年份分段下载避免单文件过大。用 xarray 打开并重采样到目标分辨率。剔除明显异常或无效区域生成 mask。计算每个变量的 mean/std保存为配置文件。生成训练样本索引每个样本记录起始时间、历史窗口长度、目标窗口长度。按时间段划分 train/val/test。很多早期失败都是因为数据预处理没做对。比如变量单位变了标准化统计范围不对mask 没有对齐导致模型训练到一半 loss 突然变成 nan。4.3 训练和评估示例训练循环的骨架可以写成for epoch in range(epochs): model.train() for x, y, t in train_loader: x, y, t x.to(device), y.to(device), t.to(device) pred model(x, t) loss latitude_weighted_rmse(pred, y, mask) loss.backward() optimizer.step() optimizer.zero_grad()评估时逐时效计算model.eval() with torch.no_grad(): pred model(x, timesteps) rmse latitude_weighted_rmse(pred, y, mask) acc latitude_weighted_acc(pred, y, climatology)其中t是一个张量里面的值代表当前样本对应的预报时效例如[6, 12, 24]。这里特别注意每个 batch 里的样本如果预报时效不同t的维度要能对应上。4.4 评估指标怎么判断RMSE、ACC 与纬度加权RMSE 的公式可以理解为对预测误差的平方做加权平均后开方。权重使用 cos(latitude)每个网格的贡献更接近真实面积。ACC 是异常相关系数通常计算预测异常场和真实异常场在空间上的相关。它比 RMSE 更能反映“空间形态”是否预测对比如槽脊位置提前或滞后。业务上通常两者都看RMSE 看幅度ACC 看形态。如果评估结果出现 RMSE 不错但 ACC 很低说明预测场和真实场形态错位可能是时间步条件没有让模型学会随预报时效调整也可能是预测结果被过度平滑。如果 ACC 在较长时效内仍能保持在一个可以参考的区间说明模型在中期尺度上有作用。具体阈值和业务区域强相关不要直接拿别人论文里的单点数值当通用标准。注意判断模型好坏时不要只盯一个时效的 RMSE要看多条预报曲线整体趋势。如果 6 小时很好但 72 小时突然爆炸模型可能存在累积误差问题。5. 结果分析和边界条件5.1 短期预报和中期预报表现全局注意力让模型能看见大半个地球的同步状态这对中期预报是有利的。比如遥相关型可以在一次注意力计算中被捕捉到不需要像卷积那样堆很多层。从实践经验看这类模型在中长期尺度上的场平滑性通常比纯卷积模型好不会出现明显的格子状噪声。但在短期预报上局地地形、海陆差异、对流系统这些过程对空间分辨率和物理约束更敏感。模型如果只在几个气压层上训练很多与降水相关的变量缺失无法精准预测短时强降水。要提升短期效果需要在输入里加入更多地表和近地面变量或者使用更高分辨率数据。5. 时间步条件Transformer：全球天气预报的关键设计与实操指南
Timestep-Conditioned Transformers for Global Weather Forecasting 这个方向核心是把时间步信息作为条件输入到 Transformer 里用全局注意力做气象场演变预测。它和常见的图像生成、文本翻译类 Transformer 不一样也不只是把气象图当成普通图片去做预测。这里的难点在于大气状态包含多个变量、多个气压层还要在不同预报时效下保持物理合理性。所以下面直接按实操视角拆一拆这类模型到底在解决什么问题输入输出怎么组织时间步条件如何注入训练和评估需要注意什么以及如果要自己复现会踩到哪些坑。如果你正在关注气象AI、时序预测或者想用 Transformer 去做科学计算任务这个方向很值得看。最值得关注的不是模型堆了多少层而是“时间步条件”这个细节怎么设计它直接影响模型能不能区分6小时预报和10天预报。下面按实际落地顺序拆。1. 为什么说时间步条件是天气预报 Transformer 的关键设计1.1 全球天气预报为什么不能只依赖空间建模天气演变是典型的时空耦合过程。同一时刻的气象场确实具有很强的空间结构比如高低压系统、锋面、急流这些结构在卫星云图上能看到。但如果只把当前时刻的气象场当作输入去预测未来模型就需要在内部隐式推断大量物理过程比如大气斜压不稳定、水汽凝结、辐射传输等。让 Transformer 纯靠空间卷积或注意力去“脑补”这些过程很难在多个预报时效之间保持一致。更重要的是不同预报时效对应的预测难度和误差来源完全不同。6小时预报主要受初始条件和局地地形影响24小时预报开始依赖中尺度系统演变5天以后则更多受到大尺度环流和遥相关影响。如果模型对所有预报时效都使用同一套特征表达它在训练时会不断被不同时效的目标拉扯结果往往是短期预报不够精细长期预报不够平稳。时间步条件就是来解决这个问题的它把“当前要预报多远”这一信息显式告诉模型。这样模型可以在不同时效选择不同的特征组合方式。从实践来看这个设计有点像给模型加了一个旋钮让它在同一个参数集合下适配不同任务。1.2 Time-step Conditioning 真正解决的是什么从名称上看Timestep-Conditioned 的核心在于“conditioned”。它不是把时间步作为普通特征拼在输入里而是作为条件向量在模型内部参与特征调制。常见实现方式有很多。最简单的做法是把预报时效的小时数归一化比如除以一个固定值然后通过一个 MLP 映射成条件向量。更精细一些的做法是借鉴 Transformer 位置编码先对时间做 sinusoidal 编码再送入一个小网络。拿到条件向量后可以在 Patch Embedding 之后直接加到每个 token 特征上在每层注意力之后做 Scale-Shift 操作与注意力权重做二次融合。为什么要建议在多层注入因为只在输入层加一次网络深层的特征可能逐渐丢失时效信息。尤其在 Transformer 很深的情况下每一层都会重新混洗特征如果不持续控制模型深层就不知道自己在预测第几天效果会打折扣。每层注入会增加一点计算量但对多时效预测更稳。1.3 和传统数值预报、常规卷积预测的区别传统数值预报基于物理方程用超级计算机做逐步积分。优点是物理一致性强缺点是计算资源消耗大且许多参数化方案存在不确定性。AI 气象预报从再分析数据中学习训练阶段重计算推理阶段一次前向传播就能给出一个或多个时刻的结果速度快很多。但代价是模型学到的规律取决于训练数据分布对历史中没有出现过的极端状态容易失效。常规卷积模型通过卷积核捕获局部空间关系。问题在于局部感受野有限要覆盖数千公里的远程相关需要堆很多层。Transformer 的全局自注意力天然支持远程依赖适合捕捉大气科学里非常关心的“遥相关”也就是相隔几千公里的气象要素之间的统计关系。这也是标题里 Global Weather Forecasting 的核心原因。不过全局注意力也带来更大的计算压力。全球网格在等经纬度下分辨率越高token 数量增长越快。因此后续所有工程优化基本都围绕着“如何让全局注意力在气象网格上跑得动”展开。2. 任务定义从再分析数据到未来变量2.1 输入变量的组织方式讲模型之前先要把任务说清楚。全球天气预报模型通常使用再分析数据作为训练标签。再分析数据是把历史观测和数值模式同化得到的“最佳历史气象场”最常用的是 ERA5。实际训练时一般不会把所有变量都塞进模型因为数据量和存储成本都会爆炸。一组常见设计是选择几个关键气压层上的变量再叠加地表变量。例如气压层变量位势高度、温度、比湿、U 风、V 风气压层次850hPa、700hPa、500hPa、250hPa 等地表变量2 米温度、海平面气压、10 米 U/V 风、地表温度。将这些变量按“通道”组织起来一个样本的输入张量可以写成[历史时间步, 变量通道, 高度, 宽度]。这里的变量通道实际是“变量类型 × 气压层”展平后的结果也可能会把地表变量单独放在一起。为什么要叠这么多变量而不是只用 500hPa 位势高度因为天气系统是三维的低层水汽和风场决定了降水和高低空配置高层位势场对应大尺度环流。模型只有同时看到多层信息才能学习垂直结构。变量类型层次气象意义位势高度500hPa / 850hPa环流形势和槽脊位置温度850hPa / 1000hPa热力分层、冷暖气团比湿850hPa / 700hPa水汽输送和湿度条件U/V 风场250hPa / 850hPa急流和低层风场2 米温度等地表贴近地面的预报结果表里没有给具体层数因为不同实现不一样。原始材料没有明确指定层次落地时先按自己手上的数据源和显存条件来定。2.2 预测目标与时间步条件的关系预测目标可以定义成回归任务给定历史时刻的气象场预测未来某个时刻的气象场。这里的“某个时刻”不是固定的可能是一组不同间隔的时效。假设输入时间窗口是过去的 12 小时也就是两个 6 小时时间步那模型可以输出未来 6 小时、12 小时、18 小时甚至更远。时间步条件在这里的作用就是告诉模型当前这次预测对应的间隔是 6 小时还是 24 小时。如果不加这个条件模型只能通过输入场的演变趋势间接推断出预报时效非常不牢靠。尤其在测试阶段如果模型训练时见过的时间步分布和测试时间步不一致表现会明显变差。有的工作还会把目标时间步作为序列输入让模型按时间顺序逐步输出多帧。这种做法的优点是模型可以共享不同时效之间的特征缺点是误差会随时间步累积。Timestep-Conditioned 的设计在这里更灵活每个输出帧都携带自己的时间条件模型可以针对每个时效单独调整预测。2.3 数据预处理和标准化预处理是整个流程里最容易被忽略但最容易出问题的部分。第一分辨率。ERA5 原始网格可以是 0.25°1440×721对 Transformer 来说直接做全局注意力几乎不可能。常见做法是先降采样到 1.40625° 或 2.8125°让网格变成 128×64 或 64×32这样 token 数量可控。也可以不变分辨率但使用 Patch 化来减少 token 数量。第二标准化。每个变量、每个气压层的物理量纲都不一样必须分别计算均值和标准差把输入输出都归一化到零均值单位方差。如果不做这一步500hPa 位势高度的数值很大2 米温度的数值也很突出模型会优先拟合数值大的变量其他变量被压制。第三气候态去除。气象学里评估 ACC 时会用到距离气候态的异常场。如果训练时没有把气候态从数据里去掉模型很容易学到“直接输出气候态”这种偷懒策略这样异常场不准确。更稳的做法是在标准化的同时保存每个网格点在每个月份和时次的气候态供训练和评估使用。第四数据泄漏。训练、验证、测试的时间段必须严格分开。如果测试期的数据参与过标准化统计或者参与过训练数据采样评估结果会虚高。常见划分是训练 1979—2015验证 2016测试 2017—2018但具体以数据版本为准。3. 模型架构中的关键设计3.1 时间步条件如何注入这一节给出一个可以借鉴的最小模块。时间步条件最核心的计算是把一维时间信息映射成与模型维度匹配的条件向量。伪代码如下import torch import torch.nn as nn import math class TimeStepFourierEmbedding(nn.Module): def __init__(self, dim): super().__init__() half dim // 2 self.freq nn.Parameter(torch.linspace(0.0, 1.0, half), requires_gradFalse) self.net nn.Sequential( nn.Linear(dim, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim) ) def forward(self, hours): # hours: [B, 1]已经归一化到 [0, 1] 或相近范围 t hours * self.freq * 2.0 * math.pi emb torch.cat([torch.sin(t), torch.cos(t)], dim-1) return self.net(emb)拿到条件向量后常用方式是在每个 Transformer Block 里做 Scale-Shiftclass ConditionedBlock(nn.Module): def __init__(self, block, dim): super().__init__() self.block block self.shift nn.Linear(dim, dim) self.scale nn.Linear(dim, dim) def forward(self, x, cond): # x: [B, L, D] scale self.scale(cond).unsqueeze(1) shift self.shift(cond).unsqueeze(1) x self.block(x) * scale shift return x我一般会在每个 Block 内部做一次 Scale-Shift而不是只在输入层加一次。这样做的原因是条件向量要能对每层特征产生作用网络越深时效信息越容易被后续特征覆盖掉。从实际使用来看加入条件注入确实会让模型在训练时更稳定。如果训练时发现短期时效和长期时效的 loss 相互干扰优先检查条件注入是否只在浅层以及条件向量的维度是否过小。3.2 空间 Patch 化与全局注意力Transformer 的 Self-Attention 复杂度是 O(n²)n 是 token 数量。全球网格如果直接用每个网格点作为 tokenn 128×64 8192还能接受如果是 0.25° 分辨率n 1440×721 ≈ 103 万直接做注意力完全不现实。所以常见的做法是“Patch 化”。把空间网格划分成 P×P 的小块例如 4×4 或 8×8每个 patch 内部展平后投影成一个特征向量。这样 token 数量可以降到原来的 1/P²同时保留局部空间相关性。气象场的相邻网格之间本身就高度相关patch 化并不会丢失太多信息。另一类做法是把经纬度两个方向分开做注意力比如先沿纬度再沿经度复杂度可以降到 O(n^{1.5}) 或线性。这类变体也适合气象数据但实现起来更复杂。如果以标题中的 Transformers 为主先按 patch 化理解更合适。Patching 还有一个好处可以把不同分辨率的数据统一到相近的 token 数量。比如 128×64 网格、patch4 时得到 32×16512 个 token64×32、patch2 时也是 32×16512 个 token。这样模型主干可以复用便于做多分辨率训练。3.3 损失函数权重、掩码与多步损失训练阶段需要指定模型优化的目标。最常见的损失是加权均方根误差RMSE因为天气预测的误差通常看空间每个点的偏差。由于等经纬度网格在高纬度的网格面积变小如果直接平均每个网格的误差会高估高纬度区域的影响。所以计算 loss 时通常按纬度的余弦值加权也就是 cos(latitude)。这和高纬地区的实际面积比例一致。另外要处理 mask。地球网格在极点附近或者某些区域可能存在无效值尤其是在重采样后插值边界附近。训练时如果 mask 不设置模型会拼命去拟合无效区域造成梯度异常。建议在 loss 计算时对 mask 区域乘以 0。多步预测时如果模型一次输出多个未来时刻可以对每个时刻分别计算 loss再取平均。这里要注意不同时效的误差尺度不同6 小时误差小5 天误差大直接平均会让模型更关注长时效挤压短期预报的优化空间。可以给每个时效设置权重但这属于超参数需要实验确定。3.4 训练策略与评估协议数据采样建议随机化。不要把连续年份按顺序喂给模型否则模型会学到时间不平移的静态分布。每次训练迭代随机选一个起始时间再截取历史窗口和目标时间可以显著增加样本多样性。训练阶段可以使用 AdamW 优化器配合余弦退火学习率。由于气象场是大尺度空间数据batch size 不宜过小否则梯度的空间统计性不稳定。但显存有限时不要硬撑可以使用梯度累积。评估协议要提前确定。常见做法是从测试期某一天开始以 6 小时为间隔每次用最近 12 小时或 24 小时数据作为输入预测未来 6 小时到 120 小时甚至 240 小时。然后画 RMSE 和 ACC 随预报时效变化的曲线。曲线中应该包含 Persistence 基线和 Climatology 基线用来判断模型是否真的学到东西而不只是跟着气候态走。4. 实操复现环境、训练和评估4.1 环境与依赖复现这类模型建议使用 Linux 系统尤其是训练全球分辨率数据时Linux 在 IO 和分布式方面更顺手。Windows 也能跑通小规模实验但处理 NetCDF 和多卡训练会更折腾。GPU 方面如果只做低分辨率验证显存 11GB 也能起步但需要把 batch size 和模型维度调小。如果目标是在 128×64 或更高分辨率下训练较深 Transformer建议至少 24GB 显存。使用混合精度和 gradient checkpointing 可以缓解显存压力但代码里要提前规划。常见 Python 依赖依赖用途torch模型训练和推理xarray读取气象 NetCDF 数据netCDF4NetCDF 格式支持numpy数值计算einops张量维度重排dask大数据并行加载matplotlib画评估曲线4.2 数据准备流程如果是从零开始不要一上来下载全量 ERA5。数据量太大光读写就会卡住。先按变量和层次筛选子集降到需要的分辨率再按年份切块。一个建议流程在数据服务页面选择要用的变量和层次。把时间范围按年份分段下载避免单文件过大。用 xarray 打开并重采样到目标分辨率。剔除明显异常或无效区域生成 mask。计算每个变量的 mean/std保存为配置文件。生成训练样本索引每个样本记录起始时间、历史窗口长度、目标窗口长度。按时间段划分 train/val/test。很多早期失败都是因为数据预处理没做对。比如变量单位变了标准化统计范围不对mask 没有对齐导致模型训练到一半 loss 突然变成 nan。4.3 训练和评估示例训练循环的骨架可以写成for epoch in range(epochs): model.train() for x, y, t in train_loader: x, y, t x.to(device), y.to(device), t.to(device) pred model(x, t) loss latitude_weighted_rmse(pred, y, mask) loss.backward() optimizer.step() optimizer.zero_grad()评估时逐时效计算model.eval() with torch.no_grad(): pred model(x, timesteps) rmse latitude_weighted_rmse(pred, y, mask) acc latitude_weighted_acc(pred, y, climatology)其中t是一个张量里面的值代表当前样本对应的预报时效例如[6, 12, 24]。这里特别注意每个 batch 里的样本如果预报时效不同t的维度要能对应上。4.4 评估指标怎么判断RMSE、ACC 与纬度加权RMSE 的公式可以理解为对预测误差的平方做加权平均后开方。权重使用 cos(latitude)每个网格的贡献更接近真实面积。ACC 是异常相关系数通常计算预测异常场和真实异常场在空间上的相关。它比 RMSE 更能反映“空间形态”是否预测对比如槽脊位置提前或滞后。业务上通常两者都看RMSE 看幅度ACC 看形态。如果评估结果出现 RMSE 不错但 ACC 很低说明预测场和真实场形态错位可能是时间步条件没有让模型学会随预报时效调整也可能是预测结果被过度平滑。如果 ACC 在较长时效内仍能保持在一个可以参考的区间说明模型在中期尺度上有作用。具体阈值和业务区域强相关不要直接拿别人论文里的单点数值当通用标准。注意判断模型好坏时不要只盯一个时效的 RMSE要看多条预报曲线整体趋势。如果 6 小时很好但 72 小时突然爆炸模型可能存在累积误差问题。5. 结果分析和边界条件5.1 短期预报和中期预报表现全局注意力让模型能看见大半个地球的同步状态这对中期预报是有利的。比如遥相关型可以在一次注意力计算中被捕捉到不需要像卷积那样堆很多层。从实践经验看这类模型在中长期尺度上的场平滑性通常比纯卷积模型好不会出现明显的格子状噪声。但在短期预报上局地地形、海陆差异、对流系统这些过程对空间分辨率和物理约束更敏感。模型如果只在几个气压层上训练很多与降水相关的变量缺失无法精准预测短时强降水。要提升短期效果需要在输入里加入更多地表和近地面变量或者使用更高分辨率数据。5.