基于机器学习的单拍冷原子成像*

应大卫 张思慧 邓书金 武海斌

(华东师范大学,精密光谱科学与技术国家重点实验室,上海 200062)

在冷原子研究中,通常采用吸收成像的方式来进行冷原子状态的探测,然而该成像方式受探测过程中光学平面干涉、激光功率、频率、空间位置抖动等方面的扰动,最终形成的冷原子图像总是会出现残余部分空间结构噪声,导致成像质量的下降.尤其是对于冷原子密度稀薄的区域或者飞行时间较久的情况,往往需要大量的重复和平均才能得到理想的信噪比,然而这样不仅导致时间周期大幅度增加,还会引入大量随机噪声.本文基于机器学习提出了一种单拍冷原子成像方案,在该方案中仅需对冷原子进行单次吸收成像,对应背景图片可以通过自动编码器的神经网络来进行生成,有效地降低了成像的空间条纹噪声,大幅度提高成像质量,可以用于冷原子单循环多次成像.

1 引言

超冷原子分子气体是高度空间隔离的纯净量子多体系统,具有很好的量子相干性,被广泛应用于量子信息[1,2]、量子计算和模拟[3-5]以及精密测量[6-8]的研究.准确测量和分析超冷原子分子的量子特性对后续研究十分重要,冷原子吸收成像是当前实验中最常采用的原子状态探测方式,该方案利用一束准直的共振激光去照射冷原子气体,将原子散射的光信号连同未被散射的光信号收集到相机中[9],通过从位置相关的散射信号中提取原子空间分布,进而分析原子热力学状态等特性.

标准的吸收成像方案需要采集两张图像: 第一次曝光采集的信号是原子存在时的吸收图像,将共振激光通过原子云后的吸收和散射信号进行收集;第二次曝光采集的是背景信号,不包含任何原子信息,用于进行原子背景参考校正,通常是在第一次曝光后等待原子移出拍摄区域或利用光学方式将原子泵浦到暗态之后再进行.通过对比吸收信号和参考背景信号得到原子的光学密度,进而得到原子的空间分布.在实验中,探测激光的相干长度通常比光学元器件之间的距离长得多,各光学元件间的残余反射会干扰图像并在记录的图像中产生条纹和牛顿环等特征图案.然而,由于噪声漂移和其他动态过程,两幅图像中的噪声模式通常并不完全相同,这导致最终图像中出现残差结构噪声模式,并最终导致成像质量的大幅度下降.

机器学习通过多层神经网络来传播信息,将给定的输入数据转换为特定的预测.近年来,机器学习在图像处理等领域取得了巨大成功[10-14].图像补全是机器学习应用的一个很好的例子.对于那些存在反复出现且不断变化的图像条纹场景,机器学习技术的优势更为明显.图像补全技术可以从部分未损坏的图像中推断出图像损坏的部分,使图像看起来更完整.机器学习技术已被应用于优化超冷原子的冷却过程和执行相关的数值计算[15-17].在这些应用中,机器学习模型可以学习如何最大限度地减少原子的热运动,从而使其达到极低的温度.

本文基于机器学习的技术方法提出了一种冷原子成像方案,该方案中背景图片采用自动编码器的神经网络来进行生成,用于消除成像中引入的空间条纹噪声,提高成像质量.本文使用一种深度神经网络(DNN)的冷原子成像替代方法,从包含原子散射信号的单个图像中生成一个理想的背景图像.在编码过程中还引入了注意力机制,使得神经网络可以自动在更加重要的图样纹路上赋予更多的权重,并最终大幅度地提高了成像质量.该单拍成像方法简化了实验操作步骤,降低了对相机的硬件要求,缩短了观测原子动力学演化的时间并有效地抑制了条纹干涉,可应用于冷原子单循环多次成像.

2 成像装置和深度神经网络结构及其训练

本实验数据库来自超冷量子气体的吸收成像,实验中首先将6Li 原子气体装载到光偶极阱中,并制备在最低的两个自旋态(|1〉=|F=1/2,MF=1/2〉和| 2〉=|F=1/2,MF=-1/2〉).接下来利用原子间的Feshbach 共振进行蒸发冷却,此时偏置磁场强度为300 G (1 G=10—4T),对应散射长度约为—290a0(玻尔半径),蒸发冷却约8 s 时间后将磁场偏转至528 G 并对|1〉态进行吸收成像,原子温度T/TF约为0.3,详细实验系统和冷却过程参考文献[18-21].成像激光的光斑直径大约1 cm,光功率密度大约0.2Is(饱和光强),成像脉冲时间10 µs.实验中使用的相机是Andor iKon-M934,具有1024×1024 个像素.考虑到电子噪声和相应速度,拍照时选择的相机参数是: 读取率为1 MHz,16 bit,转移速度为3.3 µs.由于相机本身帧转移速率限制,信号光脉冲和背景光脉冲间隔时间300 ms,该数据库包含1150 组图片,每组图片包含1 个信号光和1 个背景光.

