速度提升252倍，斯坦福/UCLA等用LSTM将二阶非线性光学仿真带入毫秒级时代

二阶非线性光学是非线性光学里最重要、应用最广泛的核心分支，主要研究强激光与无中心反演对称的特殊光学晶体相互作用时，由二阶非线性极化率 χ⁽²⁾ 主导的光学效应。简单来说，当高强度激光射入这类晶体，光子会发生「能量合并和频率重组」，直接产生全新频率、全新颜色的光束，实现倍频（Second-harmonic generation，SHG）、和频（Sum-frequency generation，SFG）、差频（Difference-frequency generation，DFG）等经典非线性转换。在现代光学研究中，二阶非线性光学是量子信息、集成光子芯片、生物医学成像、高功率激光系统等领域的关键物理基础。

当前，针对二阶非线性光学过程的理论研究已十分成熟，但在实际工程和实验落地中仍存在明显瓶颈。一方面是算力成本高昂：传统基于傅里叶的仿真算法必须精准求解超快光场振荡变化，还需采用分步傅里叶法（Split-step Fourier Method，SSFM）模拟光在介质中逐段传输过程，这使得数值仿真模拟需要海量计算支撑；而在高重复频率激光实验等实时场景下，这一问题会随着自适应调控、实时参数优化需求提升而更加凸显。

另一方面是模型与实际脱节，传统物理模型偏理想化，难以兼容实验误差、 环境噪声、设备系统漂移等真实不稳定因素，缺乏与实际情况相结合的灵活性，导致实验结果与实际应用脱节。此外，随着光学数字孪生与多级非线性链式联合仿真的日益发展，也让兼具实时仿真与实验联动能力的新工具成为行业刚需。

针对于此，来自斯坦福大学、加利福尼亚大学洛杉矶分校和 SLAC 国家加速器实验室的团队受先前将循环神经网络（RNNs）应用于光纤脉冲传播研究的启发，提出一项基于长短期记忆网络（Long short-term memory，LSTM）的代理模型（surrogate model），既能快速且精准地预测 SFG 的输出光场，又能大幅削减运算成本。该模型依托于 SLAC 实验室 LCLS-Ⅱ（Linac Coherent Light Source Ⅱ）光电阴极驱动激光全链路模型生成的 SSFM 仿真数据集进行训练，与传统模型相比运算速度提升了 252 倍，为激光系统实时优化、数据融合建模、数字孪生搭建了重要的技术基础。

相关成果以「Deep learning-assisted modeling for χ⁽²⁾ nonlinear optics」为题，收录于 Advanced Photonics。

研究亮点：

* 将深度学习模型拓展至全耦合多场二阶非线性光学极化率动力学体系，全方位提升了非线性光学建模的运算速率、灵活度和适用范围 

* 在单张 NVIDIA A100 GPU 上以 200 批次规模运行，单样本推理耗时降至 7.43 毫秒，相较 SSFM 模型实现 252 倍提升

* 本研究打通了数值仿真与实际实验应用之间的壁垒，为高效、可规模化、智能化的光子系统设计开辟了新的研究思路

论文地址：
https://go.hyper.ai/5bLoA

构建 χ⁽²⁾ 非线性光学高保真仿真数据集

本研究的数据集基于 SLAC 实验室 LCLS-Ⅱ 直线加速器相干光源的光阴极驱动激光系统全程仿真模型构建，该系统包含 1,035 纳米锁模红外激光光源、可编程光谱整形器、啁啾脉冲放大器以及非线性频率转换模块（如下图 a 所示）。为覆盖丰富的脉冲形态，研究通过随机采样二阶色散、三阶色散和光谱振幅整形参数，共生成 10,000 种脉冲整形配置方案，其中仅采用相位整形的配置不少于 400 组。之后每种配置均采用 SSFM 开展高精度仿真，最终得到非线性晶体中 100 个传输切片上的三个耦合光场数据，分别是 SHG1、SFG 和 SHG2，单场采样点数达 32,768。

数据生成与预处理

在数据预处理阶段，研究采用三阶段流程（如上图 b 所示：）：第一阶段在频域内对光场进行截断和降采样，将 SFG 光场精缩至 348 个复值，两组 SHG 光场各自精缩至 1,892 个复值；第二阶段将各光场实部与虚部相拼接，形成固定长度为 8,264 个元素的实值向量；第三阶段基于全局数据集的极值将向量内所有元素归一化至区间 [0,1]。最终数据集被划分为 890,000 训练样本、10,000 验证样本和 90,000 测试样本。

搭建基于 LSTM 的序列代理模型架构

