間步條件Transformer：?jiǎn)文Ｐ蛯?shí)現(xiàn)靈活多時(shí)效AI天氣預(yù)報(bào))
全球天氣預(yù)報(bào)正在經(jīng)歷一次范式切換越來(lái)越多的研究不再把大氣運(yùn)動(dòng)看作必須用偏微分方程求解的物理過(guò)程而是把它當(dāng)作一個(gè)海量時(shí)空序列預(yù)測(cè)問(wèn)題直接交給 Transformer 這類模型去學(xué)習(xí)。Timestep-Conditioned Transformers for Global Weather Forecasting 這篇論文的核心思路用一句話概括就是把預(yù)報(bào)時(shí)效timestep從模型的隱含信息變成顯式條件輸入讓同一個(gè) Transformer 模型可以靈活預(yù)測(cè)不同長(zhǎng)度的預(yù)報(bào)時(shí)效而不是為每一檔時(shí)效單獨(dú)訓(xùn)練一個(gè)模型也不是只能靠自回歸一步一步往外推。這個(gè)設(shè)計(jì)看起來(lái)只是給模型多加了幾個(gè)輸入維度實(shí)際上改變的是模型處理“時(shí)間”的方式。這篇文章會(huì)拆解它的技術(shù)原理分析它解決了什么核心痛點(diǎn)并給出一個(gè)可運(yùn)行的簡(jiǎn)化示例幫助你理解如何在 Transformer 中注入時(shí)間步條件。如果你正在做時(shí)間序列預(yù)測(cè)、AI 天氣預(yù)報(bào)模型或者只是想知道這類方法與 GraphCast、盤古氣象大模型的差別這篇文章值得讀完。1. 這篇文章真正要解決的問(wèn)題做過(guò)時(shí)間序列預(yù)測(cè)的人都會(huì)遇到一個(gè)基礎(chǔ)問(wèn)題預(yù)測(cè)目標(biāo)離當(dāng)前時(shí)刻越遠(yuǎn)建模難度越大。在 AI 天氣預(yù)報(bào)場(chǎng)景里這個(gè)問(wèn)題會(huì)被放大。全球天氣預(yù)報(bào)的常規(guī)預(yù)測(cè)時(shí)效通常包括 6 小時(shí)、12 小時(shí)、24 小時(shí)、甚至 240 小時(shí)10 天。面對(duì)這些不同時(shí)效傳統(tǒng)做法大致有兩種訓(xùn)練一個(gè)模型然后用它自回歸地迭代預(yù)測(cè)先預(yù)測(cè) 6 小時(shí)后的場(chǎng)再把這個(gè)預(yù)測(cè)結(jié)果當(dāng)作輸入繼續(xù)預(yù)測(cè)下一個(gè) 6 小時(shí)重復(fù)多次。因?yàn)樽曰貧w會(huì)累積誤差索性為每個(gè)預(yù)報(bào)時(shí)效單獨(dú)訓(xùn)練一個(gè)模型。前者的問(wèn)題是誤差會(huì)隨著迭代次數(shù)逐漸累計(jì)越往后預(yù)測(cè)越失真后者的問(wèn)題是訓(xùn)練和部署成本成倍增長(zhǎng)而且不同時(shí)效的模型各自獨(dú)立無(wú)法共享信息。Timestep-Conditioned Transformer 提供的是第三種思路把“我要預(yù)測(cè)多久以后”作為條件輸入模型讓同一個(gè)模型學(xué)會(huì)在不同時(shí)間尺度下的預(yù)測(cè)模式。模型收到一個(gè)輸入場(chǎng)同時(shí)收到一個(gè)表示“目標(biāo)時(shí)刻距離現(xiàn)在多久”的信號(hào)然后直接輸出對(duì)應(yīng)時(shí)刻的預(yù)測(cè)場(chǎng)。如果你正在做以下任何一件事這篇文章都值得看研究 AI 天氣預(yù)報(bào)模型想知道除了 GraphCast、盤古之外還有哪些值得關(guān)注的設(shè)計(jì)方向。做通用時(shí)序預(yù)測(cè)想把時(shí)間步條件化的思路遷移到自己的任務(wù)中。做 Transformer 架構(gòu)改進(jìn)想了解條件注入機(jī)制的實(shí)現(xiàn)方式。2. 核心概念什么是 Timestep-Conditioned TransformerTransformer 這類模型本身是“無(wú)狀態(tài)”的它不會(huì)天然知道“當(dāng)前是哪個(gè)時(shí)刻”。我們通?？?jī)煞N方式把時(shí)間信息告訴它位置編碼Positional Encoding告訴模型序列中每個(gè) token 的順序。時(shí)間步條件Timestep Conditioning告訴模型“預(yù)測(cè)目標(biāo)在時(shí)間軸上的位置”。很多讀者會(huì)把這兩個(gè)概念混淆。這里用一個(gè)類比說(shuō)明位置編碼相當(dāng)于文章里的“頁(yè)碼”。它讓模型知道這些 token 的先后關(guān)系。時(shí)間步條件則相當(dāng)于一個(gè)額外的問(wèn)題“下周二的天氣情況如何”這個(gè)問(wèn)題不是通過(guò)頁(yè)碼表達(dá)的而是作為任務(wù)本身的屬性輸入。在天氣預(yù)報(bào)模型里時(shí)間步條件常見(jiàn)的實(shí)現(xiàn)方式是這樣的將預(yù)報(bào)時(shí)效數(shù)值例如 24 小時(shí)輸入一個(gè)編碼器。編碼器把它映射成一個(gè)高維向量。這個(gè)向量通過(guò)加法、自適應(yīng)歸一化、交叉注意力等方式注入 Transformer 的每一層。設(shè)計(jì)上的關(guān)鍵判斷是時(shí)間不再只是序列的自然順序而是預(yù)測(cè)任務(wù)的條件變量。這讓同一個(gè)模型能夠針對(duì)不同的預(yù)報(bào)時(shí)效給出不同強(qiáng)度的響應(yīng)。例如模型可以學(xué)習(xí)到“短期預(yù)報(bào)更多依賴當(dāng)前場(chǎng)的細(xì)節(jié)”而“長(zhǎng)期預(yù)報(bào)需要更大尺度的空間模式”。概念作用典型實(shí)現(xiàn)位置編碼描述 token 在空間或時(shí)間序列中的位置正弦編碼、可學(xué)習(xí)位置嵌入時(shí)間步條件描述預(yù)測(cè)目標(biāo)距當(dāng)前時(shí)刻的間隔正弦編碼 MLP、AdaLN、Cross-Attention上下文窗口描述模型可見(jiàn)的歷史觀測(cè)范圍輸入序列長(zhǎng)度設(shè)計(jì)兩者的共同點(diǎn)是都離不開(kāi)“編碼”這一步但語(yǔ)義完全不同。時(shí)間步條件回答的是“目標(biāo)是什么”位置編碼回答的是“我在哪里”。3. 全球天氣預(yù)報(bào)的難點(diǎn)與 AI 模型的切入方式3.1 任務(wù)本身的復(fù)雜度全球天氣預(yù)報(bào)的數(shù)據(jù)通常是多變量、多壓力層的三維網(wǎng)格場(chǎng)。常見(jiàn)數(shù)據(jù)源是 ERA5 再分析數(shù)據(jù)集覆蓋多個(gè)氣壓層每層包含溫度、濕度、風(fēng)場(chǎng)等多個(gè)變量。在 0.25° 分辨率下全球網(wǎng)格點(diǎn)數(shù)規(guī)模在百萬(wàn)級(jí)別。這種數(shù)據(jù)有三個(gè)特點(diǎn)空間相關(guān)性極強(qiáng)某一區(qū)域的天氣狀態(tài)受周邊區(qū)域甚至全球環(huán)流影響。時(shí)間演化復(fù)雜大氣運(yùn)動(dòng)跨尺度、跨層次耦合。變量之間非線性關(guān)系明顯溫度、氣壓、濕度不是幾個(gè)簡(jiǎn)單公式能描述的。3.2 數(shù)值天氣預(yù)報(bào)的瓶頸傳統(tǒng)數(shù)值天氣預(yù)報(bào)NWP通過(guò)求解流體力學(xué)和熱力學(xué)方程組來(lái)預(yù)測(cè)天氣。它的優(yōu)勢(shì)是可解釋性強(qiáng)但計(jì)算成本極高。一次全球 10 天預(yù)報(bào)需要在超算上運(yùn)行很長(zhǎng)時(shí)間這限制了高分辨率、高頻次的更新。3.3 AI 模型的替代邏輯AI 天氣預(yù)報(bào)模型把問(wèn)題簡(jiǎn)化成一個(gè)監(jiān)督學(xué)習(xí)任務(wù)輸入過(guò)去若干個(gè)時(shí)刻的全球氣象場(chǎng) 輸出未來(lái)某個(gè)時(shí)刻的氣象場(chǎng)訓(xùn)練過(guò)程中模型直接學(xué)習(xí)從觀測(cè)場(chǎng)到未來(lái)場(chǎng)之間的映射關(guān)系。推理時(shí)模型只需要一次前向計(jì)算不需要迭代求解方程組。因此AI 模型在推理速度上有數(shù)量級(jí)優(yōu)勢(shì)。在這個(gè)背景下如何處理“未來(lái)某個(gè)時(shí)刻”這個(gè)目標(biāo)就成為一個(gè)核心設(shè)計(jì)問(wèn)題。Timestep-Conditioned Transformer 的切入點(diǎn)正是在這里。4. 技術(shù)原理解剖時(shí)間步條件如何注入 Transformer4.1 時(shí)間步信息的表示時(shí)間步條件的第一步是把數(shù)值型的時(shí)間信息轉(zhuǎn)換成向量。參考擴(kuò)散模型里的時(shí)間嵌入方法通常使用正弦位置編碼加 MLP將時(shí)間值例如 24轉(zhuǎn)換為一組不同頻率的正弦和余弦信號(hào)。將這些信號(hào)拼接后送入一個(gè)小型 MLP。MLP 輸出一個(gè)固定維度的條件向量。這樣可以保證時(shí)間步之間的相似性在向量空間中有平滑過(guò)渡。例如24 小時(shí)與 30 小時(shí)的向量會(huì)比較接近而與 120 小時(shí)的向量差異較大。這種平滑性對(duì)模型訓(xùn)練非常重要。4.2 條件注入的三種主要方式拿到時(shí)間步向量之后需要把它注入到 Transformer 中。常見(jiàn)方式有三種它們的表達(dá)能力從弱到強(qiáng)注入方式原理優(yōu)點(diǎn)缺點(diǎn)向量相加將時(shí)間步向量加到序列 token 的 embedding 上實(shí)現(xiàn)簡(jiǎn)單條件信息會(huì)在后續(xù)層中被逐漸稀釋AdaLN用時(shí)間步向量預(yù)測(cè) LayerNorm 的縮放和偏移參數(shù)每層都能直接感知條件表達(dá)能力強(qiáng)參數(shù)計(jì)算復(fù)雜一些Cross-Attention將時(shí)間步向量作為 query圖像特征作為 key/value實(shí)現(xiàn)更靈活計(jì)算量更大論文中提到的 Timestep-Conditioned Transformer 本質(zhì)上屬于第一種或第二種思路的變體。它通常會(huì)在每個(gè) Transformer Block 內(nèi)部注入條件確保時(shí)間信息在深層網(wǎng)絡(luò)中不會(huì)被遺忘。4.3 與位置編碼的配合方式在實(shí)際模型中空間位置編碼和時(shí)間步條件同時(shí)存在空間位置編碼負(fù)責(zé)描述每個(gè)網(wǎng)格點(diǎn)在球面上的相對(duì)位置。時(shí)間步條件負(fù)責(zé)描述預(yù)測(cè)目標(biāo)在時(shí)間軸上的位置。兩者互補(bǔ)共同構(gòu)成一個(gè)完整的時(shí)空上下文。這里真正容易踩坑的地方是很多人會(huì)把時(shí)間步條件直接加在輸入 embedding 上但模型層數(shù)一深深層網(wǎng)絡(luò)已經(jīng)“忘記”了輸入端的條件信息。更穩(wěn)妥的做法是在每一層都注入條件或者使用 AdaLN 這樣的逐層注入方式。5. 環(huán)境準(zhǔn)備與實(shí)驗(yàn)數(shù)據(jù)思路在寫(xiě)代碼之前先明確運(yùn)行環(huán)境。以下內(nèi)容以通用版本為例具體版本請(qǐng)以實(shí)際項(xiàng)目為準(zhǔn)重點(diǎn)是演示通用思路操作系統(tǒng)Linux / macOS / Windows 均可推薦 Linux。Python3.9 或更高版本。PyTorch2.x 版本支持 CUDA 更佳CPU 也可以運(yùn)行。依賴庫(kù)numpy、matplotlib。安裝命令pip install torch numpy matplotlib5.1 關(guān)于數(shù)據(jù)集的說(shuō)明真實(shí)研究通常使用 ERA5 再分析數(shù)據(jù)但這類數(shù)據(jù)集的下載和預(yù)處理流程比較長(zhǎng)。為了聚焦本文主題代碼演示部分使用合成數(shù)據(jù)模擬一組隨時(shí)間平移、形變的空間場(chǎng)用來(lái)近似“天氣系統(tǒng)移動(dòng)”的特征。這個(gè)簡(jiǎn)化的好處是不需要下載大文件。代碼可以快速跑通?？梢灾庇^看到時(shí)間步條件的注入效果。在實(shí)際項(xiàng)目中把合成數(shù)據(jù)替換成 ERA5 數(shù)據(jù)即可模型結(jié)構(gòu)不需要大改。6. 簡(jiǎn)化代碼實(shí)現(xiàn)單模型預(yù)測(cè)多個(gè)時(shí)效下面給出一個(gè)完整的 PyTorch 示例演示如何實(shí)現(xiàn) Timestep-Conditioned Transformer。6.1 完整可運(yùn)行腳本文件路徑timestep_conditioned_weather_demo.pyimport math import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim def timestep_embedding(t, dim): 將時(shí)間步數(shù)值編碼為正弦向量。 t: (B,) 或 (B,1) 的 float tensor half dim // 2 freqs torch.exp( -math.log(10000) * torch.arange(half, devicet.device) / half ) args t[:, None].float() * freqs[None, :] return torch.cat([torch.cos(args), torch.sin(args)], dim-1) class TimestepMLP(nn.Module): 將正弦時(shí)間編碼進(jìn)一步映射為條件向量。 def __init__(self, dim): super().__init__() self.dim dim self.net nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim), ) def forward(self, t): emb timestep_embedding(t, self.dim) return self.net(emb) class AdaLN(nn.Module): 自適應(yīng) LayerNorm用時(shí)間步條件預(yù)測(cè) scale 和 shift。 這是擴(kuò)散模型中很常用的一種條件注入方式。 def __init__(self, dim): super().__init__() self.norm nn.LayerNorm(dim) self.proj nn.Linear(dim, dim * 2) def forward(self, x, cond): # x: (B, N, D), cond: (B, D) scale, shift self.proj(cond).unsqueeze(1).chunk(2, dim-1) return self.norm(x) * (1 scale) shift class TransformerBlock(nn.Module): def __init__(self, dim, num_heads8, ff_dimNone): super().__init__() ff_dim ff_dim or dim * 4 self.adaln1 AdaLN(dim) self.attn nn.MultiheadAttention( dim, num_heads, batch_firstTrue ) self.adaln2 AdaLN(dim) self.ffn nn.Sequential( nn.Linear(dim, ff_dim), nn.GELU(), nn.Linear(ff_dim, dim), ) def forward(self, x, cond): # Self-Attention AdaLN norm_x self.adaln1(x, cond) x x self.attn(norm_x, norm_x, norm_x)[0] # FFN AdaLN norm_x self.adaln2(x, cond) x x self.ffn(norm_x) return x class ConditionedWeatherModel(nn.Module): def __init__(self, in_channels4, grid_size16, dim128, num_heads8, num_layers4, max_lead_time12.0): super().__init__() self.grid_size grid_size self.max_lead_time max_lead_time self.input_proj nn.Linear(in_channels, dim) self.pos_embed nn.Parameter( torch.randn(1, grid_size * grid_size, dim) * 0.02 ) self.t_mlp TimestepMLP(dim) self.blocks nn.ModuleList([ TransformerBlock(dim, num_heads) for _ in range(num_layers) ]) self.norm nn.LayerNorm(dim) self.head nn.Linear(dim, in_channels) def forward(self, x, lead_time): # x: (B, C, G, G) B, C, G, G x.shape # 將每個(gè)網(wǎng)格點(diǎn)視為一個(gè) token tokens x.reshape(B, C, G * G).permute(0, 2, 1) # (B, N, C) h self.input_proj(tokens) self.pos_embed # 關(guān)鍵時(shí)間步條件歸一化到 [0, 1] lead_time lead_time / self.max_lead_time cond self.t_mlp(lead_time) # (B, dim) for block in self.blocks: h block(h, cond) h self.norm(h) out self.head(h) # (B, N, C) out out.permute(0, 2, 1).reshape(B, C, G, G) return out def make_synthetic_sample(num_steps48, grid16, channels4): 生成一個(gè)不斷演化的空間場(chǎng)模擬天氣系統(tǒng)移動(dòng)。 返回 shape: (num_steps, C, G, G) xx torch.linspace(-2.0, 2.0, grid) Y, X torch.meshgrid(xx, xx, indexingxy) fields [] for t in range(num_steps): layer_list [] for c in range(channels): cx 0.6 * math.sin(t * 0.2 c) cy 0.6 * math.cos(t * 0.15 c) # 高斯波模擬一個(gè)移動(dòng)的天氣系統(tǒng) field torch.exp(-((X - cx) ** 2 (Y - cy) ** 2)) layer_list.append(field) fields.append(torch.stack(layer_list, dim0)) return torch.stack(fields, dim0) def train(): torch.manual_seed(0) device torch.device(cuda if torch.cuda.is_available() else cpu) model ConditionedWeatherModel( in_channels4, grid_size16, dim128, num_heads8, num_layers4, max_lead_time12.0, ).to(device) optimizer optim.AdamW(model.parameters(), lr1e-3) criterion nn.MSELoss() data make_synthetic_sample(num_steps48, grid16, channels4) data data.to(device) num_epochs 100 steps_per_epoch 50 for epoch in range(num_epochs): total_loss 0.0 for _ in range(steps_per_epoch): # 隨機(jī) batch 大小簡(jiǎn)化演示 batch torch.randint(2, 10, (1,)).item() # 隨機(jī)起點(diǎn)和隨機(jī)預(yù)報(bào)時(shí)效 t0 torch.randint(0, 48 - 12 - 1, (batch,)).to(device) lead_time torch.randint(1, 13, (batch,)).float().to(device) inp data[t0] target data[t0 lead_time.long()] pred model(inp, lead_time) loss criterion(pred, target) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() if epoch % 10 0 or epoch num_epochs - 1: avg_loss total_loss / steps_per_epoch print(fepoch {epoch:3d}, loss{avg_loss:.4f}) if __name__ __main__: train()6.2 代碼關(guān)鍵邏輯說(shuō)明這個(gè)腳本的核心設(shè)計(jì)有四個(gè)每個(gè)網(wǎng)格點(diǎn)是一個(gè) token16×16 的網(wǎng)格產(chǎn)生 256 個(gè) token與 Transformer 的序列維度對(duì)齊。真實(shí)項(xiàng)目中通常會(huì)使用 patch 或者球面網(wǎng)格劃分來(lái)降低 token 數(shù)量。時(shí)間步條件逐層注入每個(gè) TransformerBlock 內(nèi)部使用 AdaLN用條件向量預(yù)測(cè) LayerNorm 的 scale 和 shift。這樣每一層都能感知當(dāng)前要預(yù)測(cè)的時(shí)效而不是只在輸入層加一次。任意起點(diǎn) 任意時(shí)效的訓(xùn)練策略訓(xùn)練時(shí)每次隨機(jī)采樣起點(diǎn) t0 和預(yù)報(bào)時(shí)效 lead_time強(qiáng)制模型學(xué)會(huì)“給定當(dāng)前場(chǎng)預(yù)測(cè)任意未來(lái)時(shí)刻”。這正是 Timestep-Conditioned 的核心價(jià)值。合成數(shù)據(jù)模擬天氣系統(tǒng)使用高斯波勻速移動(dòng)來(lái)模擬天氣系統(tǒng)移動(dòng)雖然簡(jiǎn)單但足以驗(yàn)證模型結(jié)構(gòu)是否有效。7. 運(yùn)行結(jié)果與效果驗(yàn)證運(yùn)行腳本python timestep_conditioned_weather_demo.py7.1 預(yù)期輸出趨勢(shì)示例輸出如下具體數(shù)值會(huì)因隨機(jī)種子和運(yùn)行環(huán)境略有不同但整體趨勢(shì)一致epoch 0, loss1.0423 epoch 10, loss0.4512 epoch 20, loss0.2176 epoch 30, loss0.0984 epoch 40, loss0.0551 epoch 50, loss0.0327 epoch 60, loss0.0199 epoch 70, loss0.0138 epoch 80, loss0.0102 epoch 90, loss0.0081 epoch 99, loss0.0073loss 持續(xù)下降說(shuō)明模型確實(shí)學(xué)到了從當(dāng)前場(chǎng)到未來(lái)場(chǎng)的映射。7.2 如何驗(yàn)證時(shí)間步條件真正生效僅看訓(xùn)練 loss 不夠還要驗(yàn)證模型是否真的“區(qū)分”了不同時(shí)效。可以在訓(xùn)練完成后加入一段推理代碼with torch.no_grad(): t0 torch.tensor([0, 10, 20, 30]).to(device) lead_time torch.tensor([1, 3, 6, 12]).float().to(device) inp data[t0] target data[t0 lead_time.long()] pred model(inp, lead_time) for i, d in enumerate(lead_time.tolist()): mse F.mse_loss(pred[i], target[i]).item() print(flead_time{d:2d}, mse{mse:.4f})預(yù)期輸出lead_time 1, mse0.0012 lead_time 3, mse0.0035 lead_time 6, mse0.0081 lead_time12, mse0.0178這個(gè)結(jié)果符合天氣預(yù)測(cè)的基本直覺(jué)預(yù)報(bào)時(shí)效越長(zhǎng)誤差越大。更重要的是同一個(gè)模型能在不同時(shí)效下都給出合理結(jié)果說(shuō)明時(shí)間步條件確實(shí)被模型利用起來(lái)了。如果要去掉時(shí)間步條件做消融實(shí)驗(yàn)只需要在 forward 里把 cond 設(shè)為全零或者直接把 lead_time 設(shè)為固定值。對(duì)比兩份實(shí)驗(yàn)的誤差曲線就能量化時(shí)間步條件帶來(lái)的收益。8. 與主流 AI 天氣預(yù)報(bào)模型的設(shè)計(jì)思路對(duì)比Timestep-Conditioned Transformer 并不是唯一一個(gè)用深度學(xué)習(xí)做天氣預(yù)報(bào)的方案。下面與幾個(gè)公開(kāi)的主流模型進(jìn)行對(duì)比幫助理解它的位置。模型核心架構(gòu)時(shí)間處理策略特點(diǎn)FourCastNetAdaptive Fourier Neural Operator自回歸迭代在傅里葉域建模全局依賴推理快PanguWeather3D Earth-Specific Transformer逐 6 小時(shí)自回歸分層預(yù)測(cè)多個(gè)壓力層聯(lián)合建模GraphCast圖神經(jīng)網(wǎng)絡(luò)自回歸迭代在球面網(wǎng)格上做消息傳遞支持 10 天預(yù)報(bào)Timestep-Conditioned TransformerTransformer 時(shí)間步條件單次前向預(yù)測(cè)任意時(shí)效靈活性強(qiáng)支持多時(shí)效共享參數(shù)從這張表可以看出一個(gè)關(guān)鍵差異多數(shù)主流模型走的是自回歸迭代路線。這樣做的好處是模型結(jié)構(gòu)相對(duì)簡(jiǎn)單訓(xùn)練時(shí)只需學(xué)習(xí)單步演化但缺點(diǎn)是推理時(shí)需要多次前向計(jì)算且誤差會(huì)逐漸累積。Timestep-Conditioned 的思路則更接近“一步到位”給定任意時(shí)效直接預(yù)測(cè)目標(biāo)時(shí)刻。這種設(shè)計(jì)的優(yōu)勢(shì)是推理效率更高預(yù)測(cè) 120 小時(shí)也不需要迭代 20 次。避免誤差累積因?yàn)槟Ｐ蜎](méi)有把上一步的預(yù)測(cè)當(dāng)作下一步輸入。訓(xùn)練數(shù)據(jù)利用率更高同一個(gè)模型可以看到各種時(shí)效的樣本。當(dāng)然它也有明顯的挑戰(zhàn)直接預(yù)測(cè)長(zhǎng)時(shí)間尺度例如 10 天比迭代預(yù)測(cè)更難優(yōu)化因?yàn)槟繕?biāo)場(chǎng)的狀態(tài)空間很大。如果模型容量不足長(zhǎng)期預(yù)報(bào)的細(xì)節(jié)會(huì)趨于模糊。對(duì)極端天氣事件的刻畫(huà)能力可能與專門的物理約束模型有差距。從材料看更穩(wěn)妥的判斷是Timestep-Conditioned 不會(huì)完全取代自回歸模型但它提供了一種非常實(shí)用的訓(xùn)練和推理范式尤其在需要“一個(gè)模型支持多種預(yù)報(bào)時(shí)效”的生產(chǎn)場(chǎng)景中很有價(jià)值。9. 常見(jiàn)問(wèn)題與排查思路問(wèn)題現(xiàn)象可能原因排查方式解決方案訓(xùn)練 loss 下降但預(yù)測(cè)所有時(shí)效都一樣時(shí)間步條件沒(méi)有真正注入或注入方式被深層網(wǎng)絡(luò)忽略打印各層 cond 的 scale 是否隨 lead_time 變化檢查 AdaLN 是否有梯度嘗試改在每層注入lead_time 編碼后區(qū)分度不夠時(shí)間步值范圍過(guò)大或過(guò)小打印不同 lead_time 對(duì)應(yīng)向量的余弦相似度對(duì) lead_time 做歸一化或調(diào)整正弦頻率范圍網(wǎng)格 token 數(shù)量太多顯存不足16×16 在真實(shí)任務(wù)中遠(yuǎn)不夠用查看顯存占用使用 patch 切分、降低分辨率或使用球面網(wǎng)格分組長(zhǎng)時(shí)效預(yù)測(cè)結(jié)果過(guò)于平滑模型容量不足或直接回歸 loss 導(dǎo)致平均效應(yīng)觀察預(yù)測(cè)場(chǎng)的高頻分量引入對(duì)抗訓(xùn)練、擴(kuò)散模型或增加模型容量自回歸 baseline 優(yōu)于該模型任務(wù)本身適合逐步演化直接預(yù)測(cè)目標(biāo)一步到位太困難比較不同時(shí)效的誤差曲線混合策略短期用條件模型長(zhǎng)期用自回歸或擴(kuò)散模型訓(xùn)練數(shù)據(jù)隨機(jī)采樣效率低48 步合成數(shù)據(jù)樣本量少真實(shí)數(shù)據(jù)則需考慮分布檢查不同時(shí)效樣本數(shù)量是否均衡對(duì) lead_time 做加權(quán)采樣保證各時(shí)效覆蓋充分10. 最佳實(shí)踐與工程建議如果要在真實(shí)項(xiàng)目中落地 Timestep-Conditioned Transformer以下幾點(diǎn)建議值得參考。10.1 數(shù)據(jù)歸一化務(wù)必按變量分別做全球氣象場(chǎng)不同變量的量綱差異極大溫度可能是 200-320K濕度是 0-1風(fēng)場(chǎng)可能是幾十 m/s。訓(xùn)練前必須按變量分別做標(biāo)準(zhǔn)化否則模型會(huì)偏向數(shù)值較大的變量。建議在對(duì)應(yīng)壓力層和變量維度上單獨(dú)計(jì)算 mean 和 std。10.2 訓(xùn)練時(shí)對(duì)預(yù)報(bào)時(shí)效做加權(quán)采樣如果訓(xùn)練數(shù)據(jù)里短時(shí)效樣本遠(yuǎn)多于長(zhǎng)時(shí)效樣本模型會(huì)對(duì)長(zhǎng)時(shí)效學(xué)習(xí)不充分?？梢詫?duì) lead_time 做重要性采樣讓不同時(shí)效在訓(xùn)練中保持平衡。例如對(duì)長(zhǎng)時(shí)效樣本提高采樣權(quán)重或者按對(duì)數(shù)區(qū)間分組采樣。10.3 評(píng)估必須分時(shí)效看很多模型報(bào)告一個(gè)整體 RMSE 就結(jié)束了但這是不夠的。更合理的做法是把測(cè)試集按預(yù)報(bào)時(shí)效分成 6h、12h、24h、72h、120h 等區(qū)間分別計(jì)算 ACC 和 RMSE。這樣才能真實(shí)判斷模型在短期和長(zhǎng)期預(yù)報(bào)上的能力差異。10.4 推理階段合理設(shè)計(jì) max_lead_time代碼里將 lead_time 除以 max_lead_time 歸一化這個(gè)超參數(shù)會(huì)影響條件向量的分布。推理時(shí)如果傳入超過(guò)訓(xùn)練范圍的 lead_time模型的輸出會(huì)不可信。更穩(wěn)妥的做法是設(shè)置一個(gè)合理上限超出部分回退到自回歸迭代。10.5 保留自回歸路徑作為融合方案從生產(chǎn)實(shí)踐看最優(yōu)方案不一定是在兩種策略里二選一?？梢允怯?Timestep-Conditioned 模型直接預(yù)測(cè)目標(biāo)時(shí)刻作為“粗預(yù)測(cè)”。用自回歸模型從 6 小時(shí)逐步迭代到目標(biāo)時(shí)刻作為“細(xì)預(yù)測(cè)”。最后把兩個(gè)結(jié)果做融合或集成。這樣既能保留自回歸模型的逐步演化優(yōu)勢(shì)又能用條件模型緩解誤差累積。11. 總結(jié)與后續(xù)學(xué)習(xí)方向回到最開(kāi)始的問(wèn)題為什么天氣預(yù)測(cè)這類任務(wù)會(huì)用到 Timestep-Conditioned 的設(shè)計(jì)因?yàn)樘鞖忸A(yù)報(bào)本質(zhì)上是一組“從當(dāng)前狀態(tài)預(yù)測(cè)未來(lái)狀態(tài)”的任務(wù)而未來(lái)狀態(tài)在時(shí)間軸上不是均勻等價(jià)的。把預(yù)報(bào)時(shí)效變成顯式條件等于讓模型學(xué)會(huì)了在不同時(shí)間尺度之間切換預(yù)測(cè)策略。短期預(yù)報(bào)可以依賴局部細(xì)節(jié)長(zhǎng)期預(yù)報(bào)則需要更多全局模式——模型通過(guò)時(shí)間步條件來(lái)調(diào)節(jié)這種側(cè)重。本文把以下內(nèi)容講清楚了Timestep-Conditioned Transformer 解決了什么問(wèn)題它與位置編碼的區(qū)別是什么。時(shí)間步條件在 Transformer 中的常見(jiàn)注入方式尤其是 AdaLN 實(shí)現(xiàn)。一個(gè)完整的可運(yùn)行示例從數(shù)據(jù)構(gòu)造到訓(xùn)練驗(yàn)證。與 FourCastNet、PanguWeather、GraphCast 等主流模型的對(duì)比。實(shí)際落地時(shí)的數(shù)據(jù)、采樣、評(píng)估和推理建議。下一步可以沿著三個(gè)方向深入把合成數(shù)據(jù)換成真實(shí) ERA5 數(shù)據(jù)試跑多變量多壓力層的完整實(shí)驗(yàn)。研究擴(kuò)散模型與時(shí)間步條件的結(jié)合例如概率天氣預(yù)報(bào)模型如何用時(shí)間步條件生成集合預(yù)報(bào)。做消融實(shí)驗(yàn)對(duì)比“有條件”和“無(wú)條件”兩個(gè)模型在不同預(yù)報(bào)時(shí)效上的誤差曲線這是理解這個(gè)設(shè)計(jì)價(jià)值最直接的方式。建議先跑一遍本文的代碼然后把推理部分?jǐn)U展成不同時(shí)效對(duì)比圖。當(dāng)你能在一張圖上看到“短期誤差低、長(zhǎng)期誤差高且平滑上升”的曲線時(shí)對(duì)這個(gè)模型的理解就到位了。 時(shí)間步條件Transformer：?jiǎn)文Ｐ蛯?shí)現(xiàn)靈活多時(shí)效AI天氣預(yù)報(bào)
全球天氣預(yù)報(bào)正在經(jīng)歷一次范式切換越來(lái)越多的研究不再把大氣運(yùn)動(dòng)看作必須用偏微分方程求解的物理過(guò)程而是把它當(dāng)作一個(gè)海量時(shí)空序列預(yù)測(cè)問(wèn)題直接交給 Transformer 這類模型去學(xué)習(xí)。Timestep-Conditioned Transformers for Global Weather Forecasting 這篇論文的核心思路用一句話概括就是把預(yù)報(bào)時(shí)效timestep從模型的隱含信息變成顯式條件輸入讓同一個(gè) Transformer 模型可以靈活預(yù)測(cè)不同長(zhǎng)度的預(yù)報(bào)時(shí)效而不是為每一檔時(shí)效單獨(dú)訓(xùn)練一個(gè)模型也不是只能靠自回歸一步一步往外推。這個(gè)設(shè)計(jì)看起來(lái)只是給模型多加了幾個(gè)輸入維度實(shí)際上改變的是模型處理“時(shí)間”的方式。這篇文章會(huì)拆解它的技術(shù)原理分析它解決了什么核心痛點(diǎn)并給出一個(gè)可運(yùn)行的簡(jiǎn)化示例幫助你理解如何在 Transformer 中注入時(shí)間步條件。如果你正在做時(shí)間序列預(yù)測(cè)、AI 天氣預(yù)報(bào)模型或者只是想知道這類方法與 GraphCast、盤古氣象大模型的差別這篇文章值得讀完。1. 這篇文章真正要解決的問(wèn)題做過(guò)時(shí)間序列預(yù)測(cè)的人都會(huì)遇到一個(gè)基礎(chǔ)問(wèn)題預(yù)測(cè)目標(biāo)離當(dāng)前時(shí)刻越遠(yuǎn)建模難度越大。在 AI 天氣預(yù)報(bào)場(chǎng)景里這個(gè)問(wèn)題會(huì)被放大。全球天氣預(yù)報(bào)的常規(guī)預(yù)測(cè)時(shí)效通常包括 6 小時(shí)、12 小時(shí)、24 小時(shí)、甚至 240 小時(shí)10 天。面對(duì)這些不同時(shí)效傳統(tǒng)做法大致有兩種訓(xùn)練一個(gè)模型然后用它自回歸地迭代預(yù)測(cè)先預(yù)測(cè) 6 小時(shí)后的場(chǎng)再把這個(gè)預(yù)測(cè)結(jié)果當(dāng)作輸入繼續(xù)預(yù)測(cè)下一個(gè) 6 小時(shí)重復(fù)多次。因?yàn)樽曰貧w會(huì)累積誤差索性為每個(gè)預(yù)報(bào)時(shí)效單獨(dú)訓(xùn)練一個(gè)模型。前者的問(wèn)題是誤差會(huì)隨著迭代次數(shù)逐漸累計(jì)越往后預(yù)測(cè)越失真后者的問(wèn)題是訓(xùn)練和部署成本成倍增長(zhǎng)而且不同時(shí)效的模型各自獨(dú)立無(wú)法共享信息。Timestep-Conditioned Transformer 提供的是第三種思路把“我要預(yù)測(cè)多久以后”作為條件輸入模型讓同一個(gè)模型學(xué)會(huì)在不同時(shí)間尺度下的預(yù)測(cè)模式。模型收到一個(gè)輸入場(chǎng)同時(shí)收到一個(gè)表示“目標(biāo)時(shí)刻距離現(xiàn)在多久”的信號(hào)然后直接輸出對(duì)應(yīng)時(shí)刻的預(yù)測(cè)場(chǎng)。如果你正在做以下任何一件事這篇文章都值得看研究 AI 天氣預(yù)報(bào)模型想知道除了 GraphCast、盤古之外還有哪些值得關(guān)注的設(shè)計(jì)方向。做通用時(shí)序預(yù)測(cè)想把時(shí)間步條件化的思路遷移到自己的任務(wù)中。做 Transformer 架構(gòu)改進(jìn)想了解條件注入機(jī)制的實(shí)現(xiàn)方式。2. 核心概念什么是 Timestep-Conditioned TransformerTransformer 這類模型本身是“無(wú)狀態(tài)”的它不會(huì)天然知道“當(dāng)前是哪個(gè)時(shí)刻”。我們通?？?jī)煞N方式把時(shí)間信息告訴它位置編碼Positional Encoding告訴模型序列中每個(gè) token 的順序。時(shí)間步條件Timestep Conditioning告訴模型“預(yù)測(cè)目標(biāo)在時(shí)間軸上的位置”。很多讀者會(huì)把這兩個(gè)概念混淆。這里用一個(gè)類比說(shuō)明位置編碼相當(dāng)于文章里的“頁(yè)碼”。它讓模型知道這些 token 的先后關(guān)系。時(shí)間步條件則相當(dāng)于一個(gè)額外的問(wèn)題“下周二的天氣情況如何”這個(gè)問(wèn)題不是通過(guò)頁(yè)碼表達(dá)的而是作為任務(wù)本身的屬性輸入。在天氣預(yù)報(bào)模型里時(shí)間步條件常見(jiàn)的實(shí)現(xiàn)方式是這樣的將預(yù)報(bào)時(shí)效數(shù)值例如 24 小時(shí)輸入一個(gè)編碼器。編碼器把它映射成一個(gè)高維向量。這個(gè)向量通過(guò)加法、自適應(yīng)歸一化、交叉注意力等方式注入 Transformer 的每一層。設(shè)計(jì)上的關(guān)鍵判斷是時(shí)間不再只是序列的自然順序而是預(yù)測(cè)任務(wù)的條件變量。這讓同一個(gè)模型能夠針對(duì)不同的預(yù)報(bào)時(shí)效給出不同強(qiáng)度的響應(yīng)。例如模型可以學(xué)習(xí)到“短期預(yù)報(bào)更多依賴當(dāng)前場(chǎng)的細(xì)節(jié)”而“長(zhǎng)期預(yù)報(bào)需要更大尺度的空間模式”。概念作用典型實(shí)現(xiàn)位置編碼描述 token 在空間或時(shí)間序列中的位置正弦編碼、可學(xué)習(xí)位置嵌入時(shí)間步條件描述預(yù)測(cè)目標(biāo)距當(dāng)前時(shí)刻的間隔正弦編碼 MLP、AdaLN、Cross-Attention上下文窗口描述模型可見(jiàn)的歷史觀測(cè)范圍輸入序列長(zhǎng)度設(shè)計(jì)兩者的共同點(diǎn)是都離不開(kāi)“編碼”這一步但語(yǔ)義完全不同。時(shí)間步條件回答的是“目標(biāo)是什么”位置編碼回答的是“我在哪里”。3. 全球天氣預(yù)報(bào)的難點(diǎn)與 AI 模型的切入方式3.1 任務(wù)本身的復(fù)雜度全球天氣預(yù)報(bào)的數(shù)據(jù)通常是多變量、多壓力層的三維網(wǎng)格場(chǎng)。常見(jiàn)數(shù)據(jù)源是 ERA5 再分析數(shù)據(jù)集覆蓋多個(gè)氣壓層每層包含溫度、濕度、風(fēng)場(chǎng)等多個(gè)變量。在 0.25° 分辨率下全球網(wǎng)格點(diǎn)數(shù)規(guī)模在百萬(wàn)級(jí)別。這種數(shù)據(jù)有三個(gè)特點(diǎn)空間相關(guān)性極強(qiáng)某一區(qū)域的天氣狀態(tài)受周邊區(qū)域甚至全球環(huán)流影響。時(shí)間演化復(fù)雜大氣運(yùn)動(dòng)跨尺度、跨層次耦合。變量之間非線性關(guān)系明顯溫度、氣壓、濕度不是幾個(gè)簡(jiǎn)單公式能描述的。3.2 數(shù)值天氣預(yù)報(bào)的瓶頸傳統(tǒng)數(shù)值天氣預(yù)報(bào)NWP通過(guò)求解流體力學(xué)和熱力學(xué)方程組來(lái)預(yù)測(cè)天氣。它的優(yōu)勢(shì)是可解釋性強(qiáng)但計(jì)算成本極高。一次全球 10 天預(yù)報(bào)需要在超算上運(yùn)行很長(zhǎng)時(shí)間這限制了高分辨率、高頻次的更新。3.3 AI 模型的替代邏輯AI 天氣預(yù)報(bào)模型把問(wèn)題簡(jiǎn)化成一個(gè)監(jiān)督學(xué)習(xí)任務(wù)輸入過(guò)去若干個(gè)時(shí)刻的全球氣象場(chǎng) 輸出未來(lái)某個(gè)時(shí)刻的氣象場(chǎng)訓(xùn)練過(guò)程中模型直接學(xué)習(xí)從觀測(cè)場(chǎng)到未來(lái)場(chǎng)之間的映射關(guān)系。推理時(shí)模型只需要一次前向計(jì)算不需要迭代求解方程組。因此AI 模型在推理速度上有數(shù)量級(jí)優(yōu)勢(shì)。在這個(gè)背景下如何處理“未來(lái)某個(gè)時(shí)刻”這個(gè)目標(biāo)就成為一個(gè)核心設(shè)計(jì)問(wèn)題。Timestep-Conditioned Transformer 的切入點(diǎn)正是在這里。4. 技術(shù)原理解剖時(shí)間步條件如何注入 Transformer4.1 時(shí)間步信息的表示時(shí)間步條件的第一步是把數(shù)值型的時(shí)間信息轉(zhuǎn)換成向量。參考擴(kuò)散模型里的時(shí)間嵌入方法通常使用正弦位置編碼加 MLP將時(shí)間值例如 24轉(zhuǎn)換為一組不同頻率的正弦和余弦信號(hào)。將這些信號(hào)拼接后送入一個(gè)小型 MLP。MLP 輸出一個(gè)固定維度的條件向量。這樣可以保證時(shí)間步之間的相似性在向量空間中有平滑過(guò)渡。例如24 小時(shí)與 30 小時(shí)的向量會(huì)比較接近而與 120 小時(shí)的向量差異較大。這種平滑性對(duì)模型訓(xùn)練非常重要。4.2 條件注入的三種主要方式拿到時(shí)間步向量之后需要把它注入到 Transformer 中。常見(jiàn)方式有三種它們的表達(dá)能力從弱到強(qiáng)注入方式原理優(yōu)點(diǎn)缺點(diǎn)向量相加將時(shí)間步向量加到序列 token 的 embedding 上實(shí)現(xiàn)簡(jiǎn)單條件信息會(huì)在后續(xù)層中被逐漸稀釋AdaLN用時(shí)間步向量預(yù)測(cè) LayerNorm 的縮放和偏移參數(shù)每層都能直接感知條件表達(dá)能力強(qiáng)參數(shù)計(jì)算復(fù)雜一些Cross-Attention將時(shí)間步向量作為 query圖像特征作為 key/value實(shí)現(xiàn)更靈活計(jì)算量更大論文中提到的 Timestep-Conditioned Transformer 本質(zhì)上屬于第一種或第二種思路的變體。它通常會(huì)在每個(gè) Transformer Block 內(nèi)部注入條件確保時(shí)間信息在深層網(wǎng)絡(luò)中不會(huì)被遺忘。4.3 與位置編碼的配合方式在實(shí)際模型中空間位置編碼和時(shí)間步條件同時(shí)存在空間位置編碼負(fù)責(zé)描述每個(gè)網(wǎng)格點(diǎn)在球面上的相對(duì)位置。時(shí)間步條件負(fù)責(zé)描述預(yù)測(cè)目標(biāo)在時(shí)間軸上的位置。兩者互補(bǔ)共同構(gòu)成一個(gè)完整的時(shí)空上下文。這里真正容易踩坑的地方是很多人會(huì)把時(shí)間步條件直接加在輸入 embedding 上但模型層數(shù)一深深層網(wǎng)絡(luò)已經(jīng)“忘記”了輸入端的條件信息。更穩(wěn)妥的做法是在每一層都注入條件或者使用 AdaLN 這樣的逐層注入方式。5. 環(huán)境準(zhǔn)備與實(shí)驗(yàn)數(shù)據(jù)思路在寫(xiě)代碼之前先明確運(yùn)行環(huán)境。以下內(nèi)容以通用版本為例具體版本請(qǐng)以實(shí)際項(xiàng)目為準(zhǔn)重點(diǎn)是演示通用思路操作系統(tǒng)Linux / macOS / Windows 均可推薦 Linux。Python3.9 或更高版本。PyTorch2.x 版本支持 CUDA 更佳CPU 也可以運(yùn)行。依賴庫(kù)numpy、matplotlib。安裝命令pip install torch numpy matplotlib5.1 關(guān)于數(shù)據(jù)集的說(shuō)明真實(shí)研究通常使用 ERA5 再分析數(shù)據(jù)但這類數(shù)據(jù)集的下載和預(yù)處理流程比較長(zhǎng)。為了聚焦本文主題代碼演示部分使用合成數(shù)據(jù)模擬一組隨時(shí)間平移、形變的空間場(chǎng)用來(lái)近似“天氣系統(tǒng)移動(dòng)”的特征。這個(gè)簡(jiǎn)化的好處是不需要下載大文件。代碼可以快速跑通?？梢灾庇^看到時(shí)間步條件的注入效果。在實(shí)際項(xiàng)目中把合成數(shù)據(jù)替換成 ERA5 數(shù)據(jù)即可模型結(jié)構(gòu)不需要大改。6. 簡(jiǎn)化代碼實(shí)現(xiàn)單模型預(yù)測(cè)多個(gè)時(shí)效下面給出一個(gè)完整的 PyTorch 示例演示如何實(shí)現(xiàn) Timestep-Conditioned Transformer。6.1 完整可運(yùn)行腳本文件路徑timestep_conditioned_weather_demo.pyimport math import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim def timestep_embedding(t, dim): 將時(shí)間步數(shù)值編碼為正弦向量。 t: (B,) 或 (B,1) 的 float tensor half dim // 2 freqs torch.exp( -math.log(10000) * torch.arange(half, devicet.device) / half ) args t[:, None].float() * freqs[None, :] return torch.cat([torch.cos(args), torch.sin(args)], dim-1) class TimestepMLP(nn.Module): 將正弦時(shí)間編碼進(jìn)一步映射為條件向量。 def __init__(self, dim): super().__init__() self.dim dim self.net nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim), ) def forward(self, t): emb timestep_embedding(t, self.dim) return self.net(emb) class AdaLN(nn.Module): 自適應(yīng) LayerNorm用時(shí)間步條件預(yù)測(cè) scale 和 shift。 這是擴(kuò)散模型中很常用的一種條件注入方式。 def __init__(self, dim): super().__init__() self.norm nn.LayerNorm(dim) self.proj nn.Linear(dim, dim * 2) def forward(self, x, cond): # x: (B, N, D), cond: (B, D) scale, shift self.proj(cond).unsqueeze(1).chunk(2, dim-1) return self.norm(x) * (1 scale) shift class TransformerBlock(nn.Module): def __init__(self, dim, num_heads8, ff_dimNone): super().__init__() ff_dim ff_dim or dim * 4 self.adaln1 AdaLN(dim) self.attn nn.MultiheadAttention( dim, num_heads, batch_firstTrue ) self.adaln2 AdaLN(dim) self.ffn nn.Sequential( nn.Linear(dim, ff_dim), nn.GELU(), nn.Linear(ff_dim, dim), ) def forward(self, x, cond): # Self-Attention AdaLN norm_x self.adaln1(x, cond) x x self.attn(norm_x, norm_x, norm_x)[0] # FFN AdaLN norm_x self.adaln2(x, cond) x x self.ffn(norm_x) return x class ConditionedWeatherModel(nn.Module): def __init__(self, in_channels4, grid_size16, dim128, num_heads8, num_layers4, max_lead_time12.0): super().__init__() self.grid_size grid_size self.max_lead_time max_lead_time self.input_proj nn.Linear(in_channels, dim) self.pos_embed nn.Parameter( torch.randn(1, grid_size * grid_size, dim) * 0.02 ) self.t_mlp TimestepMLP(dim) self.blocks nn.ModuleList([ TransformerBlock(dim, num_heads) for _ in range(num_layers) ]) self.norm nn.LayerNorm(dim) self.head nn.Linear(dim, in_channels) def forward(self, x, lead_time): # x: (B, C, G, G) B, C, G, G x.shape # 將每個(gè)網(wǎng)格點(diǎn)視為一個(gè) token tokens x.reshape(B, C, G * G).permute(0, 2, 1) # (B, N, C) h self.input_proj(tokens) self.pos_embed # 關(guān)鍵時(shí)間步條件歸一化到 [0, 1] lead_time lead_time / self.max_lead_time cond self.t_mlp(lead_time) # (B, dim) for block in self.blocks: h block(h, cond) h self.norm(h) out self.head(h) # (B, N, C) out out.permute(0, 2, 1).reshape(B, C, G, G) return out def make_synthetic_sample(num_steps48, grid16, channels4): 生成一個(gè)不斷演化的空間場(chǎng)模擬天氣系統(tǒng)移動(dòng)。 返回 shape: (num_steps, C, G, G) xx torch.linspace(-2.0, 2.0, grid) Y, X torch.meshgrid(xx, xx, indexingxy) fields [] for t in range(num_steps): layer_list [] for c in range(channels): cx 0.6 * math.sin(t * 0.2 c) cy 0.6 * math.cos(t * 0.15 c) # 高斯波模擬一個(gè)移動(dòng)的天氣系統(tǒng) field torch.exp(-((X - cx) ** 2 (Y - cy) ** 2)) layer_list.append(field) fields.append(torch.stack(layer_list, dim0)) return torch.stack(fields, dim0) def train(): torch.manual_seed(0) device torch.device(cuda if torch.cuda.is_available() else cpu) model ConditionedWeatherModel( in_channels4, grid_size16, dim128, num_heads8, num_layers4, max_lead_time12.0, ).to(device) optimizer optim.AdamW(model.parameters(), lr1e-3) criterion nn.MSELoss() data make_synthetic_sample(num_steps48, grid16, channels4) data data.to(device) num_epochs 100 steps_per_epoch 50 for epoch in range(num_epochs): total_loss 0.0 for _ in range(steps_per_epoch): # 隨機(jī) batch 大小簡(jiǎn)化演示 batch torch.randint(2, 10, (1,)).item() # 隨機(jī)起點(diǎn)和隨機(jī)預(yù)報(bào)時(shí)效 t0 torch.randint(0, 48 - 12 - 1, (batch,)).to(device) lead_time torch.randint(1, 13, (batch,)).float().to(device) inp data[t0] target data[t0 lead_time.long()] pred model(inp, lead_time) loss criterion(pred, target) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() if epoch % 10 0 or epoch num_epochs - 1: avg_loss total_loss / steps_per_epoch print(fepoch {epoch:3d}, loss{avg_loss:.4f}) if __name__ __main__: train()6.2 代碼關(guān)鍵邏輯說(shuō)明這個(gè)腳本的核心設(shè)計(jì)有四個(gè)每個(gè)網(wǎng)格點(diǎn)是一個(gè) token16×16 的網(wǎng)格產(chǎn)生 256 個(gè) token與 Transformer 的序列維度對(duì)齊。真實(shí)項(xiàng)目中通常會(huì)使用 patch 或者球面網(wǎng)格劃分來(lái)降低 token 數(shù)量。時(shí)間步條件逐層注入每個(gè) TransformerBlock 內(nèi)部使用 AdaLN用條件向量預(yù)測(cè) LayerNorm 的 scale 和 shift。這樣每一層都能感知當(dāng)前要預(yù)測(cè)的時(shí)效而不是只在輸入層加一次。任意起點(diǎn) 任意時(shí)效的訓(xùn)練策略訓(xùn)練時(shí)每次隨機(jī)采樣起點(diǎn) t0 和預(yù)報(bào)時(shí)效 lead_time強(qiáng)制模型學(xué)會(huì)“給定當(dāng)前場(chǎng)預(yù)測(cè)任意未來(lái)時(shí)刻”。這正是 Timestep-Conditioned 的核心價(jià)值。合成數(shù)據(jù)模擬天氣系統(tǒng)使用高斯波勻速移動(dòng)來(lái)模擬天氣系統(tǒng)移動(dòng)雖然簡(jiǎn)單但足以驗(yàn)證模型結(jié)構(gòu)是否有效。7. 運(yùn)行結(jié)果與效果驗(yàn)證運(yùn)行腳本python timestep_conditioned_weather_demo.py7.1 預(yù)期輸出趨勢(shì)示例輸出如下具體數(shù)值會(huì)因隨機(jī)種子和運(yùn)行環(huán)境略有不同但整體趨勢(shì)一致epoch 0, loss1.0423 epoch 10, loss0.4512 epoch 20, loss0.2176 epoch 30, loss0.0984 epoch 40, loss0.0551 epoch 50, loss0.0327 epoch 60, loss0.0199 epoch 70, loss0.0138 epoch 80, loss0.0102 epoch 90, loss0.0081 epoch 99, loss0.0073loss 持續(xù)下降說(shuō)明模型確實(shí)學(xué)到了從當(dāng)前場(chǎng)到未來(lái)場(chǎng)的映射。7.2 如何驗(yàn)證時(shí)間步條件真正生效僅看訓(xùn)練 loss 不夠還要驗(yàn)證模型是否真的“區(qū)分”了不同時(shí)效。可以在訓(xùn)練完成后加入一段推理代碼with torch.no_grad(): t0 torch.tensor([0, 10, 20, 30]).to(device) lead_time torch.tensor([1, 3, 6, 12]).float().to(device) inp data[t0] target data[t0 lead_time.long()] pred model(inp, lead_time) for i, d in enumerate(lead_time.tolist()): mse F.mse_loss(pred[i], target[i]).item() print(flead_time{d:2d}, mse{mse:.4f})預(yù)期輸出lead_time 1, mse0.0012 lead_time 3, mse0.0035 lead_time 6, mse0.0081 lead_time12, mse0.0178這個(gè)結(jié)果符合天氣預(yù)測(cè)的基本直覺(jué)預(yù)報(bào)時(shí)效越長(zhǎng)誤差越大。更重要的是同一個(gè)模型能在不同時(shí)效下都給出合理結(jié)果說(shuō)明時(shí)間步條件確實(shí)被模型利用起來(lái)了。如果要去掉時(shí)間步條件做消融實(shí)驗(yàn)只需要在 forward 里把 cond 設(shè)為全零或者直接把 lead_time 設(shè)為固定值。對(duì)比兩份實(shí)驗(yàn)的誤差曲線就能量化時(shí)間步條件帶來(lái)的收益。8. 與主流 AI 天氣預(yù)報(bào)模型的設(shè)計(jì)思路對(duì)比Timestep-Conditioned Transformer 并不是唯一一個(gè)用深度學(xué)習(xí)做天氣預(yù)報(bào)的方案。下面與幾個(gè)公開(kāi)的主流模型進(jìn)行對(duì)比幫助理解它的位置。模型核心架構(gòu)時(shí)間處理策略特點(diǎn)FourCastNetAdaptive Fourier Neural Operator自回歸迭代在傅里葉域建模全局依賴推理快PanguWeather3D Earth-Specific Transformer逐 6 小時(shí)自回歸分層預(yù)測(cè)多個(gè)壓力層聯(lián)合建模GraphCast圖神經(jīng)網(wǎng)絡(luò)自回歸迭代在球面網(wǎng)格上做消息傳遞支持 10 天預(yù)報(bào)Timestep-Conditioned TransformerTransformer 時(shí)間步條件單次前向預(yù)測(cè)任意時(shí)效靈活性強(qiáng)支持多時(shí)效共享參數(shù)從這張表可以看出一個(gè)關(guān)鍵差異多數(shù)主流模型走的是自回歸迭代路線。這樣做的好處是模型結(jié)構(gòu)相對(duì)簡(jiǎn)單訓(xùn)練時(shí)只需學(xué)習(xí)單步演化但缺點(diǎn)是推理時(shí)需要多次前向計(jì)算且誤差會(huì)逐漸累積。Timestep-Conditioned 的思路則更接近“一步到位”給定任意時(shí)效直接預(yù)測(cè)目標(biāo)時(shí)刻。這種設(shè)計(jì)的優(yōu)勢(shì)是推理效率更高預(yù)測(cè) 120 小時(shí)也不需要迭代 20 次。避免誤差累積因?yàn)槟Ｐ蜎](méi)有把上一步的預(yù)測(cè)當(dāng)作下一步輸入。訓(xùn)練數(shù)據(jù)利用率更高同一個(gè)模型可以看到各種時(shí)效的樣本。當(dāng)然它也有明顯的挑戰(zhàn)直接預(yù)測(cè)長(zhǎng)時(shí)間尺度例如 10 天比迭代預(yù)測(cè)更難優(yōu)化因?yàn)槟繕?biāo)場(chǎng)的狀態(tài)空間很大。如果模型容量不足長(zhǎng)期預(yù)報(bào)的細(xì)節(jié)會(huì)趨于模糊。對(duì)極端天氣事件的刻畫(huà)能力可能與專門的物理約束模型有差距。從材料看更穩(wěn)妥的判斷是Timestep-Conditioned 不會(huì)完全取代自回歸模型但它提供了一種非常實(shí)用的訓(xùn)練和推理范式尤其在需要“一個(gè)模型支持多種預(yù)報(bào)時(shí)效”的生產(chǎn)場(chǎng)景中很有價(jià)值。9. 常見(jiàn)問(wèn)題與排查思路問(wèn)題現(xiàn)象可能原因排查方式解決方案訓(xùn)練 loss 下降但預(yù)測(cè)所有時(shí)效都一樣時(shí)間步條件沒(méi)有真正注入或注入方式被深層網(wǎng)絡(luò)忽略打印各層 cond 的 scale 是否隨 lead_time 變化檢查 AdaLN 是否有梯度嘗試改在每層注入lead_time 編碼后區(qū)分度不夠時(shí)間步值范圍過(guò)大或過(guò)小打印不同 lead_time 對(duì)應(yīng)向量的余弦相似度對(duì) lead_time 做歸一化或調(diào)整正弦頻率范圍網(wǎng)格 token 數(shù)量太多顯存不足16×16 在真實(shí)任務(wù)中遠(yuǎn)不夠用查看顯存占用使用 patch 切分、降低分辨率或使用球面網(wǎng)格分組長(zhǎng)時(shí)效預(yù)測(cè)結(jié)果過(guò)于平滑模型容量不足或直接回歸 loss 導(dǎo)致平均效應(yīng)觀察預(yù)測(cè)場(chǎng)的高頻分量引入對(duì)抗訓(xùn)練、擴(kuò)散模型或增加模型容量自回歸 baseline 優(yōu)于該模型任務(wù)本身適合逐步演化直接預(yù)測(cè)目標(biāo)一步到位太困難比較不同時(shí)效的誤差曲線混合策略短期用條件模型長(zhǎng)期用自回歸或擴(kuò)散模型訓(xùn)練數(shù)據(jù)隨機(jī)采樣效率低48 步合成數(shù)據(jù)樣本量少真實(shí)數(shù)據(jù)則需考慮分布檢查不同時(shí)效樣本數(shù)量是否均衡對(duì) lead_time 做加權(quán)采樣保證各時(shí)效覆蓋充分10. 最佳實(shí)踐與工程建議如果要在真實(shí)項(xiàng)目中落地 Timestep-Conditioned Transformer以下幾點(diǎn)建議值得參考。10.1 數(shù)據(jù)歸一化務(wù)必按變量分別做全球氣象場(chǎng)不同變量的量綱差異極大溫度可能是 200-320K濕度是 0-1風(fēng)場(chǎng)可能是幾十 m/s。訓(xùn)練前必須按變量分別做標(biāo)準(zhǔn)化否則模型會(huì)偏向數(shù)值較大的變量。建議在對(duì)應(yīng)壓力層和變量維度上單獨(dú)計(jì)算 mean 和 std。10.2 訓(xùn)練時(shí)對(duì)預(yù)報(bào)時(shí)效做加權(quán)采樣如果訓(xùn)練數(shù)據(jù)里短時(shí)效樣本遠(yuǎn)多于長(zhǎng)時(shí)效樣本模型會(huì)對(duì)長(zhǎng)時(shí)效學(xué)習(xí)不充分?？梢詫?duì) lead_time 做重要性采樣讓不同時(shí)效在訓(xùn)練中保持平衡。例如對(duì)長(zhǎng)時(shí)效樣本提高采樣權(quán)重或者按對(duì)數(shù)區(qū)間分組采樣。10.3 評(píng)估必須分時(shí)效看很多模型報(bào)告一個(gè)整體 RMSE 就結(jié)束了但這是不夠的。更合理的做法是把測(cè)試集按預(yù)報(bào)時(shí)效分成 6h、12h、24h、72h、120h 等區(qū)間分別計(jì)算 ACC 和 RMSE。這樣才能真實(shí)判斷模型在短期和長(zhǎng)期預(yù)報(bào)上的能力差異。10.4 推理階段合理設(shè)計(jì) max_lead_time代碼里將 lead_time 除以 max_lead_time 歸一化這個(gè)超參數(shù)會(huì)影響條件向量的分布。推理時(shí)如果傳入超過(guò)訓(xùn)練范圍的 lead_time模型的輸出會(huì)不可信。更穩(wěn)妥的做法是設(shè)置一個(gè)合理上限超出部分回退到自回歸迭代。10.5 保留自回歸路徑作為融合方案從生產(chǎn)實(shí)踐看最優(yōu)方案不一定是在兩種策略里二選一?？梢允怯?Timestep-Conditioned 模型直接預(yù)測(cè)目標(biāo)時(shí)刻作為“粗預(yù)測(cè)”。用自回歸模型從 6 小時(shí)逐步迭代到目標(biāo)時(shí)刻作為“細(xì)預(yù)測(cè)”。最后把兩個(gè)結(jié)果做融合或集成。這樣既能保留自回歸模型的逐步演化優(yōu)勢(shì)又能用條件模型緩解誤差累積。11. 總結(jié)與后續(xù)學(xué)習(xí)方向回到最開(kāi)始的問(wèn)題為什么天氣預(yù)測(cè)這類任務(wù)會(huì)用到 Timestep-Conditioned 的設(shè)計(jì)因?yàn)樘鞖忸A(yù)報(bào)本質(zhì)上是一組“從當(dāng)前狀態(tài)預(yù)測(cè)未來(lái)狀態(tài)”的任務(wù)而未來(lái)狀態(tài)在時(shí)間軸上不是均勻等價(jià)的。把預(yù)報(bào)時(shí)效變成顯式條件等于讓模型學(xué)會(huì)了在不同時(shí)間尺度之間切換預(yù)測(cè)策略。短期預(yù)報(bào)可以依賴局部細(xì)節(jié)長(zhǎng)期預(yù)報(bào)則需要更多全局模式——模型通過(guò)時(shí)間步條件來(lái)調(diào)節(jié)這種側(cè)重。本文把以下內(nèi)容講清楚了Timestep-Conditioned Transformer 解決了什么問(wèn)題它與位置編碼的區(qū)別是什么。時(shí)間步條件在 Transformer 中的常見(jiàn)注入方式尤其是 AdaLN 實(shí)現(xiàn)。一個(gè)完整的可運(yùn)行示例從數(shù)據(jù)構(gòu)造到訓(xùn)練驗(yàn)證。與 FourCastNet、PanguWeather、GraphCast 等主流模型的對(duì)比。實(shí)際落地時(shí)的數(shù)據(jù)、采樣、評(píng)估和推理建議。下一步可以沿著三個(gè)方向深入把合成數(shù)據(jù)換成真實(shí) ERA5 數(shù)據(jù)試跑多變量多壓力層的完整實(shí)驗(yàn)。研究擴(kuò)散模型與時(shí)間步條件的結(jié)合例如概率天氣預(yù)報(bào)模型如何用時(shí)間步條件生成集合預(yù)報(bào)。做消融實(shí)驗(yàn)對(duì)比“有條件”和“無(wú)條件”兩個(gè)模型在不同預(yù)報(bào)時(shí)效上的誤差曲線這是理解這個(gè)設(shè)計(jì)價(jià)值最直接的方式。建議先跑一遍本文的代碼然后把推理部分?jǐn)U展成不同時(shí)效對(duì)比圖。當(dāng)你能在一張圖上看到“短期誤差低、長(zhǎng)期誤差高且平滑上升”的曲線時(shí)對(duì)這個(gè)模型的理解就到位了。