基于该数据库,深度神经网络(DNN)建立了一个通道,输入的吸收图像首先要将原子气体所在区域进行掩码屏蔽,仅抽取背景部分的条纹信息,再经历多次卷积转换和维度变换.这些转换提取了底层空间模式的特征,对DNN 的预测起到了至关重要的作用.训练阶段主要使用背景图片库进行,用于空间噪声模式的对比和优化重构.在每个优化步骤中,将网络的预测结果与掩蔽区域的真实值进行比较,并改变模型的权重以最小化损失函数,即真实值与预测结果之间的均方误差(L2 范数).在训练结束时,将获得一个优化的模型,用于预测包含原子信息在内的新图像.然而因为原子所在区域被掩盖了,所以无论原子是否出现在原始图像中,该网络都会产生一个理想的背景图像.

卷积网络(图1)是由U-net 架构[22]的自动编码器和解码器组成的,网络中的跳层连接结构可以保证网络在很深的情况下也能够使得网络前端卷积核权重进行有效的更新.通过实验发现,加入注意力机制[7]会使得在此任务上的表现更加好.网络的输入是原子吸收成像的图像,从中掩蔽了宽度为120 个像素的条形区域,用于掩蔽原子本身的散射光信号.训练中使用一个生成器来打乱图像,在输入上应用掩码,并以8 张图片为一批输入DNN进行训练.

图1 U-net 神经网络架构Fig.1.U-net neural network architecture.

神经网络将整体数据集分为三类数据集,训练集、验证集和测试集,其中训练集占据60%,验证集占据20%,测试集占据20%,训练集用于训练和生成神经网络,验证集用于验证神经网络的准确性,测试集用于最终的实际测试.该网络由74457081 个参数组成,共34 层,其中可训练参数74433017 个,不可训练参数24064 个(主要是批归一化层中的平均值与方差).将网络输出与每张图像的原始中心部分进行比较,并最小化均方误差损失函数.使用ADAM 优化器[23]和He 初始化[24]进行参数优化,学习率为5×10—3[25],应用0.99 批归一化[26]和通道注意力机制[10,11]可以进一步使得模型的表现能力更加好.

3 实验结果分析

3.1 构建神经网络

本文中仅有训练集参与神经网络的训练,在多次训练后形成相应的神经网络.基于相应次数的神经网络,将利用验证集中的图库数据来验证其准确性.图像的损失函数通常用于衡量模型的输出图像与真实图像之间的差异.图像损失函数表示为

其中,I表示真实图像,表示生成的图像,W,H和C分别表示图像的宽度、高度和通道数.该损失函数计算了两个图像之间每个像素通道的均方误差(MSE),然后对所有像素的误差求平均,最终的损失值越小,表示生成的图像与真实图像之间的差异越小.因此,原则上训练集和验证集的损失函数都应该同步降低并最终收敛.

模型的收敛性如图2 所示,其中训练(蓝色)和验证(黄色)数据集的损失函数随着训练次数的增加而减少,两个数据集(训练集和验证集)的衰减具有相似的趋势.而当训练次数超过200 次之后,尽管训练集的损失函数仍在不停降低,验证集的损失函数不再振荡并保持收敛趋势,甚至存在非常微弱的上升趋势,这就意味着模型的进一步训练和优化将不再对真实的背景校正有本质影响,为了防止过拟合,最后取验证集表现最好时候的模型作为预测模型.

图2 损失函数随着DNN 训练次数的变化过程,即模型预测和真实值之间的均方误差作为训练迭代次数的函数,较低的值意味着更好的预测性能.黄色的曲线表示验证集的损失函数,蓝色曲线是训练集上的损失函数Fig.2.The loss function varies with the training iterations of the DNN,which is the mean square error between model predictions and true values as a function of training iterations(epochs),and a lower value indicates better prediction performance.The yellow line represents the loss function on the validation set,and the blue line represents the loss function on the training set.

而原则上,只要验证集损失函数不再发生变化,对于神经网络的训练就可以停止,推测超过200 次训练后训练集的损失函数仍在降低的现象来自于较小的数据库引起的局域化分布.当训练次数达到1000 次之后,我们停止了训练,基于此构建了整体的神经网络.

3.2 成像质量分析