LSTM 模型采用序列到序列架构，将非线性晶体的每个离散切片视为一个时间步。这套网络包含 2,048 个隐藏单元，后续衔接三层全连接层，维度依次为（2048，4096）、（4096，4096）以及（4096，8264），三层全连接层分别采用 ReLU、Tanh 和 Sigmoid activations，如下图所示。

网络架构与流程示意图

LSTM 模型使用 Adam 优化器和加权均方误差（wMSE）损失函数进行优化，如下图所示：

wMSE 损失函数

在训练过程中，LSTM 模型以 10 个空间切片组成的序列作为输入，来预测下一个切片。针对包含 100 个切片的每组模拟数据，通过在序列上滑动窗口可生成 100 组输入输出样本，前 9 组会在输入序列前端重复拼接初始切片，以此保证输入长度统一。最终输入的训练张量形状为（批量大小，10，8264），输出张量形状为（批量大小，8264）。

LSTM 模型进行推理时采用自回归运行方式。首先，将初始切片重复 10 次作为首个输入，随后 LSTM 模型预测下一个切片，将其接入输入序列末尾，同时舍弃序列中最早的切片，始终维持输入窗口的长度为 10。该流程循环执行完成全部 100 个空间步长的推演，借助序列模型的预测结果，从而完整还原非线性晶体内部整场物理量的演化过程。

精度与效率双提升，较基线模拟提速 252 倍

为评估 LSTM 模型推理性能，研究采用一种兼顾波形形态与总能量敏感度的综合误差指标。该综合指标包含三个无量纲组成部分：经面积归一化波形的余弦相似度，经过取反与缩放处理，使得波形完全一致时误差为零；依据总积分强度计算得出、与能量成正比的归一化均方误差（NMSE），用于对能量不匹配情况进行惩罚；两组强度分布曲线之间的 Wasserstein distance（Earth Mover’s Distance，EMD），可精准感知强度的局部重新分布情况。

波形重建精度

实验配置方面，LSTM 在单张 NVIDIA A10G GPU 上训练约 180 轮后完成收敛，训练时长约 160 小时，最终训练损失与验证损失分别达到 2.05 ✖️ 10⁻⁵ 和 2.03 ✖️10⁻⁵。评估方式为计算预测结果与 SSFM 仿真结果间的综合误差指标。

实验结果方面，为直观呈现定性效果，研究分别刻画了 SFG、SHG1 和 SHG2 的时间强度综合误差指标的直方图和统计量，并从 SFG 误差分布的四个四分位区间内随机选取两组测试数据集样本，同时呈现 SFG 及对应 SHG1 在时域、频域下的预测强度曲线与真实强度曲线。如下图所示：

评估误差分布

下面两组图片展示了误差分布最高四分位组中的另外两组实例，以凸显 LSTM 模型在不同整形条件下的表现。结果显示，二者综合误差指标数值分别为 0.012 和 0.003。在两种工况下，LSTM 模型均能在频域或时域中精准重构 SFG 和 SHG1，仅在光谱调制幅度较大时，SHG1 会出现局部偏差。这些实例证实了 LSTM 模型能够在各类光谱整形工况下实现良好的泛化效果。

对应以纯相位整形为主、光谱幅度调制程度极低的工况

存在明显光谱凹陷的工况

模型计算效率

下图展示了 LSTM 模型在不同硬件下的推理耗时：基线模拟是基于 SSFM 模型得出的，其中 SSFM 求解非线性传输时主要耗时项，具体来说单核 CPU 下，完整仿真总耗时 1.98 秒，而 SSFM 运算就占到约 1.875 秒。LSTM 模型在单核 CPU 上的时长与基线模拟时间相当，这是由于批处理操作产生了多余的开销；但在单张 NVIDIA A100 GPU 上，设置批次大小为 200，单样本推理耗时仅需 7.43 毫秒，较基线模拟提速 252 倍。

效率统计数据

结语

本研究既有受到过往经验的启发，又融入多项创新的想法，比如三合一复合评价体系，结合反向余弦相似度、NMSE 和 EMD，改变了传统单一指标无法适配各类脉冲形态与能量尺度的视觉和实际效果评判，使得实验结果更加可靠、可信。

当然，LSTM 最为重要的核心优势还是在于彻底改变了传统 SSFM 模型需要反复做时频域转换而造成算力成本高昂的难题，只需在简约频域直接学习映射关系，省去频繁的分步傅里叶法变换，实现实时推演。而 LSTM 模型除针对 SFG 过程外，还能通用至各类二阶非线性光学场景，具备广泛的应用潜力，适配当前实施实验+实际应用的需求。

最后，LSTM 模型的成功验证，更加揭示机器学习模型在非线性光学仿真领域所具备的高效赋能能力，打通纯数值仿真与实体实验的壁垒，为高效、规模化、智能化光电系统设计提供了全新技术范式。