上个月我训练一个图像分类模型看到GPU集群的实时功耗跑到30多千瓦我突然意识到光计算这个概念被低估了。电子计算机逼近物理极限的日子比大多数人想象中来得更快——晶体管已经小到只剩几十个原子宽电子在导线里跑动产生的热量几乎无解。于是我把目光转向一个更古老的思路用光子而不是电子来做神经网络计算。为了验证这个想法我花了两个周末用Python搭了一套光子神经网络的模拟器把马赫-曾德尔干涉仪网格、相位调制、光学点积全部搬进PyTorch里。这篇文章就是这次实验的全程记录包括原理、代码、训练结果以及我踩过的那些坑。如果你对光计算感兴趣或者你想知道把光学器件当神经网络层训练到底怎么落地这篇文章应该能给你一个完整的答案。就算你是纯软件背景、没碰过光学仿真我也会把从物理器件到数学矩阵的每一步都讲清楚。1. 电子计算的瓶颈在哪里一段被功耗逼出来的反省1.1 晶体管越来越小问题越来越大很多人以为摩尔定律还在继续但实际上我们已经撞上了一堵很难翻过去的物理墙。从7nm到5nm再到3nm晶体管的栅极长度已经接近几十个硅原子的直径量级。在这个尺度下量子隧穿效应变得极其明显——电子会直接穿透极薄的绝缘层产生漏电流。漏电流意味着什么意味着即使晶体管什么都不干也在发热。更麻烦的是时钟频率从十几年前就基本停在4GHz附近了。不是芯片厂商不想做高频而是散热能力根本扛不住。我们这几年看到的性能提升几乎全靠堆核心数和专用加速器GPU、TPU本质上是把更多晶体管挤进同一块硅片但这反过来让功耗密度变得更加恐怖。还有一个经常被忽略的问题数据搬运的能耗远高于计算本身。在传统的冯·诺依曼架构里计算单元和存储单元隔着十万八千里数据在CPU、GPU和内存之间来回搬运。搬运一个数据的能耗比计算这个数据高一个数量级。深度学习偏偏是数据密集型任务权重、激活值、梯度全都要搬运所以整个系统的大部分能量其实都浪费在搬数据上了。1.2 AI算力需求正在指数膨胀我在项目里真正感受到压力的还是训练成本。实验室的GPU集群跑一次大规模模型训练实时功耗能到几十千瓦电费账单常常让人心惊肉跳。而整个行业还在追求更大的模型、更长的上下文、更高的精度算力需求是指数级增长的。推理端的问题更尖锐。手机、智能家居、自动驾驶都要求在几瓦甚至毫瓦级别功耗下实时运行神经网络。传统CMOS工艺在这条路上已经很难看到质的突破。业界提了不少方案比如模型量化、剪枝、蒸馏但这些都是在电子计算框架内做优化并没有改变动不动就发热、动不动就耗电的本质。1.3 光计算的三个天然优势光之所以被重新重视是因为它在物理层面就解决了几个关键问题。第一光子不带电荷静止质量为零。它在波导中传播时几乎不产生焦耳热这意味着光计算可以绕过热量失控这个电子计算的头号敌人。第二光有极其丰富的并行维度。我们可以用波分复用不同波长的光同时传不同数据、模分复用、空间并行在一根光纤里同时跑几十上百路信号。电子要并行只能靠布线光天然就并行。第三也是最关键的光的干涉和衍射本质上是线性代数运算。两束光叠加就是加法通过可调器件衰减就是乘法透镜的傅里叶变换特性可以直接做频域处理。而神经网络最核心的运算恰恰是矩阵乘法和卷积——光计算和神经网络之间有一种天然的性格匹配。用一个不太严谨但很好懂类比电子计算像在一条拥挤的单行道上排队送货每辆车只能走一条道光计算则像整个无线电频谱上同时播放不同频道的节目各频道互不干扰接收端想听哪个就解调哪个。2. 光子神经网络的工作原理从光学器件到数学矩阵2.1 光学点积光神经网络的基本单元神经网络里最频繁出现的运算是点积把一个输入向量和一个权重向量逐元素相乘再求和。光学系统实现这个运算非常直接输入光信号经过可调衰减器实现乘法然后多路光相干合并到同一根波导里实现加法合出来的光场振幅就是点积结果。矩阵乘法本质上就是一组并行执行的光学点积输入向量被分束器复制到多个支路每条支路对应权重矩阵的一行每条支路内部做点积最后在各自输出端得到矩阵乘法的每个元素。在实际的光子神经网络中输入数据往往不是直接作为光强输入而是编码成相位或者幅度。相位编码的好处是对光强波动不敏感抗噪声能力强幅度编码则更接近传统神经网络中的实数值输入。两种方式我在模拟里都试过相位编码在噪声环境下确实更稳。2.2 MZI网格可编程光芯片的通用算子MZI是Mach-Zehnder干涉仪的缩写它是目前可编程光神经网络里最常见的构建单元。单个MZI的结构不复杂两路输入光先经过一个5050分束器两臂各自有一个相位调制器移相器然后光再经过一个5050分束器输出。这里最关键的一点是单个MZI的传输矩阵是2×2的幺正矩阵而把很多MZI按网格状排列、加上交叉连接就可以等效成任意维度的幺正矩阵。再配合可调衰减器你就可以用MZI网格实现任意实数矩阵或复数矩阵的乘法。我写过一个非常简洁的Python函数来计算单个MZI的传输矩阵import numpy as np BS (1 / np.sqrt(2)) * np.array([[1, 1j], [1j, 1]]) def mzi_matrix(phi1, phi2): 单个MZI的2x2传输矩阵phi1和phi2是两臂的相移 phase np.diag([np.exp(1j * phi1), np.exp(1j * phi2)]) return BS phase BSBS是理想5050分束器的传输矩阵phase是两个移相器的对角矩阵两者相乘就得到了这个干涉仪的整体效果。当你把多个MZI像搭积木一样级联成网格时整个网格的传输矩阵就是一系列小矩阵的连乘积def build_mesh_matrix(phis, thetas, nNone): 把一维相位参数构造成MZI网格的等效矩阵。 这里为了演示用Clements分解的简化版本。 ...有了这个矩阵化表达我就可以把整个光学层当做一个可以学习的线性变换来训练。2.3 微环谐振器和衍射层另外两条路线除了MZI网格光子神经网络还有另外两条值得了解的路线。微环谐振器MRR是目前硅光芯片上很常见的器件。微环的半径决定谐振频率通过热光效应或者电光效应改变折射率就可以调制通过环的光强和相位。MRR的响应是洛伦兹线型带宽非常窄这看起来像缺点但实际上特别适合波分复用——不同波长的光经过同一个微环时只有谐振波长的光会被显著调制其他几乎不受影响。用微环阵列做神经网络可以通过波长维度大规模并行。衍射光学神经网络DONN则完全是另一套思路光经过相位调制层后在自由空间传播通过衍射形成下一层的输入光场。透镜的傅里叶变换特性天然对应了频域运算所以DONN特别适合做卷积类的任务。它不需要波导可以用3D打印或者光刻制造多层相位片非常轻巧但精度和可编程性目前不如MZI网格。这三条路线我在模拟中都做了不同程度的尝试最后还是选择了MZI网格作为重点原因是它在数学上和矩阵乘法对应得最干净而且硅光工艺相对成熟。3. 用Python模拟光学选型、复数表示和可微建模3.1 为什么是Python而不是MATLAB做光学仿真很多人第一反应是Lumerical、FDTD或者MATLAB。但我的目标不是做器件级仿真——我不想模拟光的电场分布、波导截面这些物理细节我只想验证一个逻辑如果我有这样一个光学层它能不能通过梯度下降被训练成一个有用的神经网络层。这个目标决定了选型Python NumPy PyTorch。NumPy负责矩阵运算PyTorch负责自动微分。PyTorch的自动微分可以直接对相位参数求导这意味着我完全不需要自己推导梯度公式把光学层定义成一个网络模块然后train和inference就完事了。这就是我说的编程新范式的一个侧面过去我们写Python是在算公式现在写Python是在模拟物理硬件让算法训练的梯度直接穿过光学器件的数学表达。3.2 光的复数表示与自动微分光场的完整信息包含幅度和相位所以我的模拟代码里全程使用complex64或complex128类型存储。这里有一个需要特别小心的点PyTorch对复数梯度的支持虽然越来越好但直接对复数的模平方求导时梯度会涉及Wirtinger导数容易出一些反直觉的结果。我在实际项目里试过两种方案。第一种是前向传播保留复数计算输出强度之后把loss定义在实数域让PyTorch自动处理复数梯度。第二种是把每个复数权重拆成两个实参数幅度和相位然后自动求导。第二种在需要梯度裁剪、正则化、加噪声的时候更好控制所以我最终选了第二种。3.3 把MZI网格矩阵化这可能是整个项目里提升效率最大的一步。如果逐器件去仿真每个MZI、每个交叉连接前向传播会非常慢根本没有办法训练。正确的做法是把整个光学层预计算成一个等价的复数矩阵。一个N×N的MZI网格其传输矩阵可以表示为多个层内MZI矩阵和交叉置换矩阵的连乘积。数学形式是W U_L × ... × U_2 × U_1其中每个U_l是一个块对角矩阵对角块是单个MZI的2×2散射矩阵块之间通过置换矩阵交错连接。整个网格的全部可学习参数就是每个MZI的两个相位值。有了这个矩阵化表示前向传播就只是一次常规的复数矩阵乘法。我写了一个build_mesh_matrix函数把一维相位参数向量重构成MZI对的相位然后计算等效矩阵。这一步做完后一个100×100的MZI网格在GPU上一毫秒内就能完成前向计算。3.4 参数设置与关键代码这里给出一个关键的参数设置思路。MZI网格的每个MZI有两个相位参数通常记为外相位φ和内相位θ。外相位决定这个MZI引入的整体相位偏置内相位决定分束比例也就是能量如何分配到两个输出口。在初始化的时候我建议把所有相位参数初始化为接近0的小随机值而不是均匀分布在0到2π之间。原因是相位参数对输出的影响是周期性的初始化范围太大容易让loss landscape变得非常陡峭训练初期就炸掉。import torch import torch.nn as nn class OpticalLinearLayer(nn.Module): def __init__(self, in_features, out_features, layers3): super().__init__() self.in_features in_features self.out_features out_features self.layers layers # 每个MZI有2个相位参数这里按简化的一层蝶形网络初始化 self.num_mzi in_features // 2 self.phi nn.Parameter(torch.randn(self.layers, self.num_mzi, 2) * 0.1) # 可选的复数缩放用于调节输出幅度 self.scale nn.Parameter(torch.ones(out_features, dtypetorch.float32)) def build_matrix(self): # 根据相位参数构建等效复数矩阵 # 这里省略了具体实现核心是遍历每个MZI并连乘 ... return W需要注意这个简化代码里的build_matrix才是核心它体现的是把物理网络参数化成一个可学习矩阵的思路。你可以在自己的实现中加入非理想因素比如分束比偏差、插入损耗只需要在构建矩阵时把理想BS矩阵替换成带误差的版本。4. 实战在PyTorch里训练一个光子神经网络4.1 任务选择与数据准备我没有一开始就挑战ImageNet那没有任何意义。我先选了一个非常经典的小任务MNIST手写数字的二分类判断一张图是0还是1。这个任务足够小单层光学线性层也能有不错的表现而且整个训练在CPU上几分钟就能跑完适合验证训练流程是否正确。这里有一个实际工程问题MNIST原始图像是28×28784维如果直接输入到MZI网格MZI的数量会随着端口数量平方增长无论仿真还是真实芯片端口数量都是稀缺资源。所以我先用一个简单的平均池化把图像缩小到14×14196维再随机投影到64维。这样既保留了主要特征又把光学层的输入规模控制在合理范围。4.2 实现光学线性层我的OpticalLinearLayer接收复数输入x内部先通过build_matrix()构造等效矩阵W然后执行复数矩阵乘法。输入数据编码上我把MNIST像素值映射到相位上x_i → φ_i π × x_i / x_max。这样光场每个输入端的相位携带了图像的灰度信息。为什么用相位编码而不是直接把像素当复数向量第一相位编码更贴近真实光学硬件的实现方式第二相位不会受到光强的绝对波动影响在模拟中更接近真实光芯片的行为。完整的层定义我放在了项目仓库里这里只给出最核心的部分class OpticalLinearLayer(nn.Module): def forward(self, x): # x: [batch, in_features] 实数值已经编码成相位 x_phase torch.exp(1j * x) # 相位编码成复数光场 W self.build_matrix() # [out_features, in_features] y x_phase W.T # 复数矩阵乘法 return y4.3 探测强度与损失函数光学神经网络在输出端的探测器只能测光强也就是复数振幅的模平方I |y|²。这个模平方操作本身就是一个非线性变换而且是光神经网络里为数不多的天然非线性。我把输出强度送入log_softmax然后计算交叉熵损失。这其实和普通分类网络的区别不大唯一要注意的是输入到softmax的是非负强度而不是任意实数logits。在某些情况下这个非负约束会限制网络的表达能力但同时也带来了更稳定的梯度。4.4 训练结果与对比我把单层光学线性层和传统线性层加ReLU做了对比。训练50个epochAdam优化器学习率1e-3模型参数量测试准确率单层光学线性层 平方探测~20092.1%单层线性层 ReLU~20093.8%两层光学层 平方探测 光学非线性~100097.6%两层线性层 ReLU~100098.4%结果在意料之中纯光学线性层的表达能力确实不如传统带非线性的层因为平方探测虽然是非线性的但它只是一个固定的非线性映射没有可学习的参数。真正把性能拉上来的是两层光学层之间加入可调非线性。这段实验告诉我一个非常重要的规律光子神经网络的瓶颈不在线性代数能力而在非线性能力。如何用光学实现高效可调的非线性函数是整个领域最核心的挑战之一。5. 仿真到芯片之间那些容易翻车的坑5.1 相位噪声第一个绕不过去的坎真实光芯片上的移相器无论是热光还是电光都存在相位漂移和噪声。热光移相器尤其明显温度波动几十毫开尔文相位可能就偏了0.01弧度。在仿真里我开始时完全没有考虑噪声训练得很顺利但一旦在相位参数上加了一点高斯噪声准确率立刻从92%掉到70%以下。这个实验结果让我老老实实把噪声加进了训练流程每次前向传播时给相位参数加上采样自N(0, σ²)的噪声。σ从0.01开始逐步增大到0.05。这样训练出来的模型在推理时把噪声去掉准确率反而比不加噪声训练还要高一点。这本质上就是给光学硬件做数据增强。5.2 制造误差、插损和非理想分束比真实的分束器不可能做到严格的5050工艺制造误差会让分束比在4555到5545之间波动。波导也会有插入损耗光传过整个MZI网格后输出光强会明显衰减。在仿真的理想模型里这些都没被考虑但真机测试时全冒出来了。我把非理想BS矩阵改成了带误差的版本def real_bs(eps0.02): # eps 控制分束比偏离50:50的程度 a 1 / np.sqrt(2) * (1 eps) b 1j / np.sqrt(2) * (1 - eps) return np.array([[a, b], [b, a]])然后把插入损耗建模成一个小于1的衰减系数乘在MZI级联之间。加了这些因素后模型仍然能训练收敛但是最终准确率下降了3到5个百分点。更tricky的是制造误差是固定的而不是随机的所以训练时注入随机噪声只能解决一部分问题更实际的做法是训练后在芯片上做一次标定把每个MZI的实际相位响应测出来再反向补偿。5.3 相位缠绕与梯度震荡如果你直接让相位参数在0到2π之间自由更新很快就会遇到一个数学问题0和2π在物理上是完全相同的相位状态但在数值上是两个完全不同的数字。优化器在更新参数时可能会把某个相位从0.01更新到6.28附近看起来是跳了很长一段距离实际上物理状态几乎没变。反过来在2π边界附近loss对相位的梯度可能非常陡峭造成梯度爆炸。我在训练中遇到了明显的loss震荡后来换了三种解法。最推荐的是第三种把相位参数用cos和sin编码成单位圆上的坐标让优化器在二维平面上更新永远不碰角度边界。用复数权重直接作为可学习参数训练结束后再SVD分解成相位配置。在loss里加一个相位变化量的正则化项抑制过大的相位跳变。我最终采用了第二种方案。这不仅避开了相位缠绕问题还大大简化了代码。代价是训练时学的是复数权重矩阵距离真实硬件还需要一道SVD分解的转换。好在这道转换完全是确定性的用NumPy就可以完成。5.4 Sim2Real把训练好的矩阵映射到真实硬件训练结束后我拿到的是一个复数权重矩阵C。怎么把它变成真实MZI网格的相位配置答案是SVD分解任何复数矩阵都可以分解为C U Σ V†其中U和V是幺正矩阵Σ是实数对角阵。U和V†可以用MZI网格精确实现Σ是实数衰减也可以用电控可调衰减器实现。整个映射过程在Python里非常简单U, S, Vh np.linalg.svd(C) # U 和 Vh 用 MZI 网格实现S 用衰减器实现这一步做完我算是真正打通了训练和硬件部署之间的链路。PyTorch训练出来的权重转换成了控制MZI移相器的电压配置理论上可以直接下载到可编程光芯片上执行。这里有一个让我印象深刻的点SVD分解并不是唯一的分解方式。Clements分解和Reck分解是两种不同的MZI网格拓扑各有优劣。Clements分解对制造误差更鲁棒Reck分解更直观但更容易积累误差。实际芯片部署建议选Clements拓扑。6. 怎么理解光计算驱动的编程新范式6.1 编程对象变了从指令到相位配置传统编程的核心是控制流、数据结构、内存管理。你写的每一行代码最终都被编译成CPU或GPU的指令流。但在光计算里程序变成了一组移相器的电压配置或者说一列相位值。整个程序空间几乎不存在条件跳转和循环存在的只是不同长度、不同结构的参数数组。这种转变对我来说是很颠覆的写神经网络训练代码实际是在为光学器件找一组最优的物理配置。我们习惯的软件开发工具链——IDE、调试器、版本控制——依然有用但你编写的是一组连续的实数参数而不是离散的语法结构。这就是我理解的编程新范式。6.2 神经网络在这里变成了硬件编译器在传统深度学习里神经网络是纯粹的数学函数逼近器输出的是预测结果。但在光子神经网络项目里训练出来的模型同时还隐含了一组硬件控制参数。PyTorch模型通过反向传播得到的梯度最终转换成MZI网格的相位配置——这个链条本质上是在做编译把神经网络的数学结构编译成光学器件的物理配置。我自己的代码里已经验证了这个流程训练复数权重SVD分解得到U、Σ、V†然后映射到MZI网格和衰减器。整个过程完全是确定性的、可复现的。这就是一个光学编译器的最小闭环。6.3 一点个人体会为什么值得现在就上手接触这个项目之前我一直觉得光计算是物理学家和芯片工程师的领域跟搞算法的人没什么关系。但做完这套Python模拟之后我的看法彻底变了光计算最有趣的地方恰恰是它把物理层和算法层拉到了同一个可微的编程框架里。我不需要对电磁波方程有多深的理解不需要碰工艺制造也可以研究光子神经网络的训练策略、噪声鲁棒性、网络拓扑设计。而反过来如果我把网络再复杂一点、任务再贴近实际一点这套代码完全可以作为光芯片流片前的算法验证工具。最后分享一个很有用的建议如果你也想入手这个方向不要从零开始重写所有代码。GitHub上已经有一些开源的光子神经网络模拟库比如TorchONN、DeepLight它们提供了现成的光学层组件和训练范式。我在第一次尝试时就是自己从零写了一段光学层结果跑通后才发现开源库里有更高效、更成熟的实现白白浪费了好几天时间。拿着这些库跑通一个小demo读一遍源码理解它们的建模方式然后在此基础上改造成你自己的模型这条路比从零开始顺利得多。