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_grad=False) 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-Shift：
class 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 数量。全球网格如果直接用每个网格点作为 token，n = 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 网格、patch=4 时得到 32×16=512 个 token；64×32、patch=2 时也是 32×16=512 个 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 数据 |
| netCDF4 | NetCDF 格式支持 |
| 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 短期预报和中期预报表现
全局注意力让模型能看见大半个地球的同步状态，这对中期预报是有利的。比如遥相关型可以在一次注意力计算中被捕捉到，不需要像卷积那样堆很多层。从实践经验看，这类模型在中长期尺度上的场平滑性通常比纯卷积模型好，不会出现明显的格子状噪声。
但在短期预报上，局地地形、海陆差异、对流系统这些过程对空间分辨率和物理约束更敏感。模型如果只在几个气压层上训练，很多与降水相关的变量缺失，无法精准预测短时强降水。要提升短期效果，需要在输入里加入更多地表和近地面变量，或者使用更高分辨率数据。
