自动微分在光学逆向设计中的应用:从光束偏转到颜色路由

自动微分在光学逆向设计中的应用:从光束偏转到颜色路由 1. 项目概述当自动微分遇上光学设计在光学计算设计OCD和超表面Metasurface领域工程师们一直面临着一个核心挑战如何从期望的光学功能比如将特定波长的光偏转60度或者将白光精准地分离成红绿蓝三色出发逆向推导出实现该功能的最优物理结构这就像给你一个谜底让你反推出谜面而且这个“谜面”是纳米尺度下、由硅和空气构成的复杂几何图案。传统方法比如参数扫描或者基于有限差分法的梯度计算要么计算量巨大到不切实际要么在复杂、高维的设计空间中步履维艰精度也难以保证。近年来一个源自深度学习领域的技术——自动微分Automatic Differentiation, AD正在彻底改变这个游戏规则。它并非简单的数值差分或符号求导而是一种基于链式法则在程序执行过程中同时计算函数值及其导数的精确方法。其核心价值在于对于任意一个由代码定义的光学仿真模型比如基于严格耦合波分析RCWA的衍射计算AD能够自动、精确且高效地计算出目标函数如衍射效率相对于成千上万个设计变量如每个纳米像素的折射率的梯度。这相当于为优化算法装上了“高精度导航仪”使其能在复杂的设计空间中沿着梯度方向快速、准确地找到性能最优的结构。本文将以一个具体的工程实践为例深入探讨如何利用自动微分技术完成从一维、二维光束偏转器到一维光栅颜色路由器的逆向设计与优化全过程。我们将不仅展示如何借助现代计算框架如JAX实现这一流程更会拆解每一个关键步骤背后的物理考量与工程权衡分享在实际操作中关于参数选择、优化策略以及结果后处理的经验与技巧。无论你是刚刚接触计算光子学的学生还是寻求提升设计效率的光学工程师这篇来自一线的实践记录都将为你提供一条清晰、可复现的技术路径。2. 核心原理与工具链拆解2.1 自动微分AD为何是光学逆向设计的“神兵利器”要理解AD的优势首先要看看它替代了什么。在光学逆向设计中我们需要优化一个目标函数FoMFigure of Merit例如特定衍射级的效率。假设我们的设计变量是结构内N个单元的折射率分布那么目标函数FoM就是这N个变量的函数FoM f(n1, n2, ..., nN)。为了用梯度下降法优化我们需要计算梯度向量 ∇f [∂f/∂n1, ∂f/∂n2, ..., ∂f/∂nN]。传统方法主要有两种一是有限差分法通过给每个变量施加微小扰动来近似梯度计算一次梯度需要运行N1次仿真当N很大时动辄数万成本无法承受二是手动推导解析梯度公式这对于复杂的电磁仿真模型如RCWA来说极其繁琐且容易出错几乎不可行。自动微分则另辟蹊径。它将整个仿真程序从前向的场计算到后向的衍射效率提取视为一系列基本运算加、减、乘、除、三角函数等的组合。AD系统如JAX、PyTorch会记录下前向计算过程中的所有中间变量和运算关系形成一个计算图。当需要求梯度时它从输出FoM开始沿着计算图反向传播利用链式法则自动、精确地计算出每个输入变量折射率的梯度。这个过程只需要一次前向仿真和一次反向传播其计算复杂度与一次前向仿真同量级与设计变量N的数量几乎无关。这意味着无论你的结构多复杂、变量有多少计算梯度的代价都是可接受的。注意这里说的“精确”是指数学上的精确没有截断误差其精度与函数值计算本身的双精度浮点误差同量级远高于有限差分法。这是AD在优化稳定性上的一大优势。2.2 技术栈选择JAX与RCWA的强强联合在本次实践中我们选择了JAX作为自动微分的核心引擎。JAX不仅提供了强大的、可组合的函数变换grad,jit,vmap其基于NumPy的API也使得科学计算代码非常直观。更重要的是JAX的jit即时编译功能可以将我们的仿真和梯度计算代码编译成高效的XLA加速线性代数指令在CPU或GPU/TPU上获得巨大的速度提升这对于需要成千上万次迭代的优化过程至关重要。光学仿真部分我们采用了严格耦合波分析RCWA。RCWA是分析周期性衍射结构如光栅、超表面的经典频域方法。它将结构在傅里叶空间展开通过求解麦克斯韦方程组得到衍射效率。选择RCWA是因为它对于本文涉及的1D/2D周期性结构计算速度快、精度高且其数学形式矩阵运算非常适合用JAX实现和自动微分。我们将RCWA的核心计算步骤构建Toeplitz矩阵、求解本征值问题、计算散射矩阵等都用JAX的数组操作重写。这样整个“输入结构参数 - 运行RCWA仿真 - 输出衍射效率”的流程就变成了一个纯粹的、由JAX可追踪tracable函数组成的计算图从而无缝地支持自动微分。2.3 优化目标与设计空间定义逆向设计的起点是明确“我们要什么”。对于光束偏转器目标很直接最大化特定衍射级通常是1级的效率。FoM可以简单地定义为FoM η_{1}。对于颜色路由器目标则更复杂需要将不同波段如R、G、B的光场能量分别集中到指定的子像素区域。因此其FoM是一个复合函数例如文中公式(74)它计算了在每个目标波段内目标子像素区域的平均电场强度与总透射能量的加权比值再对不同波长取平均。定义清晰、可微的FoM是成功优化的第一步。设计空间决定了我们“能改变什么”。在超表面设计中通常将器件的单元结构离散化为一个个“像素”或称体素。每个像素可以被一种材料如硅或另一种材料如空气填充形成一个二元结构。在优化初期我们允许每个像素的“材料填充率”或等效折射率是一个介于空气和硅之间的连续值即灰度值这被称为“灰度优化”grayscale optimization。这样做的原因是连续的设计空间更平滑梯度信息更丰富有利于优化算法找到更好的性能极值点。在优化收敛后我们再通过一个“二值化”binary push过程将这些灰度值“推”向0空气或1硅得到最终可制造的二元结构。3. 一维光束偏转器逆向设计实战3.1 问题建模与初始化我们首先从相对简单的一维光束偏转器开始。目标设计一个工作在波长λ₀900nm下的一维光栅将垂直入射的TM偏振光以50度角偏转。器件单元周期由光栅方程决定P λ₀ / sin(θ) ≈ 1175 nm。我们将一个周期内的结构在高度方向假设厚度固定上离散化为256个单元格cells。每个单元格的折射率nᵢ就是我们的设计变量初始值在硅的折射率~3.5和空气的折射率1.0之间随机均匀分布。这种随机初始化有助于探索设计空间的不同区域避免陷入平庸的局部最优解。前向仿真模型基于RCWA。我们需要设置足够的傅里叶谐波数量Fourier orders来保证计算精度。经过收敛性测试后文会详述对于此结构我们选择FTO100即保留100个谐波。仿真输出包括各个衍射级的效率我们关注的是1级效率η_{1}将其作为FoM。3.2 基于ADAM优化器的迭代更新有了可微的仿真模型和FoM优化过程就转化为一个标准的梯度下降问题。我们采用深度学习领域广泛使用的ADAM优化器。ADAM结合了动量Momentum和自适应学习率RMSProp的优点在非凸优化问题上通常表现稳定、收敛快。关键参数设置学习率 (Learning Rate)设置为0.5。这是一个相对较大的值但在光学逆向设计初期梯度往往很大较大的学习率有助于快速下降。我们通常会配合学习率衰减策略使用。迭代次数 (Epochs)进行足够多的迭代直到FoM收敛。实践中我们可能会设置一个较大的值如500并配合早停Early Stopping策略。每次迭代的步骤清晰而自动化前向传播将当前折射率分布n_array输入RCWA模型计算得到FoMη_{1}。自动微分求梯度调用jax.grad函数自动计算FoM对n_array中每一个元素的梯度grad_array。这个过程在后台完成了复杂的链式法则求导。参数更新将n_array和grad_array输入ADAM优化器得到更新后的n_array_new。循环用n_array_new替换n_array重复步骤1-3。3.3 灰度优化与二值化处理在优化过程中折射率值可以自由变化因此我们得到的是一个“灰度”结构即每个像素的折射率是介于硅和空气之间的某个值。这在物理上可能对应一种混合材料或亚波长结构但通常为了简化制造如采用标准的刻蚀工艺我们需要一个纯硅和纯空气的二元结构。因此在优化收敛后我们需要一个“二值化”步骤。这不是简单地将灰度值四舍五入到0或1因为那样做通常会严重破坏性能。文中提到的“binary push”是一种常用技巧在优化的最后若干轮或者在得到灰度优化结果后引入一个额外的“二值化力”。例如可以修改FoM增加一个惩罚项鼓励折射率值远离中间值、靠近0或1也可以在每次更新后对折射率施加一个Sigmoid类的函数进行挤压。最终我们得到一个清晰的、由硅柱和空气间隙组成的二元光栅结构。3.4 结果分析与实操要点经过优化我们得到了性能优异的设计。图27a展示了优化曲线其中蓝线是灰度优化过程中的效率橙线是经过二值化处理后的最终效率。可以看到二值化过程会带来一定的性能损失但通过合理的“binary push”策略可以将损失降到最低。最终我们找到了一个效率高达89.4%的设计。实操心得随机初始化的必要性由于设计空间高度非凸存在大量局部最优解。仅从一种初始状态如均匀分布、全硅、全空气开始优化很可能收敛到一个次优解。因此必须进行多次随机初始化的优化。文中提到实验重复了100次。在实际工程中根据计算资源至少进行10-20次独立优化然后从中挑选性能最好的设计。图27a中的阴影区域±标准差正说明了不同初始点导致的性能分布这有助于我们评估设计的鲁棒性和优化算法的有效性。图27b展示了最终结构内部的电场分布。我们可以清晰地看到入射波如何与光栅结构相互作用能量如何被有效地耦合到1级衍射方向。这种可视化不仅是结果验证也是理解器件物理工作机制的宝贵工具。4. 二维光束偏转器与设计复杂性管理4.1 从1D到2D问题扩展与网格划分将问题从一维扩展到二维设计自由度呈平方级增长同时也带来了新的挑战和考量。我们以设计一个将λ1000nm的光以60度角偏转的二维超光栅metagrating为例。器件单元在x和y方向具有不同的周期P_x λ / sinθ ≈ 1154 nm, P_y λ/2 500 nm。这个非正方形的周期意味着器件在x方向偏转方向需要更长的相位调制距离而在y方向则保持紧凑。我们将这个矩形单元离散化为256 (x方向) × 128 (y方向) 32,768个单元格。这意味着我们有超过3万个设计变量是之前1D案例256个变量的128倍。如此高维的优化对梯度计算的效率和优化算法的稳定性提出了更高要求也凸显了自动微分仅需一次前向仿真即可获得全部梯度的巨大优势。4.2 傅里叶谐波收敛性测试精度与效率的权衡RCWA的精度和计算成本直接取决于所使用的傅里叶谐波数量Nx, Ny。谐波数越多仿真越精确但计算时间尤其是矩阵运算的规模也急剧增加。因此在开始正式优化前必须进行收敛性测试。如图28b所示我们对初始随机结构进行测试计算不同谐波数下的衍射效率。可以看到当Nx和Ny增加到一定程度后效率的变化趋于平缓。我们需要选择一个“拐点”在保证足够精度的前提下最小化计算量。文中最终选择了Nx13, Ny10。这个选择基于经验通常需要保证仿真精度在0.1%量级以内同时使得单次仿真时间在可接受范围内例如在GPU上小于1秒以满足成千上万次迭代的需求。注意事项收敛性测试的陷阱收敛性测试应在与目标结构“相似”的初始结构上进行。如果用一个全硅或全空气的简单结构测试收敛所需的谐波数可能远低于一个复杂的、高对比度的优化后结构。最稳妥的方法是用一个中等复杂度的随机结构或者用优化中间过程的某个结构进行测试。此外对于宽带优化如颜色路由器需要在所有关注波长上进行测试取最苛刻的情况作为标准。4.3 优化流程与性能表现二维优化的流程与一维类似随机初始化、ADAM优化、二值化处理。经过110个epoch的优化最终结构实现了92%的偏转效率。图28c的学习曲线显示优化过程平稳收敛。图28d的电场分布图在XZ平面上清晰地展示了光束被成功偏转60度的效果。这个案例的成功表明基于自动微分的优化框架能够有效地处理高维数万变量、非凸的设计问题并且能够找到物理上合理、性能接近理论极限的复杂二维结构。5. 一维光栅颜色路由器处理复杂目标函数5.1 复杂FoM的定义与物理意义颜色路由器的设计目标比简单的偏转器复杂得多。它需要将宽谱入射光例如白光中不同波段的成分红、绿、蓝分别引导到空间上不同的位置对应的子像素。因此其FoM不能是单一波长、单一衍射级的效率。文中定义的FoM公式74是一个精心构造的复合函数空间选择性对于每个目标波长λ计算目标子像素区域x1到x2内的平均电场强度积分。这衡量了光能量是否被“路由”到了正确的位置。归一化将该积分除以整个周期内的总电场强度积分。这是为了消除绝对透射率的影响专注于能量的相对分布。透射率加权乘以该波长下的总透射率T(λ)。这确保了在将能量路由到正确位置的同时器件本身也有高的透过率避免单纯追求分布而牺牲了整体光通量。宽带平均最后对每个颜色通道R、G、B在其对应的波长范围内如R:600-700nm选取多个采样点文中是9个计算FoM的平均值。这迫使优化出的器件在目标波段内具有平坦、稳定的性能而不是只在某个孤立波长点工作。这个FoM的定义充分体现了逆向设计的精髓将高阶的、功能性的系统级指标“把红光分到左边绿光分到中间”转化为一个可微的、基于底层物理场电场的数学表达式。5.2 多层堆叠结构与优化策略该颜色路由器采用了一个垂直堆叠的8层1D光栅结构总高2μm每层有64个横向单元格。这种多层结构提供了更丰富的相位和振幅调控自由度能够实现更复杂的光谱和空间响应。设计变量总数是8层 × 64单元 512个虽然数量上少于2D案例但由于FoM计算涉及多个波长和空间积分每次前向仿真的计算量更大。优化策略与之前保持一致。图30a展示了优化曲线其中包含了灰度优化和二值化后的结果对比。可以看到对于如此复杂的FoM优化过程依然能够有效收敛。图30b展示了最终器件的光谱性能在R、G、B三个波段分别有大部分能量被路由到了正确的子像素区域实现了颜色分离的功能。图30c的电场分布图直观地展示了不同颜色的光在器件内部被导向不同方向的过程。5.3 处理复杂FoM的梯度计算技巧当FoM变得复杂时自动微分依然可以无缝工作这是其强大之处。但工程师需要注意以下几点计算图复杂度复杂的FoM意味着更长的计算图。虽然JAX等框架可以处理但可能会占用更多内存用于存储前向传播的中间变量。合理使用jit编译可以优化计算和内存访问。梯度噪声当FoM涉及对空间或波长的积分、平均时数值离散化可能会引入微小的“噪声”。虽然AD梯度在数学上是精确的针对离散化的模型但优化算法如ADAM对噪声有一定的鲁棒性。如果发现优化不稳定可以检查积分网格是否足够细密或者尝试在FoM中加入微小的正则化项如对折射率变化的平滑性约束来稳定梯度。多目标权衡颜色路由器的FoM本质上是R、G、B三个通道性能的加权平均。权重的选择会影响最终设计的侧重点。在实践中可能需要根据具体应用需求如对哪个颜色通道的灵敏度要求更高来调整权重并进行多次优化来探索帕累托前沿Pareto Frontier。6. 工程实践中的常见问题与排查指南在实际操作基于自动微分的光学逆向设计流程时会遇到各种各样的问题。以下是一些典型问题及其排查思路来源于多次“踩坑”后的经验总结。6.1 优化不收敛或收敛至不良解这是最常见的问题。现象可能是FoM震荡、不上升或者收敛到一个很低的性能值。排查学习率学习率过大导致震荡过小导致收敛慢甚至停滞。建议策略从一个中等值如0.1开始观察初期迭代。如果FoM剧烈震荡立即降低学习率除以10。如果FoM几乎不变可适当增大。更高级的做法是使用学习率调度器如指数衰减或余弦退火。检查梯度量级在迭代开始时打印出梯度向量的范数jnp.linalg.norm(grad)。如果梯度范数异常小如1e-10可能是FoM定义有问题或者仿真模型存在数值问题如矩阵奇异导致梯度无法有效回传。如果梯度范数异常大可能需要降低学习率或对梯度进行裁剪gradient clipping。验证前向仿真在优化循环外手动设置几个简单的测试结构如均匀介质、已知解的光栅运行前向仿真检查输出的衍射效率是否物理合理。确保RCWA等仿真器的核心计算正确无误。审视FoM的可微性确保你定义的FoM在整个设计空间内是连续可微的。例如如果你直接取某个探测点的电场强度绝对值在电场为零点附近可能不可微。通常使用电场强度的平方与能量成正比作为基础量。局部最优解陷阱如前所述必须进行多次随机初始化。如果只有少数几次优化找到了好解而多数结果很差说明设计空间非常崎岖。可以考虑引入模拟退火、增加随机扰动等策略来帮助跳出局部最优。6.2 二值化导致性能严重下降灰度优化结果很好但一进行二值化效率就暴跌。二值化过程太激进如果简单地进行硬阈值处理0.5设为硅否则为空气性能损失是必然的。应采用渐进式二值化在优化的后期逐渐引入一个偏向二值化的惩罚项到FoM中或者每迭代若干步后对折射率施加一个逐渐变陡的Sigmoid函数让优化器有机会在二值化的约束下重新调整结构。制造约束考虑不足最终的二值结构可能包含一些非常细小的特征如孤立的单个像素点或极窄的线条这些在制造中可能无法实现或极其脆弱。可以在FoM中加入制造约束例如通过卷积操作引入最小特征尺寸约束或者在优化后期进行形态学开闭运算来平滑结构。灰度结构本身“脆弱”有时灰度优化找到的解强烈依赖于中间折射率值这种解物理上可能对应一种有效的梯度折射率分布但无法用二元材料近似。这提示我们或许在优化初期就应该施加更严格的材料约束或者探索三材料如硅、二氧化硅、空气的设计空间。6.3 仿真结果与预期或文献不符优化出的结构看起来奇怪或者性能远低于理论预期或已发表结果。收敛性测试不足这是首要怀疑对象。用你优化出的最终结构重新进行傅里叶谐波数的收敛性测试。很可能当前设置的谐波数对于这个复杂结构来说不够导致仿真结果虚高或虚低。用不足的谐波数进行优化等于在错误的“地图”上导航。物理参数错误反复检查所有物理参数波长、入射角、材料折射率是否随波长变化、偏振态、周期、厚度等。一个单位的错误如nm写成μm就会导致全盘皆输。边界条件与对称性检查RCWA仿真中设置的边界条件是否正确。对于正常入射是否利用了对称性来减少计算量如果用了要确保你的FoM定义和梯度计算与对称性假设兼容。与已知案例对比找一个结构简单、有解析解或公认结果的案例如亚波长光栅的等效介质理论用你的仿真代码复现进行交叉验证。6.4 计算速度过慢或内存溢出当设计变量很多如2D案例或FoM很复杂如宽带颜色路由器时计算可能成为瓶颈。充分利用jit编译确保你的前向仿真函数和梯度计算函数都被jit装饰器编译。第一次调用会有编译开销但后续调用速度极快。将循环体一次迭代整体jit编译通常比只编译内部函数更有效。批处理与vmap颜色路由器需要对多个波长计算FoM然后平均。使用JAX的vmap函数可以自动将单个波长的计算向量化在GPU/TPU上并行执行极大提升速度。梯度计算模式JAX的grad默认计算标量函数对输入的第一个参数的梯度。如果你需要计算梯度矩阵如多个输出对多个输入的梯度了解并使用jacfwd前向模式雅可比或jacrev反向模式雅可比更常用等更高级的变换。内存管理高分辨率2D或3D仿真可能消耗大量内存存储中间变量。可以考虑使用jit的donate_argnums参数来允许函数复用输入缓冲区或者检查代码中是否无意间创建了不必要的中间数组副本。对于极大的问题可能需要采用分块计算或降低精度如使用float32而非float64的策略。这个基于自动微分的逆向设计流程将光学工程师从繁琐的梯度推导和低效的优化中解放出来让我们能够更专注于物理问题的定义和工程目标的实现。它代表了一种“设计即代码”的新范式其中物理模型、优化目标和制造约束都被编码进一个可微的计算图中通过梯度下降自动寻找最优解。随着计算硬件和算法工具的不断进步这套方法正变得越来越强大和普及成为现代光学与光子学设计的核心技能之一。