[{"content":"传输计算是大气传输过程的重要组成部分，总体上可以分为垂直扩散、垂直对流、水平扩散、水平对流四个部分，而四个部分计算范式截然不同，在 CMAQ 当中是污染物传输和扩散的重要求解过程，其求解精度影响多个下游模块。优化计算速度的同时，保证传输部分高精度以不影响下游模块计算是大气模拟的重要条件。\n而水平传输是传输模块中的计算密集型部分，其使用双向交替的二遍算子分裂，每一个子时间步都运行 X 和 Y 方向的分段抛物线计算。同时其嵌套使用同步步长和子步步长，受 Courant 数约束，大部分子步被切得很小，在一个同步步内要循环很多次，每一步都要进行 X 和 Y 方向的分段抛物线计算。所以单位时间（每小时）的计算量不是恒定的，而是随天气强度（风速）动态增长的——风越大，算得越慢、越费时。\n同时计算过程会产生大量 I/O 操作，读取逐步长、逐层的输入数据，缺乏必要的算子融合。PPM 抛物线算法也缺乏向量化（SIMD 范式计算），编译器基本上只能生成标量代码，而且访问模式是那种很吃内存带宽的 gather/scatter（聚集/散射）模式。整个传输计算过程高度碎片化，存在大量条件判断、子步迭代、非向量化操作，以此用统一的机器学习模型对传输模型进行替代，将迭代求解转换成单次前向传播，配合 GPU 的高速张量计算，我们可以把传输部分的同步时间步的单次计算耗时降低至原来的 1/10。\n模型与数据构造 我们用一个简单的 3D CNN 网络学习 CMAQ 输运算子（transport-only，不含化学/排放），通过 rollout training 解决一个关键问题：单步精度好但自回归链式推理时误差累积爆炸。替代目标是传输模块当中的 HADV、ZADV、HDIFF 三个模块，将其融合为一个统一的机器学习模型。通过选取真实 CMAQ 运行案例，在模块开始与结束位置进行插桩，转储输入输出数据，构建真实运行场景数据集。输入变量包含 U、V 方向风场（来自 MCIP 数据）、浓度张量、层数、Sigma Layer Thickness 等变量，输出 Arakawa-C 网格边界面物质传输通量，通过将物种作为 Batch 维度，让所有物种共享同一套权重，从而实现物种并行（Species Parallel）。\n对于回归链式推理时误差累积爆炸，我们给出的答案是大规模长程轨迹合成数据。通过将目标模块单独隔离出来，用 Python 重写成一个独立运行的模块，采样转储的输入数据，测试重写的部分能与原代码保持 bit 级一致，通过真实运行案例的传输数据作为输入，定义为 $\\mathrm{input}_{t0}$，由此得到 $\\mathrm{output}_{t0}$，将 $\\mathrm{output}_{t0}$ 作为下一步迭代的输入 $\\mathrm{input}_{t1}$，由此不断循环，最终我们能得到一段 $K=N$ 的高保真自迭代轨迹数据，起始点由真实环境采样得到，同时通过拉丁超立方采样，我们可以探索输入向量空间的更多可能性，从而形成拓展合成轨迹数据（Extended synthetic trajectory data）。\n损失函数与 rollout training 对于 loss 函数来说：\n$$ L_{\\mathrm{ss}} = \\frac{1}{N}\\sum_{k=0}^{N-1}\\left\\Vert NN\\!\\left(C_b^{(k)}\\right) - C_a^{(k)} \\right\\Vert_2^2 \\tag{1} $$加入 rollout 项之后：\n$$ L = L_{\\mathrm{ss}} + \\lambda_{\\mathrm{roll}} \\cdot L_{\\mathrm{roll}}, \\qquad \\lambda_{\\mathrm{roll}} = 0.5 \\tag{2} $$同时在 rollout training 当中，关键的训练配方是同步计算多步自迭代预测的损失，而只针对第一步的梯度来更新权重，这样模型既能学习到第一步的预测精度，同时又对多步预测做好准备，综合损失函数在多步损失当中会表现平滑。需要注意的是：单步的高精度未必是多步预测稳定的必要条件，机器学习模型对输入存在不同程度的敏感性，在自迭代链式预测过程中，微小误差会随着多步累积而放大，这在控制论当中被称之为自激震荡。犹如一个高保真的麦克风，无论多还原音色，靠近音响的时候也难免产生刺耳的啸叫声。\n所以在长程合成轨迹模型训练当中，一个在第一步 loss 很高的模型未必会在多步预测的最后一步 loss 很高，模型会逐渐学会多步轨迹预测损失的平滑值，从而避免因为自激震荡导致数值爆炸的问题。我们采用混合步长作为合成数据配方，从 $K=3$ 到 $K=20$ 不等，在训练过程中混合调配样本比例，避免在长程监督学习中产生\u0026quot;灾难性遗忘\u0026quot;。下图很好地解释了不同训练步长导致的 loss 变化：\n图 1：代表物种自迭代链式推理是否使用多步训练在离线评测当中的差异增长\n如果我们不选择使用 rollout training，对于以 NO2 为例的物种，其相对误差在自迭代 4 个时间步后便急剧放大，以模型 Sync Step = 600 seconds 的条件下，最多只能撑过 40 分钟便会产生灾难性的数值爆炸现象。这里为方便说明起见，我们选择一个不采用混合步长训练的 $K=8$ 的小参数模型作为训练目标，可以看到 $K=8$ 的训练效果能让多步误差累积更加平滑，差不多能抑制约 5 倍的相对误差。而 $K=8$ 的 v6 版本在第一个时间步的相对误差甚至更高！关于 $K=N$ 的步长数量和样本比例的训练配方，并没有一个比较科学的依据，受训练时间限制，最终产出模型采用了从 $K=5$ 到 $K=20$ 的混合比例，通过 Top_K 随机采样在训练集当中进行训练。\n图 2：仅移植传输模块于在线模拟中单个同步步耗时对比\n我们以一块 NVIDIA A800 GPU + 4 个英特尔 Intel® Xeon® Gold 6326 CPU 核心作为测试硬件基准，于在线测试中测试模型在传输模块的墙钟模拟时间，可以发现在单模块实现了约 8.9 倍的加速。我们的测试案例采用 $292 \\times 273 \\times 14$ 约 116 万个网格的模拟配置，总有效物种数 867 个，模拟网格分辨率 27 公里。通过有效的机器学习做算子融合和 GPU 并行推理加速，将物种放入 Batch 维度，实现了高效的物种并行，单次前向传播计算可同时预测 867 个物种的界面通量。配合尾端的物种质量守恒性校验，可以更加高效地验证物理约束是否吻合。\n启发式强化学习：如何改进大气机器学习代理的自迭代链式推理？ 我们的目标是构建这样一个链条：NN surrogate → physics diagnosis → correct/retry/fallback → continue rollout。\n短期预测看似准确，但在 7 天至 1 个月的运行过程中，误差会逐渐累积，尤其对于 NOx、铵、硝酸盐、PM 等化学类别的预测更为明显。以传输模块为例，顺序是 HADV → ZADV → HDIFF，state（状态）→ action（动作）→ criticize（评价，奖励或惩罚，这里用物理信息引导，可以参考的标准有原子守恒、前体物依赖顺序、散度/梯度等）→ strategy update（策略更新，我们的案例中用 PPO）→ next state（下一步的状态）。\n图 3：物理诊断引导的自迭代修正链路\n对于一个长程轨迹优化问题来说，传统的监督学习（supervised learning）能给模型一个比较好的监督信号，这种信号很密集、通过优化器我们能知道梯度变化方向，从而优化模型权重来使得 loss 函数更小，这是一个机器学习的入门知识。但是监督学习不能解决的问题是：优化器结合当前梯度优化的方向，是否真的符合真实大气环境的物理规律？ 权重迭代的方向会使得模型预测趋于平滑，这样会让全局损失函数更小，但这犹如农场里的火鸡，能观察到\u0026quot;农场主打开门 → 会有食物\u0026quot;这样一个前后关系，但是却没有真的学会物理规律——有食物吃的原因是每天农场主都要养着它们，而总有一天打开门的不是食物，而是屠刀。机器学习从海量的数据当中学会了统计规律，却不能真的理解背后的逻辑和因果关系。\n由此：机器学习模型能够在我们的大气传输模型中进行长程训练，但是仅仅有监督信号是不够的，还要有物理信息引导的强化学习，按照物理规律设置奖励/惩罚函数。\n一个具体的问题：物理引导的 PPO（近端策略优化）如何处理 图 4：物理引导 PPO 的训练流程\n也就是说，模型从当前状态 $s_t$ 出发，根据策略 $\\pi_\\theta$ 产生连续动作 $a_t$（连续动作空间，因为大气内部的变量都是连续的，我们这里的动作空间输出的是高斯分布），经过数值演化得到下一状态 $s_{t+1}$，并将预测误差与质量守恒、原子守恒、梯度一致性和散度等物理约束共同构造成 physics-guided reward，即：\n$$ r_t = -L_{\\mathrm{pred}} - \\lambda L_{\\mathrm{physics}} \\tag{3} $$与传统单步监督学习不同，模型会利用自身预测结果不断向前推进，形成完整的 rollout trajectory：\n$$ \\tau = (s_0, a_0, s_1, a_1, \\ldots, s_T) \\tag{4} $$因此训练目标关注的是整个长期轨迹的累计回报，而不仅是下一时间步的预测精度。最后，PPO 根据 rollout 得到的 reward 和 advantage 判断哪些动作有利于长期预测准确性和物理一致性，提高这些动作在相似状态下出现的概率，同时通过 clipping 限制每次策略更新幅度，从而逐步训练出一个在长期积分过程中既保持预测精度，又满足关键物理约束并减少误差累积的 surrogate model。\n对于 PPO 来说，核心目标函数是 $L_{\\mathrm{PPO}}(\\theta)=\\mathbb{E}_t\\left[\\min\\left(r_t(\\theta)\\hat{A}_t,\\ \\operatorname{clip}\\!\\left(r_t(\\theta),1-\\epsilon,1+\\epsilon\\right)\\hat{A}_t\\right)\\right]$，其中最关键的是两个量。\n首先是新旧策略的概率比 $r_t(\\theta)=\\pi_\\theta(a_t\\,|\\,s_t)\\,/\\,\\pi_{\\theta_{\\mathrm{old}}}(a_t\\,|\\,s_t)$，例如某个连续 correction $a_t$ 在 rollout 中表现很好，新策略就希望增加产生这个 action 的概率，因此 $r_t(\\theta)\u003e1$，反之则降低其概率。\n第二个量是 advantage $\\hat{A}_t$，它回答的是：在状态 $s_t$ 下采取 $a_t$，从长期 rollout 的结果来看，比平均水平好还是差？框架里，这个\u0026quot;好坏\u0026quot;最终来自前面定义的 physics-guided reward $R_t=-\\mathcal{L}_{\\mathrm{pred}}-\\lambda\\mathcal{L}_{\\mathrm{physics}}$，因此可以把整个逻辑理解成 Physics Reward → Rollout Return → $\\hat{A}_t$ → PPO Update → $\\pi_{\\theta_{\\mathrm{new}}}$：例如一个 action 当前时刻预测误差很小，而且使后续几十、几百步的质量守恒和原子守恒都更好，那么它的长期 advantage 可能满足 $\\hat{A}_t\u003e0$，PPO 就会提高类似 action 的概率；如果某个 action 单步看起来不错，但最终导致 NOy 在长期 rollout 中持续漂移，那么它可能得到 $\\hat{A}_t\u003c0$，PPO 就会降低类似 action 的概率。\n而公式里的 $\\operatorname{clip}$ 是 PPO 很重要的一点：即使发现某个 action 很好或很差，也不允许策略一次改变得太剧烈，例如 $\\epsilon=0.2$，就限制每轮优化对策略进行相对保守的调整，从而提高训练稳定性。\n","permalink":"https://cmaq-notes.pages.dev/posts/cmaq-transport-surrogate-rollout/","summary":"\u003cp\u003e传输计算是大气传输过程的重要组成部分，总体上可以分为垂直扩散、垂直对流、水平扩散、水平对流四个部分，而四个部分计算范式截然不同，在 CMAQ 当中是污染物传输和扩散的重要求解过程，其求解精度影响多个下游模块。优化计算速度的同时，保证传输部分高精度以不影响下游模块计算是大气模拟的重要条件。\u003c/p\u003e","title":"大气传输模型机器学习代理的长程轨迹优化"},{"content":"本站记录大气化学传输模式（CMAQ）与机器学习交叉方向的工作\n","permalink":"https://cmaq-notes.pages.dev/about/","summary":"关于本站","title":"关于"}]