Task-Adaptive Parameter-Efficient Fine-Tuning for Weather Foundation Models¶
会议: ICLR 2026
OpenReview: https://openreview.net/forum?id=eFExhM3tKr
代码: https://github.com/ShileiCao/WeatherPEFT
领域: 地球科学 / 气象基础模型 / 参数高效微调
关键词: 天气基础模型, PEFT, 任务自适应 prompt, Fisher 信息, 参数选择
一句话总结¶
针对天气基础模型(WFM)微调,提出 WeatherPEFT:前向用 Task-Adaptive Dynamic Prompting(TADP)从编码器嵌入权重里抽出"变量×分辨率×时空"的任务特征生成软 prompt,反向用 Stochastic Fisher-Guided Adaptive Selection(SFAS)只更新 Fisher 信息最高的少量参数,在三个下游任务上用 ~0.3%–4% 的可训练参数追平甚至超过全量微调(Full-Tuning)。
研究背景与动机¶
领域现状:天气基础模型(如 Aurora 1.3B、Prithvi-WxC 2.3B)在大规模气象数据上预训练后,能迁移到降尺度、集合预报后处理、降水预测等多种下游任务,正在取代传统数值天气预报(NWP)。但模型越做越大,每个下游任务都做全量微调(Full-Tuning)在算力和存储上都不可持续——每个任务都要存一整套 10 亿级参数。
现有痛点:自然语言/视觉里成熟的 PEFT 方法(LoRA、DoRA、AdaptFormer、SSF、VPT 等)直接搬到 WFM 上效果明显掉点。论文在降尺度任务上量化了这个差距:DoRA 用 3.75M 参数,T2m 的 RMSE 比 Full-Tuning 高约 36%(1.228 vs 0.906);在区域降水任务上 LoRA 的 12h SEEPS 比 Full-Tuning 高 62.8%。
核心矛盾:气象数据和 RGB 图像/词向量根本不同——它在变量(温度、湿度等不同物理量,且变量间相关性随任务变)、分辨率(5.625° vs 0.25° 不只是空间尺度,而是从静力大尺度动力学切到非静力对流尺度的物理 regime 切换)、时空覆盖(全球 vs 区域)三个维度上高度异质。而主流 PEFT 对所有输入、所有任务用同一套可训练参数,既无法感知当前任务的物理特性,也没意识到不同参数在不同任务里重要性不同(降水预测关键的参数未必是降尺度关键的)。已有的任务选择型 PEFT(Child-Tuning、SAM、SCT)虽然会挑参数,但都是训练前静态选定,无法动态适配气象的变量耦合和 regime 切换。
本文目标:让 PEFT 既能在前向感知任务(context-aware),又能在反向按任务挑参数(task-adaptive),且这两件事都要针对气象异质性专门设计。
核心 idea:把编码器嵌入层当成"任务信息的仓库"——它天然编码了输入变量、分辨率、天气现象,于是从嵌入权重里抽特征生成动态软 prompt(TADP);再用 Fisher 信息加退火随机扰动来稳健地挑出当前任务最敏感的少量参数更新(SFAS)。两者一个管前向一个管反向,协同工作。
方法详解¶
整体框架¶
WeatherPEFT 在微调时冻结整个预训练 backbone(3D Swin Transformer U-Net),只训练两个轻量模块,它们作用在微调流程的两个不同阶段:
- 前向阶段——TADP:不直接拿数据,而是拿编码器嵌入层的权重 \(E\in\mathbb{R}^{D\times V\times P_h\times P_w}\) 作为输入(\(V\) 变量数、\(P_h\times P_w\) patch 空间核、\(D\) 隐藏维),先用三个层级化适配器抽"内部模式",再用自注意力抽"外部模式",生成软 prompt token \(E_P\in\mathbb{R}^{P\times D}\),把它拼接到输入 token 前面送进 backbone 每一层。这样模型在每一层计算时都带着"当前是什么任务"的上下文。
- 反向阶段——SFAS:backbone 参数虽然冻结,但仍需选一部分微调。SFAS 用 Fisher 信息衡量每个参数对损失的敏感度,加退火随机扰动后取 Top-k 生成 Fish Mask,只让被选中的参数接收梯度更新,其余保持预训练值不变。
两个模块互补:TADP 让前向"看得懂任务",SFAS 让反向"只动该动的参数",消融显示二者缺一都掉点。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["编码器嵌入权重<br/>E (D×V×Ph×Pw)"] --> B["TADP 内部模式<br/>HW→V→D 三级适配器"]
B --> C["TADP 外部模式<br/>自注意力耦合变量×分辨率"]
C --> D["软 prompt token EP"]
E["输入 X<br/>编码为 token T"] --> F["T 与 EP 拼接"]
D --> F
F --> G["冻结 backbone 逐层<br/>SFAS Fisher-Top-k 选参更新"]
G --> H["解码器 → 预测输出"]
关键设计¶
1. TADP 内部模式:三级适配器按"空间→变量→特征"层级抽取嵌入权重
针对气象数据三维异质(变量/分辨率/时空),TADP 不在数据上加 prompt,而是去挖嵌入权重里已经编码好的任务特征。它把嵌入权重的空间维展平为 \(\hat{E}\in\mathbb{R}^{D\times V\times P_hP_w}\),然后串联三个适配器,每个都是 LayerNorm + 降维 + GELU + 升维的轻量结构,按由低到高的层级递进:HW-Adapter 先处理空间与分辨率信息(\(P_h\times P_w\)),建立"特征在空间位置上怎么交互"的基础语境;V-Adapter 在此基础上建模不同物理变量(温度、湿度等)之间的相互依赖;D-Adapter 处理由天气特性抽象出的高层属性(\(D\) 维),整合前两阶段输出得到解释大气响应机制的普适模式。形式上:
其中 \(\pi\) 是把张量最后一维移到最前的换轴操作,恰好让下一个适配器作用在新的维度上,从而实现"空间→变量→特征"逐级聚焦。这种分层正对应气象数据"局部空间交互 → 变量耦合 → 抽象大气机制"的物理层次,比统一加一个 prompt 更能刻画当前任务的物理特性。
2. TADP 外部模式:自注意力耦合变量与分辨率,生成注入式软 prompt
光有内部分层特征还不够——变量和空间分辨率之间存在跨维度的耦合(例如不同分辨率下变量相关性会变),需要显式建模。TADP 把 \(E_D\) 的前两维合并为 \(\hat{E}_D\in\mathbb{R}^{VP_hP_w\times D}\),做一次自注意力把"物理量 \(V\)"和"空间分辨率特征 \(P_hP_w\)"耦合起来,再经线性投影压成 \(P\) 个 prompt token:
输入 \(X\) 被编码器编成 \(M\) 个 token \(T\in\mathbb{R}^{M\times D}\) 后,生成的软 prompt \(E_P\) 被拼接到 \(T\) 前面,并送进 backbone 的每一个 block。这样任务上下文不是只在输入端注入一次,而是贯穿整个前向计算,使冻结的 backbone 在每一层都"知道自己在做哪个气象任务",实现 context-aware 的特征重标定。
3. SFAS:Fisher 信息 + 退火随机扰动,动态挑选并稳健更新任务关键参数
不同天气任务对参数的依赖不同(有些参数编码降水的混沌模式,有些专管降尺度的空间关系),所以反向时不能均匀更新所有参数。SFAS 用 Fisher 信息矩阵 \(F_\theta\) 量化参数对学习目标的敏感度(Fisher 越大参数越关键),并用对角近似把它简化为向量。在有监督设定下用真实标签近似:
但天气任务异质性强,微调早期噪声大、Fisher 容易失真——早期高 Fisher 分的参数可能只是抓到了瞬时噪声而非真正任务相关特征。为此 SFAS 加一个退火随机分量稳定选择:
其中 \(\gamma\) 是初始系数,\(M_{sc}\sim\mathrm{Uniform}(0,1)\) 是随机向量,\(n_s/t_s\) 是当前步与总步数(每个 batch 算一步)。训练早期随机项大,给低 Fisher 参数一些被选中的机会、避免被初期噪声误导;随训练推进线性衰减到 0,最终回归纯 Fisher 选择。每步按 \(\bar{F}_\theta\) 取 Top-k 置 Fish Mask 为 1(其余为 0),只更新这 k 个参数,\(k\) 控制稀疏度。如此既保住不变的预训练物理先验,又把有限的更新预算花在当前任务真正关键的权重上。
损失函数 / 训练策略¶
沿用各下游任务自身的回归/概率评分目标(降尺度用纬度加权 RMSE、后处理用 CRPS、降水用 SEEPS/ACC/RMSE),WeatherPEFT 不改任务损失,只改"前向带什么 prompt、反向更新哪些参数"。backbone 全程冻结,只有 TADP 模块与被 Fish Mask 选中的 backbone 参数参与训练。
实验关键数据¶
三个任务均以 Aurora(1.3B,3D Swin-UNet)为基础模型,并在附录用更大的 Prithvi-WxC 验证;只统计 backbone 内的可训练参数。
主实验¶
降尺度(ERA5 5.625°→1.40625°,同时降尺度 68 个变量),纬度加权 RMSE(越低越好):
| 方法 | 可训练参数(M) | T2m | U10 | V10 | T850 | Z500 |
|---|---|---|---|---|---|---|
| LoRA | 3.63 | 1.190 | 1.130 | 1.118 | 0.998 | 50.421 |
| DoRA | 3.75 | 1.228 | 1.140 | 1.120 | 1.024 | 50.061 |
| SSF | 3.92 | 1.180 | 1.106 | 1.094 | 0.987 | 48.342 |
| TADP Only | 2.22 | 1.183 | 1.118 | 1.105 | 0.996 | 49.809 |
| SFAS Only | 1.26 | 1.161 | 1.090 | 1.081 | 0.973 | 47.000 |
| WeatherPEFT | 3.48 | 1.119 | 1.057 | 1.051 | 0.950 | 44.922 |
| Full-Tuning | 1239.94 | 0.906 | 0.882 | 0.884 | 0.836 | 35.821 |
WeatherPEFT 用 3.48M 参数(约 backbone 的 0.3%)在所有 PEFT 里 RMSE 最低;把预算放大到 ~4%(52.47M)后 T2m RMSE 0.916 已逼近 Full-Tuning 的 0.906,部分变量(V10 0.875 vs 0.884)甚至更好。
区域降水预测(ERA5-CH 中国区 0.25°,预测未来 6h 累计降水),SEEPS↓ / ACC↑ / RMSE↓,这里 WeatherPEFT 超过了 Full-Tuning:
| 方法 | 参数(M) | 12h SEEPS↓ | 12h ACC↑ | 24h SEEPS↓ | 36h ACC↑ |
|---|---|---|---|---|---|
| LoRA | 3.63 | 0.495 | 0.592 | 0.634 | 0.294 |
| Child-Tuning\(_D\) | 3.39 | 0.407 | 0.694 | 0.565 | 0.364 |
| SAM | 3.39 | 0.421 | 0.673 | 0.598 | 0.299 |
| WeatherPEFT | 3.38 | 0.368 | 0.742 | 0.515 | 0.443 |
| Full-Tuning | 1246.77 | 0.304 | 0.797 | 0.452 | 0.481 |
| WeatherPEFT(~4%) | 52.37 | 0.302 | 0.805 | 0.437 | 0.518 |
放大预算后 WeatherPEFT(52.37M)在所有指标上全面超过 1.2B 的 Full-Tuning(12h SEEPS 0.302 vs 0.304、ACC 0.805 vs 0.797)。在集合预报后处理任务上,WeatherPEFT 用 3.18M 参数也在 Z500 上反超 Full-Tuning(CRPS 72.701 vs 73.760)。
消融实验¶
| 配置 | 降尺度 T2m RMSE↓ | 降水 12h SEEPS↓ | 说明 |
|---|---|---|---|
| TADP Only | 1.183 | 0.549 | 只前向加 prompt |
| SFAS Only | 1.161 | 0.459 | 只反向选参数 |
| WeatherPEFT(full) | 1.119 | 0.368 | 两者协同 |
关键发现¶
- 两个模块协同 > 单用:TADP Only 和 SFAS Only 都比完整模型差,说明前向的任务感知和反向的参数选择是互补的,缺一个就掉点。
- 不同任务两模块的重要性不同:在稀疏、局部性强的降水任务上 SFAS 比 TADP 更关键(SFAS Only 0.459 远好于 TADP Only 0.549);这印证了"挑对参数"对捕捉降水稀疏信号的价值。
- 任务自适应选参确实有用,但还不够:SCT、SAM、Child-Tuning 这些任务选择型 baseline 整体优于 LoRA,验证了选任务相关参数的直觉;但它们仍被 WeatherPEFT 大幅超过,说明只有自适应选参、缺了 TADP 提供的气象上下文感知是不够的。
- PEFT 在气象域确实有大 gap:常规 PEFT(LoRA/DoRA)即便放大参数预算也追不上 Full-Tuning,凸显科学领域微调的特殊性,而 WeatherPEFT 几乎补齐了这个差距。
亮点与洞察¶
- 拿"嵌入权重"而非数据生成 prompt:TADP 把编码器嵌入层视作任务信息仓库,从权重而非输入特征里抽 prompt,巧妙地把"变量/分辨率/时空"这类任务级元信息显式注入——这个思路可迁移到任何输入异质、需要任务条件化的基础模型微调。
- 退火随机扰动稳住 Fisher 早期噪声:直接用 Fisher 选参在早期会被瞬时噪声误导,加一个线性衰减的随机项让选择在早期"敢探索"、后期"收敛到纯 Fisher",是个简单但针对性强的稳健化技巧,可复用到其他基于梯度敏感度的参数选择场景。
- 前向+反向双管齐下:大多数 PEFT 只在一侧动手(要么加模块、要么选参数),WeatherPEFT 同时在前向(感知任务)和反向(挑参数)两个阶段设计,二者正交互补,这种"分阶段设计"框架本身值得借鉴。
- PEFT 反超 Full-Tuning:在区域降水上放大预算后全面超过 1.2B 全量微调,说明对异质科学任务,"少而精地更新对的参数 + 任务条件化"可能比"无差别更新全部参数"更不易过拟合。
局限与展望¶
- 方法在 Aurora / Prithvi-WxC 两个 WFM 上验证,但都是 Transformer 类 backbone;对其他架构(如纯图神经网络或谱方法的气象模型)是否同样有效未知。
- TADP 依赖编码器嵌入层确实编码了足够的任务信息这一假设;若某些 WFM 的嵌入层信息不充分,prompt 质量可能下降。
- Fisher 信息需要额外前向/反向来估计,\(k\) 和退火系数 \(\gamma\) 是新引入的超参,论文虽在附录做了敏感性分析,但跨任务自动选取仍需人工调。
- 只覆盖降尺度、后处理、区域降水三类任务,对中长期全球预报、数据同化等更复杂时序任务的适配性还有待验证。
相关工作与启发¶
- vs 通用 PEFT(LoRA / DoRA / AdaptFormer / SSF / VPT):它们对所有任务/输入用同一套参数、面向 RGB 或词向量这类同质输入;本文针对气象的变量-分辨率-时空异质,前向做任务自适应 prompt、反向做任务自适应选参,因此在 WFM 上显著缩小与 Full-Tuning 的差距。
- vs 任务选择型 PEFT(Child-Tuning / SAM / SCT):它们在训练前静态挑参数;SFAS 是动态+退火随机地按 Fisher 在线挑参数,能跟随气象 regime 切换重标定,且额外配了 TADP 的上下文感知,本文实验显示比纯静态选参更强。
- vs 全量微调(Full-Tuning):Full-Tuning 每任务存一整套 12 亿参数、算力存储不可持续;WeatherPEFT 用 0.3%–4% 参数追平甚至超越它,是面向多任务 WFM 部署的现实方案。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首个系统研究 WFM 高效微调,TADP 从嵌入权重生成 prompt + SFAS 退火 Fisher 选参都针对气象异质性原创设计。
- 实验充分度: ⭐⭐⭐⭐⭐ 三类任务、十余个 PEFT baseline、两个 WFM、消融与放大预算设置齐全。
- 写作质量: ⭐⭐⭐⭐ 动机和方法层次清晰、公式完整,部分换轴/注入细节稍需对照图才好理解。
- 价值: ⭐⭐⭐⭐⭐ 让 10 亿级天气基础模型可在小参数预算下落地多任务,对气象 AI 实用部署意义大。