训练生成的神经网络可以应用于原子信息重构,首先进行的是背景图片重构.选取测试集中无原子信息存在的区域进行对比研究,由于这部分特定选取的区域不存在原子分布,可以用于对比研究生成的背景图片与实际拍照获取的背景图片的噪声分布.为了比较两者之间的差异,同时采用冷原子光学密度(optical density,OD)的表征方法来进行描述,O D=log[Iref/Isig] ,其中Isig为实际拍摄的信号图片,Iref为相应的背景图片,可以为预测生成或者实际拍摄.

图3 展示了一组无原子信息存在时的实验结果,其中图3(a)和图3(b) 分别展示的是预测生成和实际拍摄的光学密度分布,可以发现背景图片扣除后图3(a)中的噪声明显小于图3(b)中的噪声,且几乎看不到存在空间关联的条纹噪声.这说明利用已训练的神经网络进行预测生成的背景图片更加贴近于真实噪声分布,而双次拍照获得的图像信息由于残差干涉噪声的存在信噪相对较差.

图3 神经网络预测的成像信息和实际拍摄的结果对比图(无原子信息部分) (a),(c) 神经网络预测的光学密度图及其傅里叶空间分布;(b),(d) 实际双拍成像的光学密度图及其傅里叶空间分布;(e) 频谱空间的噪声对比图Fig.3.Comparison of images predicted by neural networks and actual captured results (excluding atomic information): (a),(c) The optical density and Fourier space distribution of images generated by neural networks;(b),(d) the optical density and Fourier space distribution of actual double-shot imaging;(e) comparison of noise in frequency space.

将残余噪声转换为傅里叶空间的分布上区分则更加清楚,如图3(c)和图3(d)所示,在相同的幅度显示下,双拍成像带来了更多的空间条纹噪声,对其进行角平均后可以发现其空间频率主要分布于0.05 µm—1附近(图3(e)),而采用神经网络生成的背景图片进行背景扣除后OD 的测量结果则可以大幅度抑制空间条纹噪声,对于某些特定的空间频率,其噪声强度甚至可以抑制到双拍成像的1/20 左右.

该神经网络构建的吸收成像优化模型不仅可以优化背景图片中存在的空间条纹噪声,同样也可以用于优化原子信息处的噪声分布.本次测试中原子数目总计大约 20 万个,但是在一般双拍吸收成像方式中由于残差噪声的存在导致原子处能看到清晰的条纹结构,如图4(b)所示,表现出原子分布被空间噪声所调制,一般只能通过多次平均来克服这种噪声干扰,但是往往需要耗费大量的时间成本的同时还会引入随机噪声.图4(a)展示了基于神经网络构建背景图片优化后的原子密度分布,可以看到原子的密度分布变得更加的平滑,空间条纹噪声被大幅度的抑制,成像得到了很好的优化.

图4 神经网络预测的成像信息和实际拍摄的结果对比图(原子信息部分) (a) 经网络预测生成的无条纹的原子密度分布图;(b) 使用传统双拍成像得到的原子密度分布;(c) 白框中原子沿x 方向积分后的一维密度分布图Fig.4.Comparison of images predicted by neural networks and actual captured results (atomic information included): (a) Stripefree atomic density distribution map generated by neural networks;(b) atomic density distribution obtained using traditional double-shot imaging;(c) one-dimensional density distribution in the white square (integrated along x direction).

由于原子处空间噪声有多种频率成分,条纹幅度和方向也有很大不同,随机选取了图4(a)和图4(b)白框中的部分原子来展示算法对原子噪声的抑制作用,对白框中的原子沿着x方向进行积分得到结果,如图4(c)所示,经算法优化后的原子密度分布更加平滑,信噪比大幅度提高.同时值得注意的是,传统双拍成像得到的密度分布比算法校正后的原子密度分布更高,意味着部分条纹噪声被误计入原子数目中,影响实验中对原子实际参数的评估.经计算实际大约13%左右的原子应该被移出总原子基数,这将最终导致实际费米面缩小2%左右.

4 总结与展望

本文展示了一个基于深度卷积神经网络背景还原的单拍冷原子成像技术,证实了这种方案可以大幅度抑制由于空间干涉带来的空间条纹噪声干扰,准确地重建比传统的双拍成像技术质量更好的冷原子成像.同时利用该单拍成像方案可以降低相机高帧转移速率的要求,结合偏振成像技术、相衬成像技术等无损探测技术可应用于冷原子单循环多次成像,便于更快地还原原子的动力学行为.该机器学习方案具有很好的鲁棒性,将来可以通过在线学习方案进行进一步的优化,通过把图像定期添加到数据集,使得模型可以一边预测一边不断更新,从而获得更高的预测精度[27].