1. 项目概述当Transformer遇见系外行星大气物理如果你从事系外行星大气研究或者对计算天体物理感兴趣那么“辐射传输”这个词对你来说一定不陌生。它几乎是所有大气光谱模拟、能量平衡计算和气候模型的核心引擎。简单来说它描述的是光辐射在大气中如何被气体分子、云和气溶胶吸收、散射和发射的过程。我们通过望远镜接收到的系外行星光谱就是辐射传输过程的最终“指纹”。然而这个物理过程的数值模拟长期以来都是计算上的“拦路虎”。传统的辐射传输求解器比如业内广泛使用的DISORT离散纵标法或PICASO中集成的算法虽然精度高但计算成本极其昂贵。模拟一个包含几十个大气层的行星光谱往往需要数秒甚至数分钟。当我们需要进行大规模参数空间扫描比如探索不同化学成分、温度剖面、云层假设的组合或者将辐射传输模块耦合到需要成千上万次迭代的全球气候模型GCM中时这种计算开销就变得完全不可接受。这就像你想用一台老式算盘去处理现代大数据分析理论可行但效率会让你崩溃。近年来深度学习特别是Transformer架构在序列建模和物理信息学习方面展现出了惊人的潜力。我们团队就在思考能否用Transformer来“学习”辐射传输这个复杂的物理过程构建一个高速、高精度的“模拟器”或“代理模型”这个想法并非天方夜谭。辐射传输问题本质上可以看作是一个序列到序列的映射输入是一系列描述大气状态如每层的温度、压力、成分浓度的序列输出是每层的净辐射通量或最终逃逸到太空的光谱。Transformer的自注意力机制天生擅长捕捉序列元素间的长程依赖关系而这恰恰是辐射传输中“非局域性”的关键——高层大气的状态会直接影响低层的辐射场反之亦然。我们基于这个思路开发了一个基于Transformer的辐射传输模拟器。实测下来这个模型在预测大气层净热通量和散射星光通量时平均误差可以控制在1%左右与PICASO这样的高精度标准工具的结果高度吻合。而最令人兴奋的是它的推理速度达到了毫秒级相比传统方法提升了数个数量级。这意味着以前需要跑几天的参数研究现在可能喝杯咖啡的功夫就完成了。这对于处理詹姆斯·韦布空间望远镜JWST等新一代设备产生的海量高精度光谱数据以及构建更复杂的耦合气候模型无疑是一个游戏规则的改变者。2. 核心思路为什么是Transformer而不是LSTM或全连接网络在决定使用Transformer之前我们实际上评估过多种神经网络架构包括全连接网络FNN和长短期记忆网络LSTM。这是一个关键的技术选型决策其背后的逻辑直接关系到模型的成败。2.1 传统方法的瓶颈与代理模型的必要性首先为什么我们需要一个代理模型以经典的DISORT算法为例它通过求解辐射传输方程在多个离散角度和波长上的积分来获得通量。这个过程涉及大量的矩阵运算和迭代计算复杂度与大气层数、角度离散数和光谱分辨率呈多项式甚至指数增长。在PICASO等工具中为了获得一个光谱需要在数百个波长点上分别调用这样的计算。当我们进行大气反演从观测光谱推断大气参数时需要成千上万次的正向模型调用计算负担成为主要瓶颈。代理模型或称为仿真器、替代模型的思路是用大量的“输入-输出”数据对由高精度但缓慢的模型生成来训练一个快速的机器学习模型。训练阶段耗时但一旦训练完成推理预测阶段可以极快。问题的核心就变成了我们选择的机器学习模型能否足够精确地学习到从大气参数到辐射通量之间复杂的物理映射关系2.2 架构对比FNN、LSTM与Transformer全连接网络FNN这是最直接的尝试。我们将所有层的参数压力、温度、混合比等展平成一个长向量输入网络。但这种方法存在明显缺陷它完全破坏了大气层的垂直结构信息序列顺序。大气物理具有强烈的垂直梯度底层和高层的物理过程截然不同。FNN难以有效学习这种具有空间顺序和层次关系的特征对于非局部相互作用如高层云对底层热通量的影响的捕捉能力很弱。长短期记忆网络LSTM作为经典的循环神经网络变体LSTM是为序列数据设计的。它按顺序处理每一层大气信息并通过隐藏状态传递上下文。这比FNN更符合物理直觉。我们初期也复现了类似Ukkonen2022工作中使用的LSTM架构。然而在实际测试中尤其是当大气层数较多我们用了75层时LSTM暴露出两个问题一是训练速度较慢因为其顺序处理机制无法充分利用GPU的并行计算能力二是对于非常长程的依赖关系信息在传递过程中容易衰减或丢失即长期依赖问题。尽管有门控机制但在捕捉跨越数十个大气层的非局部辐射相互作用时其表现不够稳健。Transformer编码器最终我们选择了仅使用编码器部分的Transformer架构类似BERT但不用于生成而是用于回归/预测。它的核心优势在于自注意力机制和并行计算。自注意力对于输入的大气层序列每层用一个特征向量表示自注意力机制允许每一层“直接看到”序列中所有其他层并计算一个加权和。这个权重是动态学习得到的代表了该层与其他所有层在辐射传输过程中的“关联强度”。这完美契合了辐射传输的“非局域性”——某一层的辐射源函数不仅取决于本地温度还受到上下所有层辐射场的影响。Transformer能显式地建模这种全序列的相互作用这是LSTM隐式传递所难以媲美的。并行计算由于自注意力不依赖于前一时刻的计算结果整个序列可以同时进行处理。这使得Transformer在GPU上的训练和推理效率远高于LSTM尤其适合我们这种需要处理大量训练样本的场景。位置编码为了弥补自注意力本身不具备顺序感知能力的缺点我们向输入特征中添加了正弦余弦位置编码。这为模型提供了每一层在大气中所处位置如压力坐标的对数的信息使其能够理解“顶层”和“底层”的物理差异。注意模型选择的核心权衡。选择Transformer并非否定LSTM。对于层数较少如30层或序列依赖模式更局部化的问题LSTM可能更简单有效。但针对系外行星大气这种层数多、非局部效应强的物理问题Transformer在精度和效率上的综合优势更为明显。我们的对比实验显示在相同的数据集和训练时长下Transformer在测试集上的均方误差比LSTM降低了约30-40%。3. 模拟器设计与实现细节我们的目标不是构建一个通用的、包罗万象的模型而是在一个定义明确的物理框架内证明Transformer作为辐射传输核的可行性。因此我们首先限定了问题的范围。3.1 问题定义与数据生成我们聚焦于一维、平面平行、热辐射与恒星散射辐射分离处理的大气模型。这是目前系外行星大气光谱模拟中最常用的近似。对于每一层大气我们关心两个核心的输出量层净热通量和层净射星光通量。这两个通量是计算大气温度结构和最终观测光谱的基础。训练数据由高精度的参考模型PICASO生成。我们固定了一系列物理参数行星类型类木星气态巨行星、恒星类型类太阳、大气化学平衡化学并暂时不考虑云和气溶胶。这样做的目的是控制变量首先在最“干净”的场景下验证架构的有效性。我们通过拉丁超立方采样等方法在合理的参数空间内随机生成了数万个不同的大气温度-压力剖面。对于每一个剖面调用PICASO计算其75个大气层中每一层的向上和向下热通量、向上和向下散射星光通量。最终我们的训练数据集大小不到5GB包含了输入温度、压力、混合比剖面和输出各层通量的对应关系。3.2 Transformer模型架构详解我们采用了一个相对轻量化的编码器-解码器结构实际上解码器是一个简单的多层感知机MLP。输入嵌入层每个大气层共75层的原始物理参数如对数压力、温度、主要气体H2/He的混合比等通过一个线性层映射到更高维的模型表示空间例如128维。同时我们将该层对应的正弦位置编码基于归一化的压力层级加到嵌入向量上。Transformer编码器堆叠嵌入后的序列被送入一个由N个例如6个相同的编码器层堆叠而成的模块。每个编码器层包含多头自注意力层我们使用了8个注意力头。每个头可以学习到不同子空间下的层间依赖关系。例如有的头可能专注于捕捉温度梯度对热辐射的影响有的头则可能学习到不透明度突变层对散射光的遮挡效应。所有头的输出被拼接并线性变换。前馈神经网络一个简单的两层MLP作用于每个位置的向量上用于进行非线性特征变换。残差连接与层归一化每个子层自注意力和前馈网络周围都应用了残差连接和层归一化。这是稳定深层Transformer训练的关键能有效缓解梯度消失问题。序列池化与输出层编码器输出的是一个75x128的序列。我们需要为每一层预测两个标量值热通量和散射通量。这里有两种常见做法一是为每一层位置单独接一个小的MLP位置相关的输出二是将整个序列的信息通过全局平均池化或添加一个特殊的[CLS]标记来汇聚再映射到输出。我们选择了前者因为我们需要的是逐层的物理量。具体来说我们取编码器输出的第i个位置向量通过一个两层的MLP直接回归出该层的两个通量值。这样模型在计算第i层的输出时已经通过自注意力机制充分考虑了所有其他层的信息。3.3 训练策略与超参数选择训练这样的物理代理模型与训练普通的分类模型有所不同我们更关注泛化能力和在极端情况下的物理合理性。损失函数我们使用平滑L1损失Huber损失而不是均方误差MSE。这是因为通量值可能跨越多个数量级例如大气底层热通量远大于顶层MSE容易被大值样本主导。平滑L1损失对大误差的敏感度较低能使训练更稳定同时对小误差也有良好的区分度。优化器使用AdamW优化器它修正了Adam的权重衰减方式能更好地防止过拟合。初始学习率设置为3e-4并配合余弦退火学习率调度器让学习率在训练过程中平滑下降。正则化除了权重衰减我们还使用了较高的Dropout率如0.1在编码器的前馈网络之后以增强模型的泛化能力。此外梯度裁剪是必须的尤其是在训练初期可以防止梯度爆炸保证训练稳定性。数据划分我们将生成的数据集按85:5:10的比例划分为训练集、验证集和测试集。验证集用于早停和超参数调整测试集用于最终评估模型在“从未见过”的大气剖面上的表现。硬件与训练时间我们在一块NVIDIA A10080GBGPU上进行训练。得益于Transformer的并行性整个训练过程在8小时内就完成了。这说明了该方法的可扩展性——生成数据和训练模型的成本远低于用传统方法进行大规模参数研究所需的时间。4. 性能评估与结果分析模型训练完成后我们将其在完全独立的测试集上进行评估并与PICASO的“真实”结果进行对比。4.1 精度评估1%误差的达成图5在原始论文中展示了一个典型的测试案例。黑色实线是PICASO计算出的各层净热通量和散射星光通量红色虚线是我们的Transformer模拟器的预测结果。可以看到两条曲线几乎完全重合。顶部和底部的误差面板显示对于热通量和散射通量预测误差在整个大气柱中都保持在很低的水平。我们进行了全面的统计分析。如图6所示我们计算了在整个测试集上每个压力层级大气层预测误差的平均百分位数分布。黑色实线代表所有层误差的平均值而彩色色带展示了误差分布的不同百分位如25%-75%区间。关键结论是对于热通量和散射星光通量全大气平均的绝对百分比误差中位数都在1%左右。这意味着在绝大多数情况下我们的模拟器能以99%的精度复现高精度物理模型的结果。这个精度水平对于许多天体物理应用已经足够了。例如在光谱反演中观测误差和模型简化如化学平衡假设、一维近似带来的不确定性往往远大于1%。因此用这个快速模拟器替代传统方法不会成为整个分析流程中的主要误差来源。4.2 速度飞跃毫秒级推理与批量加速精度只是故事的一半速度才是革命性的部分。我们在苹果M3集成GPU上测试了推理速度。如图7所示我们测量了处理一个完整的75层大气序列同时计算热和散射通道所需的时间并考察了其与批量大小的关系。单次推理即使批量大小为1即一次只算一个大气剖面推理时间也仅在几毫秒量级。作为对比PICASO计算一个类似剖面可能需要数秒到数十秒。速度提升了1000倍以上。批量加速效应Transformer的另一个优势是批量处理的效率极高。图7曲线显示随着批量大小的增加平均到每个序列的推理时间急剧下降。这是因为模型加载和计算图执行的开销可以被批量内的所有样本分摊。当批量大小达到128或256时平均每个序列的推理时间可以降至亚毫秒级。这对于需要处理成千上万个大气剖面的任务如马尔可夫链蒙特卡洛反演来说意味着将计算时间从“天”缩短到“分钟”。4.3 物理一致性检验一个优秀的代理模型不仅要数值上准确还要在物理上合理。我们进行了几项额外的检查通量守恒在没有内部能源和散射的情况下大气顶部的向上热通量应等于底部的向上热通量能量守恒。我们检查了模型预测的通量剖面是否近似满足这一物理约束结果令人满意。极端外推我们输入了一些训练数据分布之外的、物理上可能不合理的大气剖面如温度逆增过于剧烈。模型虽然误差增大但其预测的趋势仍然是物理的如通量随高度的变化方向正确而不会产生完全荒谬的数值振荡。这说明模型确实学到了一些底层的物理规律而非简单的数据插值。梯度检验我们通过自动微分计算了模型输出如顶层通量对输入参数如某一层的温度的梯度。这个梯度反映了该参数对辐射通量的敏感度。我们将此梯度与通过PICASO进行有限差分法计算的数值梯度进行对比发现两者基本一致。这表明模型不仅学会了函数映射还近似学会了其导数即雅可比矩阵这对于需要梯度信息的优化算法如在大气反演中是一个非常有价值的特性。5. 应用场景与未来拓展方向这个初步成功的Transformer辐射传输模拟器为系外行星天体物理研究打开了一扇新的大门。它的应用场景远不止于快速光谱计算。5.1 集成到一维气候模型中最直接的应用是替代PICASO或类似工具中的辐射传输核心。在一维气候模型中需要反复调用辐射传输代码来计算大气各层的加热/冷却率从而迭代求解大气的温度剖面即“辐射对流平衡”计算。传统上这是计算中最耗时的部分。我们的毫秒级模拟器可以无缝嵌入将整个气候模拟的速度提升数个数量级使得研究人员可以在短时间内探索更广泛的参数空间研究不同恒星照射、行星内部热流、大气成分下的气候状态。5.2 赋能三维全球气候模型GCM这是更具挑战性但也更有前景的方向。在三维GCM中每个水平网格柱在每一个时间步都需要进行辐射传输计算。目前由于计算限制GCM中的辐射方案往往采用高度参数化或简化的“二流近似”等方法牺牲了精度。如果能在GCM的每个网格点嵌入我们这个高精度、快速的辐射传输模拟器将极大地提升GCM预测的物理真实性和可靠性。这将帮助我们更好地理解系外行星上的大气环流、能量输送和极端天气现象。我们未来的工作重点之一就是实现这个集成。5.3 加速大气反演与参数估计从JWST观测到的光谱中反演出行星大气的温度剖面、化学成分和云的性质是一个高维、非线性的优化问题通常需要运行数百万次正向模型。传统方法的计算成本是制约反演深度和精度的主要瓶颈。我们的快速模拟器可以完美地集成到马尔可夫链蒙特卡洛或嵌套采样等反演算法中使得在个人工作站上对高分辨率JWST光谱进行全贝叶斯反演成为可能。这将极大推动对系外行星大气的精确表征。5.4 模型能力的拓展我们目前的模型是在一个相对简化的框架下训练的固定化学、无云。未来的拓展方向非常明确扩充输入参数空间将更多的物理参数作为模型输入例如气体丰度将H2O, CO, CO2, CH4, NH3等关键分子的混合比剖面作为输入特征。云与霾引入云层参数如云顶压力、粒子尺寸、光学厚度等。这需要生成包含云效应的大量训练数据。恒星参数改变恒星的有效温度和光谱类型。行星参数考虑不同重力加速度和行星半径的影响。输出端拓展目前模型输出是层通量。我们可以训练它直接输出高分辨率的光谱作为波长序列或者输出加热率剖面以满足不同下游任务的需求。使用更复杂的训练数据目前数据由PICASO生成。未来可以使用更高精度的辐射传输代码如基于DISORT的模型来生成训练数据让模拟器在逼近“黄金标准”的同时保持高速。处理内存与计算瓶颈虽然当前模型训练很快但若将参数空间极大扩展数据集可能增长到TB级别。我们可以采用数据流式加载、模型并行或混合精度训练等技术来解决。正如原文提到的A100有80GB显存且数据可以流式读取而非全部预加载内存不会成为根本性障碍。实操心得从研究到部署的关键点。在尝试复现或扩展此类工作时有几点经验值得分享第一数据质量决定上限。用于训练的高精度模型数据必须覆盖足够广且物理合理的参数空间否则模型外推能力会很差。第二特征工程很重要。对输入物理量进行标准化如取对数的压力、归一化的温度能加速训练并提升稳定性。第三验证需多维度。不能只看测试集上的平均误差一定要检查模型在物理边界条件如通量守恒和极端情况下的行为。第四考虑不确定性量化。一个理想的代理模型应该能给出预测的不确定性估计如通过贝叶斯神经网络或集成学习这对于后续的反演等科学应用至关重要。6. 常见问题与挑战在实际开发和测试过程中我们遇到并克服了一系列挑战这里总结出来供后来者参考。6.1 训练不稳定性与过拟合问题初期训练时损失函数震荡剧烈甚至出现NaN。或者在训练集上误差很快下降但在验证集上误差早早就开始上升过拟合。排查与解决梯度裁剪这是解决训练初期震荡和NaN问题的首要措施。将梯度范数限制在一个阈值如1.0内。学习率与热身使用学习率热身策略在训练的前几千步从小学习率线性增加到初始学习率有助于模型稳定起步。更严格的归一化确保输入和输出数据都进行了充分的归一化。对于跨越多个数量级的通量值我们尝试了对数变换后再归一化效果更好。增加Dropout和权重衰减对抗过拟合的有效手段。我们最终在前馈网络后使用了0.1的Dropout并设置了适中的权重衰减系数。早停密切监控验证集损失当其连续多个epoch不再下降时停止训练并回滚到最佳模型。6.2 模型无法捕捉物理细节问题模型在整体趋势上正确但在某些特定压力区域如对流层顶、平流层误差明显偏大。排查与解决检查数据分布可能是训练数据在这些物理区域采样不足。我们需要检查生成训练数据时是否充分覆盖了温度反转、不透明度突变等复杂情况。引入物理引导在损失函数中加入简单的物理约束项作为正则化。例如可以加入一个惩罚项鼓励预测的通量剖面满足单调性在某些区域热通量应随高度递减。这属于“物理信息神经网络”的范畴能有效提升模型的物理一致性。调整模型容量可能是模型太简单层数少、隐藏维度过小无法表达复杂的函数关系。可以尝试增加Transformer编码器的层数或隐藏维度。但要注意与过拟合的平衡。6.3 推理速度未达预期问题在部署时发现单次推理速度没有达到论文中报告的毫秒级。排查与解决硬件与后端确保使用了GPU进行推理并且深度学习框架如PyTorch, TensorFlow已正确配置CUDA。苹果M芯片上应使用Metal后端。模型优化使用框架提供的工具如PyTorch的torch.jit.trace或torch.jit.script以及ONNX Runtime对训练好的模型进行图优化和序列化可以去除动态图的开销提升推理速度。批量处理尽可能将多个大气剖面组成一个批次进行推理这是提升吞吐量的最有效方法。如图7所示批量处理能极大分摊固定开销。6.4 外推性能差问题模型在训练数据分布范围内表现良好但一旦输入参数稍微超出范围如更高的温度、更极端的成分预测结果就完全失真。排查与解决这是代理模型的固有局限。机器学习模型本质上是复杂的插值器其外推能力有限。解决方案是尽可能扩大训练数据的覆盖范围。在生成数据时要有意地将参数空间采样扩展到比科学兴趣区域更广的范围为模型提供一个“缓冲区”。不确定性估计开发能够输出预测不确定性的模型如使用Dropout作为近似贝叶斯推断或训练模型集成。当输入位于数据分布之外时模型应给出很大的不确定性这样用户就知道预测结果不可信。混合建模在极端参数区域可以设置一个“安全开关”当模型不确定性超过阈值时自动回退到运行缓慢但可靠的传统物理模型。
