End-to-End Hybrid Refractive-Diffractive Lens Design with Differentiable Ray-Wave Model

一句话总结

提出一种可微的”光线追踪 + 波动传播”（ray-wave）混合成像模型，把折射透镜的像差与衍射光学元件（DOE）的相位调制在同一框架内精确、可微地模拟，从而支持折射-衍射混合透镜与图像重建网络的端到端联合优化。

研究背景

混合折射-衍射透镜同时具备折射透镜的高光效与 DOE 的强信息编码能力，被视为下一代成像系统的重要方向。但要对这类系统做端到端设计，必须有一个既精确又可微的成像模型，而已有方法都存在硬伤：

• 傍轴波动光学模型把折射透镜理想化为一块薄相位板，完全忽略光学像差，且无法优化折射透镜，只能用于小视场、低像差场景。
• 光线追踪模型（如 Zemax 及部分近期工作）把衍射面近似为局部光栅，无法真实刻画衍射现象，还要求相位图缓慢变化；而端到端设计中的 DOE 往往具有快速变化甚至不连续的相位分布，这类模型直接失效。

因此，此前没有工作能够在同时存在复杂/不连续衍射面的情况下，精确地模拟并优化折射透镜。本文针对 DOE 位于最后一个折射面与传感器之间（远离常用的傅里叶平面）这一配置，填补该空白。

方法

整体框架分三步串成一条全可微的成像流水线：先用相干光线追踪算出 DOE 平面处的复振幅波场（承载折射透镜引入的幅度与相位像差），再对该波场施加 DOE 相位调制，最后做自由空间传播到传感器平面计算点扩散函数（PSF）。PSF 与输入图像卷积得到仿真原始图像，喂入重建网络，梯度可一路回传至光学参数。

flowchart LR
    A[点光源/输入图像] --> B[相干光线追踪<br/>折射透镜像差]
    B --> C[DOE 平面复波场 U_DOE-]
    C --> D[DOE 相位调制 U_DOE+]
    D --> E[角谱法自由空间传播]
    E --> F[传感器平面 PSF]
    F --> G[与图像卷积得原始图]
    G --> H[NAFNet 图像重建]
    H -. 梯度回传 .-> B

关键设计一：相干光线追踪构建像差波场。沿用蒙特卡洛光线追踪，逐面记录光线的位置、方向、相位与波长，相位由光程差得到，交点用牛顿法、折射用斯涅尔定律求解。到达 DOE 面后用相干叠加合成复波场：

\[U_{DOE-}=\sum_{i=1}^{spp} u_i \exp(j\phi_i)\cos\langle d_i, n\rangle\]

其中 \(cos\langle d_i, n\rangle\) 是 Rayleigh-Sommerfeld 定理中的倾斜因子，\(spp\) 为每个点源采样光线数（实验取 \(10^6\)）。每条光线的复振幅通过反向双线性插值分配到相邻 4 个像素，既保证可微又缓解亚像素相移带来的混叠。由于波长远小于透镜物理尺度，相位计算采用双精度。

关键设计二：DOE 相位调制。基于标量衍射理论与 Kirchhoff 边界条件，DOE 对波场引入相位改变：

\[U_{DOE+}=U_{DOE-}\exp\!\left(j\,\frac{2\pi}{\lambda}(n_\lambda-1)h(x,y)\right)\]

为与 Zemax 兼容，内部用设计波长 \(\lambda_0\) 下的相位 \(\phi_0\) 表示，并换算到其他波长：

\[U_{DOE+}=U_{DOE-}\exp\!\left(j\,\frac{n_\lambda-1}{n_0-1}\,\frac{\lambda_0}{\lambda}\,\phi_0\right)\]

这里 \(\phi_0\) 是可优化参数（\(\lambda_0=0.55\,\mu m\)），且允许 \(\phi_0\) 为不连续函数以刻画多级 DOE，从而摆脱已有方法对相位连续性的限制。

关键设计三：PSF 计算与混合精度训练。调制后的波场用角谱法传播到传感器：

\[U_{Sensor}=\mathcal{F}^{-1}\big(\mathcal{F}(U_{DOE+})\,H\big),\qquad PSF=|U_{Sensor}|^2\]

为满足奈奎斯特采样，波场以 DOE 特征尺寸的整数倍（通常 2 倍）采样并做零填充覆盖离轴波。同时提出混合精度策略：光学仿真中的相位计算用双精度，PSF 转单精度供网络训练，反传时再转回双精度继续在光学部分回传，兼顾精度与效率。模型基于开源可微光线追踪器 DeepLens 与 PyTorch 实现。

实验结果

在端到端混合透镜设计仿真中（折射透镜 + DOE + NAFNet 联合优化），与傍轴波动光学和 Zemax 光线追踪两种设计对比。Zemax 因直接以 RMS 光斑尺寸为目标而在该指标上最优，但无法做端到端设计；本文模型直接优化最终成像质量，在图像 PSNR/SSIM/1-LPIPS 上全面领先：

设计方法	RMS 半径(轴上/离轴/平均) ↓	PSNR/SSIM/1-LPIPS(原始) ↑	PSNR/SSIM/1-LPIPS(重建) ↑
傍轴波动光学	16.1/55.0/32.0	17.7/0.520/0.443	25.041/0.688/0.569
光线追踪(Zemax)	7.6/19.6/10.1	24.8/0.728/0.706	36.462/0.963/0.924
Ray-wave 模型(本文)	10.7/28.4/15.4	27.1/0.795/0.713	39.9/0.982/0.963

模型精度上，用”理想薄透镜 + DOE”的可解析系统作为地面真值，本文模型在方孔、连续 Kinoform、不连续多级三类相位图下都与波动光学真值几乎一致，而光线追踪模型忽略边缘衍射、无法处理不连续相位。真实原型方面，作者用一枚焦距 7.5 mm 的现成非球面镜配合自制 16 级 DOE（特征尺寸 \(1\,\mu m\)、物理尺寸 3 mm）搭建 35° 对角视场的紧凑混合透镜：像差校正任务中混合透镜重建 PSNR 从 28.8 提升到 32.5；大视场 EDoF 任务（20 cm 至 10 m）中，相比傍轴模型设计的 DOE，本文 DOE 在离轴区域仍能保持清晰成像。

亮点与局限

亮点：

• 首次给出同时精确模拟折射像差与衍射相位调制、且全程可微的混合透镜成像模型，支持折射透镜、DOE、重建网络三者联合优化。
• 支持不连续/多级 DOE 相位图，突破了以往方法对相位连续性的强假设。
• 有真实拍摄验证，把 DOE 插入现成折射镜后不增加体积，对手机相机等尺寸受限系统有实用价值。

局限：

• 显存开销巨大，单视场 RGB PSF 计算与反传约需 35 GB GPU 内存，需依赖多 GPU、梯度检查点、分块反传等策略缓解。
• 模型假设 DOE 位于最后折射面与传感器之间这一特定配置，且为便于对比忽略了 Fresnel 透射的能量损失。
• 原型受制于折射镜加工能力，只优化 DOE 而未真正端到端优化折射透镜。

延伸思考

• 该 ray-wave 框架把”光线追踪管像差、波动传播管衍射”的分工做成可微流水线，思路可推广到超表面、多层衍射面乃至偏振/矢量衍射等更复杂的元件建模。
• 混合精度处理相位数值精度问题的做法，对其他涉及大动态范围物理量的可微仿真（如全息、超分辨光场）有借鉴意义。
• 显存瓶颈是端到端光学设计规模化的核心障碍，如何在保持波动精度的同时降低采样分辨率或引入自适应采样，是值得深入的工程与算法方向。