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

1. 项目概述：当自动微分遇上光学设计

在光学计算设计（OCD）和超表面（Metasurface）领域，工程师们一直面临着一个核心挑战：如何从期望的光学功能（比如将特定波长的光偏转60度，或者将白光精准地分离成红绿蓝三色）出发，逆向推导出实现该功能的最优物理结构？这就像给你一个谜底，让你反推出谜面，而且这个“谜面”是纳米尺度下、由硅和空气构成的复杂几何图案。传统方法，比如参数扫描或者基于有限差分法的梯度计算，要么计算量巨大到不切实际，要么在复杂、高维的设计空间中步履维艰，精度也难以保证。

近年来，一个源自深度学习领域的技术——自动微分（Automatic Differentiation, AD），正在彻底改变这个游戏规则。它并非简单的数值差分或符号求导，而是一种基于链式法则，在程序执行过程中同时计算函数值及其导数的精确方法。其核心价值在于，对于任意一个由代码定义的光学仿真模型（比如基于严格耦合波分析RCWA的衍射计算），AD能够自动、精确且高效地计算出目标函数（如衍射效率）相对于成千上万个设计变量（如每个纳米像素的折射率）的梯度。这相当于为优化算法装上了“高精度导航仪”，使其能在复杂的设计空间中，沿着梯度方向快速、准确地找到性能最优的结构。

本文将以一个具体的工程实践为例，深入探讨如何利用自动微分技术，完成从一维、二维光束偏转器到一维光栅颜色路由器的逆向设计与优化全过程。我们将不仅展示如何借助现代计算框架（如JAX）实现这一流程，更会拆解每一个关键步骤背后的物理考量与工程权衡，分享在实际操作中关于参数选择、优化策略以及结果后处理的经验与技巧。无论你是刚刚接触计算光子学的学生，还是寻求提升设计效率的光学工程师，这篇来自一线的实践记录都将为你提供一条清晰、可复现的技术路径。

2. 核心原理与工具链拆解

2.1 自动微分（AD）为何是光学逆向设计的“神兵利器”

要理解AD的优势，首先要看看它替代了什么。在光学逆向设计中，我们需要优化一个目标函数FoM（Figure of Merit），例如特定衍射级的效率。假设我们的设计变量是结构内N个单元的折射率分布，那么目标函数FoM就是这N个变量的函数：FoM = f(n1, n2, ..., nN)。为了用梯度下降法优化，我们需要计算梯度向量 ∇f = [∂f/∂n1, ∂f/∂n2, ..., ∂f/∂nN]。

传统方法主要有两种：一是有限差分法，通过给每个变量施加微小扰动来近似梯度，计算一次梯度需要运行N+1次仿真，当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）来保证计算精度。经过收敛性测试（后文会详述），对于此结构，我们选择FTO=100（即保留100个谐波）。仿真输出包括各个衍射级的效率，我们关注的是+1级效率η_{+1}，将其作为FoM。

3.2 基于ADAM优化器的迭代更新

有了可微的仿真模型和FoM，优化过程就转化为一个标准的梯度下降问题。我们采用深度学习领域广泛使用的ADAM优化器。ADAM结合了动量（Momentum）和自适应学习率（RMSProp）的优点，在非凸优化问题上通常表现稳定、收敛快。

关键参数设置：

• 学习率 (Learning Rate)：设置为0.5。这是一个相对较大的值，但在光学逆向设计初期，梯度往往很大，较大的学习率有助于快速下降。我们通常会配合学习率衰减策略使用。
• 迭代次数 (Epochs)：进行足够多的迭代直到FoM收敛。实践中，我们可能会设置一个较大的值（如500），并配合早停（Early Stopping）策略。

每次迭代的步骤清晰而自动化：

1. 前向传播：将当前折射率分布n_array输入RCWA模型，计算得到FoM（η_{+1}）。
2. 自动微分求梯度：调用jax.grad函数，自动计算FoM对n_array中每一个元素的梯度grad_array。这个过程在后台完成了复杂的链式法则求导。
3. 参数更新：将n_array和grad_array输入ADAM优化器，得到更新后的n_array_new。
4. 循环：用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增加到一定程度后，效率的变化趋于平缓。我们需要选择一个“拐点”，在保证足够精度的前提下，最小化计算量。文中最终选择了Nx=13, Ny=10。这个选择基于经验：通常需要保证仿真精度在0.1%量级以内，同时使得单次仿真时间在可接受范围内（例如，在GPU上小于1秒），以满足成千上万次迭代的需求。

注意事项：收敛性测试的陷阱收敛性测试应在与目标结构“相似”的初始结构上进行。如果用一个全硅或全空气的简单结构测试，收敛所需的谐波数可能远低于一个复杂的、高对比度的优化后结构。最稳妥的方法是用一个中等复杂度的随机结构，或者用优化中间过程的某个结构进行测试。此外，对于宽带优化（如颜色路由器），需要在所有关注波长上进行测试，取最苛刻的情况作为标准。

4.3 优化流程与性能表现

二维优化的流程与一维类似：随机初始化、ADAM优化、二值化处理。经过110个epoch的优化，最终结构实现了92%的偏转效率。图28c的学习曲线显示，优化过程平稳收敛。图28d的电场分布图在XZ平面上清晰地展示了光束被成功偏转60度的效果。

这个案例的成功表明，基于自动微分的优化框架能够有效地处理高维（数万变量）、非凸的设计问题，并且能够找到物理上合理、性能接近理论极限的复杂二维结构。

5. 一维光栅颜色路由器：处理复杂目标函数

5.1 复杂FoM的定义与物理意义

颜色路由器的设计目标比简单的偏转器复杂得多。它需要将宽谱入射光（例如白光）中不同波段的成分（红、绿、蓝）分别引导到空间上不同的位置（对应的子像素）。因此，其FoM不能是单一波长、单一衍射级的效率。

文中定义的FoM（公式74）是一个精心构造的复合函数：

1. 空间选择性：对于每个目标波长λ，计算目标子像素区域（x1到x2）内的平均电场强度积分。这衡量了光能量是否被“路由”到了正确的位置。
2. 归一化：将该积分除以整个周期内的总电场强度积分。这是为了消除绝对透射率的影响，专注于能量的相对分布。
3. 透射率加权：乘以该波长下的总透射率T(λ)。这确保了在将能量路由到正确位置的同时，器件本身也有高的透过率，避免单纯追求分布而牺牲了整体光通量。
4. 宽带平均：最后，对每个颜色通道（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）的策略。

这个基于自动微分的逆向设计流程，将光学工程师从繁琐的梯度推导和低效的优化中解放出来，让我们能够更专注于物理问题的定义和工程目标的实现。它代表了一种“设计即代码”的新范式，其中物理模型、优化目标和制造约束都被编码进一个可微的计算图中，通过梯度下降自动寻找最优解。随着计算硬件和算法工具的不断进步，这套方法正变得越来越强大和普及，成为现代光学与光子学设计的核心技能之一。