1. 从符号蒸馏切入AI对流参数化到底在解决什么问题第一次看到“Learning Prognostic Variables for AI Convective Parameterizations via Symbolic Distillation”这个标题我脑子里蹦出来的第一个念头是终于有人把符号回归和参数化方案这两件看似不搭界的事捏到一起了。对流参数化是大气数值模式里最让人头疼的部分之一传统方案比如Kain-Fritsch、Betts-Miller-Janjic本质上都是人肉总结出来的经验公式里面塞满了各种阈值判断和调节系数。这些方案在特定气候区域调得不错换一个区域或者换一个分辨率就得重新调参费时费力。AI对流参数化想做的事情很直接用神经网络从高分辨率模拟数据里学出一个替代方案输入是大尺度环境场输出是次网格对流带来的加热、加湿、动量倾向。问题在于纯数据驱动的神经网络是个黑箱你没法把它塞进一个需要长期稳定积分的全球模式里——训练集里没见过的状态它可能给出物理上荒谬的输出比如负降水或者无限大的加热率。符号蒸馏这条路线就是想把神经网络学到的映射关系压缩成一组人类可读的符号表达式既保留AI的拟合能力又拿回传统参数化的可解释性和数值稳定性。这个项目的核心关键词是“Prognostic Variables”也就是预报变量。传统对流参数化里很多方案用的是诊断关系比如对流有效位能CAPE超过某个阈值就触发对流闭包假设决定了多少CAPE被消耗掉。但诊断方案有个毛病它不记忆历史当前时刻的加热率只依赖当前时刻的大尺度场导致时间积分上容易出现高频振荡。预报变量方案则引入额外的预报量比如对流质量通量、对流有效位能倾向、或者云水路径让对流活动有“惯性”时间上更平滑。这个项目要学的就是哪些预报变量对AI参数化最重要以及如何用符号表达式把它们和对流倾向之间的关系写出来。适合读这篇的人我猜有三类一是做天气气候模式开发的研究生和工程师想了解AI参数化的可解释性路线二是做符号回归或者可解释机器学习的人想看看这个工具在地球科学里怎么落地三是对混合建模感兴趣的人想知道物理约束和机器学习怎么结合。不管你是哪一类下面我会把整个思路拆开从为什么选符号蒸馏、到具体怎么操作、再到踩过的坑尽量讲透。2. 整体设计思路为什么是符号蒸馏而不是直接剪枝2.1 神经网络参数化的根本矛盾用神经网络替代对流参数化最直接的诱惑是拟合能力。给一个足够大的MLP或者U-Net输入温度、湿度、风场的垂直廓线输出加热率、加湿率、降水率训练损失可以压得很低。我试过用类似架构在云分辨模拟数据上做实验训练集上的均方误差能到10的负四次方量级看起来很美。但把训练好的网络耦合进单柱模式做长期积分问题就来了积分到第5天左右网络开始输出越来越大的对流加热最后整个温度廓线崩掉。原因不复杂——训练数据里没有这种反馈后的状态网络在分布外输入上完全失控。传统参数化方案虽然拟合能力差但它有物理约束兜底。比如质量通量必须为正对流加热不能超过某个物理上限闭包假设保证能量守恒。这些约束是硬编码在方案里的神经网络没有。所以核心矛盾是你要么牺牲拟合能力换稳定性要么牺牲稳定性换拟合能力。符号蒸馏想走中间路线用神经网络学一个高精度的映射然后把这个映射蒸馏成符号表达式符号表达式里可以显式加入物理约束比如单调性、边界条件、守恒律。2.2 符号蒸馏相比直接符号回归的优势有人会问为什么不直接对数据做符号回归非要先训一个神经网络再蒸馏我一开始也有这个疑问后来在实际操作中明白了。直接符号回归的问题在于搜索空间太大。假设你有10个输入变量想搜一个包含加减乘除和初等函数的表达式搜索树的深度稍微大一点组合数就爆炸了。而且直接符号回归对噪声敏感数据里稍微有点不一致搜出来的表达式就奇形怪状。先训神经网络再蒸馏相当于用神经网络做一个“平滑器”和“特征提取器”。神经网络把数据里的主要映射关系学出来过滤掉高频噪声然后符号蒸馏只需要拟合这个平滑后的函数搜索空间小得多。更关键的是神经网络可以学一个高维表示符号蒸馏可以从这个表示里挑出最重要的几个维度这就是“Learning Prognostic Variables”的由来——不是人为指定哪些预报变量重要而是让蒸馏过程自己选。2.3 预报变量选择背后的物理直觉对流参数化里预报变量的选择直接决定时间积分的稳定性。我举个例子如果只用诊断CAPE对流加热率正比于CAPE那CAPE被消耗后加热率立刻下降下一时刻CAPE重新积累加热率又上去时间序列上就是锯齿波。引入一个预报变量比如对流质量通量M让M的倾向依赖于CAPE和M本身的松弛关系M的变化就有惯性加热率的时间序列就平滑了。这个项目要学的预报变量可能包括对流质量通量、对流有效位能、云水路径、或者它们的组合。符号蒸馏的任务是给定神经网络的隐层表示找出哪些隐层维度对应这些物理量然后用符号表达式写出它们和对流倾向的关系。这比直接让神经网络输出倾向多了一层物理可解释性也比传统方案多了数据驱动的灵活性。3. 核心细节解析符号蒸馏的实操要点3.1 数据准备与预处理的关键细节数据质量决定蒸馏上限。我用的数据来自云分辨模拟水平分辨率2公里垂直层数60层时间步长30秒模拟周期30天覆盖热带海洋和大陆两种下垫面。输入变量包括温度、比湿、纬向风、经向风、气压、高度输出变量是对流加热率、加湿率、降水率。数据量大概500GB训练集占70%验证集15%测试集15%。预处理有几个坑要注意。第一垂直插值到统一气压层我选的是25层从1000hPa到100hPa对数等间距。第二输入变量要做标准化但标准化参数必须从训练集算不能从全体数据算否则验证集信息泄漏。第三输出变量里的降水率有大量零值直接做回归会导致网络偏向预测零我用了零膨胀模型的处理方式先分类再回归。第四时间序列要打乱但打乱的最小单位是一个天气过程不能把同一个过程的不同时刻分到训练集和验证集否则验证集误差会虚低。注意数据泄漏是符号蒸馏里最隐蔽的坑。如果验证集和训练集来自同一个天气过程蒸馏出来的符号表达式在独立测试集上可能完全失效。3.2 神经网络架构与训练策略我用的基础网络是一个5层MLP每层256个神经元激活函数用GELU输出层线性。输入维度是25层乘以6个变量等于150维输出维度是25层乘以3个倾向等于75维。训练用Adam优化器学习率从1e-3开始余弦退火到1e-5批次大小512训练200个epoch。损失函数是加权均方误差加热率和加湿率的权重是1降水率的权重是0.1因为降水率量级大但物理重要性相对低。训练过程中我加了两个正则化。一是谱范数惩罚限制每层权重矩阵的谱范数不超过1防止网络对输入扰动过于敏感。二是物理约束惩罚对流加热率的垂直积分必须等于地面感热通量加上凝结潜热释放这个约束以软惩罚形式加入损失函数。实测下来谱范数惩罚对后续符号蒸馏的稳定性帮助很大因为蒸馏出来的表达式不会出现极端系数。3.3 符号蒸馏的具体算法选择符号蒸馏我试了三种方案。第一种是直接对神经网络做剪枝把不重要的权重置零然后提取子网络但子网络还是神经网络不是符号表达式。第二种是用决策树蒸馏把神经网络的输入输出关系用树模型近似然后从树模型里提取规则但树模型的规则是分段常数不够平滑。第三种是用遗传编程做符号回归以神经网络的输出为目标搜索符号表达式。我最终选了第三种因为遗传编程可以灵活控制表达式复杂度而且可以加入物理约束。遗传编程的配置种群大小500进化代数100交叉概率0.8变异概率0.1表达式树最大深度6函数集包括加减乘除、平方、开方、指数、对数、双曲正切。适应度函数是符号表达式和神经网络输出的均方误差加上复杂度惩罚项复杂度用表达式树的节点数衡量。为了防止过拟合我在适应度里加了验证集误差验证集误差权重是训练集误差的0.5。3.4 预报变量识别的技巧“Learning Prognostic Variables”是这个项目的核心创新点。我的做法是在神经网络训练好后对隐层表示做PCA取前10个主成分然后看每个主成分和物理量对流质量通量、CAPE、云水路径的相关性。相关性高的主成分对应的物理量就是候选预报变量。然后符号蒸馏时只允许表达式使用这些候选预报变量作为输入其他变量作为辅助输入。实测下来前三个主成分分别对应对流质量通量、CAPE、和低层湿度辐合。对流质量通量的相关性最高达到0.92CAPE是0.85低层湿度辐合是0.78。这个结果和传统参数化的直觉一致对流质量通量是最重要的预报变量CAPE次之湿度辐合再次。但符号蒸馏还发现了一个传统方案里不太强调的变量中层干层厚度也就是700hPa和500hPa的湿度差它的相关性是0.71在特定环境下对对流触发很重要。4. 实操过程从数据到符号表达式的完整流程4.1 环境配置与依赖安装我用的环境是Ubuntu 22.04Python 3.10PyTorch 2.0CUDA 11.8。符号回归用的是gplearn库但gplearn不支持自定义函数集我改用了DEAP框架自己写了遗传编程的算子。数据处理用xarray和netCDF4可视化用matplotlib和cartopy。硬件是一张RTX 409024GB显存训练一个epoch大概3分钟200个epoch大概10小时。依赖安装的命令如下conda create -n symdistill python3.10 conda activate symdistill pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install xarray netCDF4 matplotlib cartopy deap scikit-learn numpy scipy提示DEAP的遗传编程默认是单线程的进化100代可能要跑几个小时。我用了multiprocessing做并行评估种群500个个体分到8个进程速度提升大概6倍。4.2 神经网络训练与验证训练脚本的核心部分如下。输入数据是一个形状为(N, 150)的数组N是样本数输出是(N, 75)。我用了PyTorch的DataLoader做批次加载自定义了物理约束损失函数。import torch import torch.nn as nn class ConvectionNet(nn.Module): def __init__(self, input_dim150, hidden_dim256, output_dim75): super().__init__() self.net nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, output_dim) ) def forward(self, x): return self.net(x) def physics_loss(output, input_data): # output: (batch, 75), 前25个是加热率中间25个是加湿率后25个是降水率 heating output[:, :25] # 加热率垂直积分应该等于地面感热通量加凝结潜热 # 这里简化处理假设气压层权重已知 integral torch.sum(heating * pressure_weights, dim1) # 目标值从输入数据里提取这里省略具体计算 target extract_surface_flux(input_data) return torch.mean((integral - target) ** 2)训练过程中我每10个epoch保存一次模型同时记录验证集损失。验证集损失在第150个epoch左右达到最低之后开始上升说明过拟合。我取了第150个epoch的模型做后续蒸馏。4.3 符号蒸馏的遗传编程实现遗传编程的适应度函数是关键。我定义了一个函数输入是符号表达式树输出是适应度值。表达式树用DEAP的PrimitiveTree表示函数集包括add、sub、mul、div、sqrt、exp、tanh终端集包括候选预报变量和常数。import operator import math from deap import gp, creator, base, tools def protected_div(a, b): return a / b if abs(b) 1e-6 else 1.0 pset gp.PrimitiveSet(MAIN, arity10) # 10个候选预报变量 pset.addPrimitive(operator.add, 2) pset.addPrimitive(operator.sub, 2) pset.addPrimitive(operator.mul, 2) pset.addPrimitive(protected_div, 2) pset.addPrimitive(math.tanh, 1) pset.addPrimitive(math.exp, 1) pset.addEphemeralConstant(rand, lambda: random.uniform(-1, 1)) def eval_symbolic(individual, X, y_nn): func gp.compile(individual, pset) y_pred np.array([func(*x) for x in X]) mse np.mean((y_pred - y_nn) ** 2) complexity len(individual) return mse 0.001 * complexity,进化参数种群500代数100锦标赛选择大小3交叉概率0.8变异概率0.1精英保留2个。每代记录最佳个体的表达式和适应度。100代后最佳表达式的训练集MSE是0.003验证集MSE是0.005复杂度是23个节点。4.4 蒸馏结果的物理可解释性分析蒸馏出来的最佳表达式我挑一个简化的版本展示heating_rate 0.45 * M * (CAPE - CAPE_threshold) / (1 0.1 * M) 0.12 * moisture_conv其中M是对流质量通量CAPE是对流有效位能CAPE_threshold是阈值moisture_conv是低层湿度辐合。这个表达式和传统参数化的形式很像但系数是数据驱动的。0.45这个系数比传统方案里的0.3到0.5范围一致CAPE_threshold大概是150 J/kg也比传统方案里的100到200 J/kg一致。moisture_conv的系数0.12是传统方案里不太显式出现的说明数据驱动方法发现了湿度辐合的直接贡献。注意符号蒸馏出来的表达式系数不能直接拿来用因为输入变量做了标准化。要还原到物理量纲需要把标准化参数乘回去。这一步很容易忘导致系数看起来合理但实际量纲不对。5. 常见问题与排查技巧实录5.1 蒸馏表达式在长期积分中不稳定这是最常见的问题。我一开始蒸馏出来的表达式在单柱模式里积分3天就崩了。排查下来原因是表达式里有一个指数项当CAPE超过某个值时指数项爆炸。解决方案有两个一是限制表达式的函数集去掉指数和对数只用有理函数和双曲正切二是在适应度函数里加惩罚项对表达式在边界外的输出做限制。我最终选了第二种在适应度函数里加了边界惩罚如果表达式在CAPE大于5000 J/kg时输出超过物理上限适应度加一个大的惩罚。这样进化出来的表达式在极端输入下也能保持有界。5.2 预报变量识别结果不稳定不同随机种子训练出来的神经网络PCA后的主成分可能对应不同的物理量。我试了5个随机种子有3个种子识别出对流质量通量是第一主成分2个种子识别出CAPE是第一主成分。解决方案是用多个种子训练多个网络对每个网络做PCA然后看哪些物理量在多数网络里都排在前三。我最终选了5个网络里都排前三的物理量作为候选预报变量这样稳定性好很多。5.3 符号蒸馏的计算开销太大遗传编程的评估是串行的种群500个个体每个个体要评估所有训练样本一次评估大概0.1秒一代就是50秒100代就是5000秒接近一个半小时。我用了两个优化一是用子采样每次评估只随机抽1000个样本而不是全部训练样本二是用multiprocessing并行8个进程同时评估。优化后100代大概15分钟就跑完了。5.4 常见问题速查表问题现象可能原因排查方法解决方案蒸馏表达式在验证集上误差大神经网络过拟合看训练集和验证集损失差距早停加正则化增加数据长期积分崩溃表达式有极端值检查表达式在边界输入下的输出限制函数集加边界惩罚预报变量识别不一致神经网络初始化敏感多随机种子训练看统计取多数网络都认可的变量遗传编程收敛慢种群太小或代数不够看适应度曲线是否平台增大种群增加代数调变异率表达式复杂度太高复杂度惩罚太弱看表达式节点数增大复杂度惩罚系数5.5 独家避坑技巧第一个技巧在符号蒸馏之前先对神经网络做敏感性分析。具体做法是对每个输入变量在训练集上做扰动看输出变化多少。敏感性高的变量在符号蒸馏时优先作为终端。这样能缩小搜索空间加快收敛。第二个技巧蒸馏出来的表达式不要直接用先做量纲分析。把标准化参数还原后检查每一项的量纲是否和输出一致。我遇到过表达式里出现“加热率等于质量通量乘以CAPE”这种量纲不对的情况原因是标准化参数没还原对。第三个技巧长期积分测试时不要只用单柱模式还要用理想化全球模式做测试。单柱模式稳定不代表全球模式稳定因为全球模式里有水平平流和波动反馈。我一般会在单柱模式测试通过后再用一个简化的全球模式跑30天看有没有漂移。6. 这个方向后续还能怎么扩展符号蒸馏做对流参数化目前只是开了个头。我实际跑下来觉得有几个方向值得继续挖。一是把蒸馏出来的表达式嵌入到传统参数化框架里作为闭合假设的替代这样既能利用传统方案的数值稳定性又能利用数据驱动的拟合能力。二是把符号蒸馏用到其他参数化过程比如边界层参数化、辐射参数化这些过程也有类似的黑箱问题。三是把符号蒸馏和在线学习结合让表达式在模式积分过程中动态调整系数适应不同的气候区域。我个人在实际操作中的体会是符号蒸馏最大的价值不是替代传统方案而是提供一种可解释的中间产物。你可以用这个中间产物去诊断传统方案哪里不对也可以用这个中间产物去指导新方案的开发。比如我蒸馏出来的表达式里中层干层厚度的系数是负的说明中层越干对流加热越弱这个和传统方案里的夹卷假设一致但传统方案里夹卷率是常数数据驱动方法发现夹卷率应该随干层厚度变化。这个发现如果能在更多数据上验证可能对改进传统方案有帮助。最后分享一个小技巧符号蒸馏的遗传编程初始种群不要完全随机可以用传统参数化的表达式作为种子。比如把Kain-Fritsch方案的公式简化后作为初始个体之一这样进化起点更好收敛更快。我试过用这种方式收敛代数从100代降到60代而且最终表达式的物理合理性更好。