基于機器學習的單拍冷原子成像*

應大衛(wèi) 張思慧 鄧書金 武海斌

(華東師范大學,精密光譜科學與技術(shù)國家重點實驗室,上海 200062)

在冷原子研究中,通常采用吸收成像的方式來進行冷原子狀態(tài)的探測,然而該成像方式受探測過程中光學平面干涉、激光功率、頻率、空間位置抖動等方面的擾動,最終形成的冷原子圖像總是會出現(xiàn)殘余部分空間結(jié)構(gòu)噪聲,導致成像質(zhì)量的下降.尤其是對于冷原子密度稀薄的區(qū)域或者飛行時間較久的情況,往往需要大量的重復和平均才能得到理想的信噪比,然而這樣不僅導致時間周期大幅度增加,還會引入大量隨機噪聲.本文基于機器學習提出了一種單拍冷原子成像方案,在該方案中僅需對冷原子進行單次吸收成像,對應背景圖片可以通過自動編碼器的神經(jīng)網(wǎng)絡(luò)來進行生成,有效地降低了成像的空間條紋噪聲,大幅度提高成像質(zhì)量,可以用于冷原子單循環(huán)多次成像.

1 引言

超冷原子分子氣體是高度空間隔離的純凈量子多體系統(tǒng),具有很好的量子相干性,被廣泛應用于量子信息[1,2]、量子計算和模擬[3-5]以及精密測量[6-8]的研究.準確測量和分析超冷原子分子的量子特性對后續(xù)研究十分重要,冷原子吸收成像是當前實驗中最常采用的原子狀態(tài)探測方式,該方案利用一束準直的共振激光去照射冷原子氣體,將原子散射的光信號連同未被散射的光信號收集到相機中[9],通過從位置相關(guān)的散射信號中提取原子空間分布,進而分析原子熱力學狀態(tài)等特性.

標準的吸收成像方案需要采集兩張圖像: 第一次曝光采集的信號是原子存在時的吸收圖像,將共振激光通過原子云后的吸收和散射信號進行收集;第二次曝光采集的是背景信號,不包含任何原子信息,用于進行原子背景參考校正,通常是在第一次曝光后等待原子移出拍攝區(qū)域或利用光學方式將原子泵浦到暗態(tài)之后再進行.通過對比吸收信號和參考背景信號得到原子的光學密度,進而得到原子的空間分布.在實驗中,探測激光的相干長度通常比光學元器件之間的距離長得多,各光學元件間的殘余反射會干擾圖像并在記錄的圖像中產(chǎn)生條紋和牛頓環(huán)等特征圖案.然而,由于噪聲漂移和其他動態(tài)過程,兩幅圖像中的噪聲模式通常并不完全相同,這導致最終圖像中出現(xiàn)殘差結(jié)構(gòu)噪聲模式,并最終導致成像質(zhì)量的大幅度下降.

機器學習通過多層神經(jīng)網(wǎng)絡(luò)來傳播信息,將給定的輸入數(shù)據(jù)轉(zhuǎn)換為特定的預測.近年來,機器學習在圖像處理等領(lǐng)域取得了巨大成功[10-14].圖像補全是機器學習應用的一個很好的例子.對于那些存在反復出現(xiàn)且不斷變化的圖像條紋場景,機器學習技術(shù)的優(yōu)勢更為明顯.圖像補全技術(shù)可以從部分未損壞的圖像中推斷出圖像損壞的部分,使圖像看起來更完整.機器學習技術(shù)已被應用于優(yōu)化超冷原子的冷卻過程和執(zhí)行相關(guān)的數(shù)值計算[15-17].在這些應用中,機器學習模型可以學習如何最大限度地減少原子的熱運動,從而使其達到極低的溫度.

本文基于機器學習的技術(shù)方法提出了一種冷原子成像方案,該方案中背景圖片采用自動編碼器的神經(jīng)網(wǎng)絡(luò)來進行生成,用于消除成像中引入的空間條紋噪聲,提高成像質(zhì)量.本文使用一種深度神經(jīng)網(wǎng)絡(luò)(DNN)的冷原子成像替代方法,從包含原子散射信號的單個圖像中生成一個理想的背景圖像.在編碼過程中還引入了注意力機制,使得神經(jīng)網(wǎng)絡(luò)可以自動在更加重要的圖樣紋路上賦予更多的權(quán)重,并最終大幅度地提高了成像質(zhì)量.該單拍成像方法簡化了實驗操作步驟,降低了對相機的硬件要求,縮短了觀測原子動力學演化的時間并有效地抑制了條紋干涉,可應用于冷原子單循環(huán)多次成像.

2 成像裝置和深度神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)及其訓練

本實驗數(shù)據(jù)庫來自超冷量子氣體的吸收成像,實驗中首先將6Li 原子氣體裝載到光偶極阱中,并制備在最低的兩個自旋態(tài)(|1〉=|F=1/2,MF=1/2〉和| 2〉=|F=1/2,MF=-1/2〉).接下來利用原子間的Feshbach 共振進行蒸發(fā)冷卻,此時偏置磁場強度為300 G (1 G=10—4T),對應散射長度約為—290a0(玻爾半徑),蒸發(fā)冷卻約8 s 時間后將磁場偏轉(zhuǎn)至528 G 并對|1〉態(tài)進行吸收成像,原子溫度T/TF約為0.3,詳細實驗系統(tǒng)和冷卻過程參考文獻[18-21].成像激光的光斑直徑大約1 cm,光功率密度大約0.2Is(飽和光強),成像脈沖時間10 μs.實驗中使用的相機是Andor iKon-M934,具有1024×1024 個像素.考慮到電子噪聲和相應速度,拍照時選擇的相機參數(shù)是: 讀取率為1 MHz,16 bit,轉(zhuǎn)移速度為3.3 μs.由于相機本身幀轉(zhuǎn)移速率限制,信號光脈沖和背景光脈沖間隔時間300 ms,該數(shù)據(jù)庫包含1150 組圖片,每組圖片包含1 個信號光和1 個背景光.

基于該數(shù)據(jù)庫,深度神經(jīng)網(wǎng)絡(luò)(DNN)建立了一個通道,輸入的吸收圖像首先要將原子氣體所在區(qū)域進行掩碼屏蔽,僅抽取背景部分的條紋信息,再經(jīng)歷多次卷積轉(zhuǎn)換和維度變換.這些轉(zhuǎn)換提取了底層空間模式的特征,對DNN 的預測起到了至關(guān)重要的作用.訓練階段主要使用背景圖片庫進行,用于空間噪聲模式的對比和優(yōu)化重構(gòu).在每個優(yōu)化步驟中,將網(wǎng)絡(luò)的預測結(jié)果與掩蔽區(qū)域的真實值進行比較,并改變模型的權(quán)重以最小化損失函數(shù),即真實值與預測結(jié)果之間的均方誤差(L2 范數(shù)).在訓練結(jié)束時,將獲得一個優(yōu)化的模型,用于預測包含原子信息在內(nèi)的新圖像.然而因為原子所在區(qū)域被掩蓋了,所以無論原子是否出現(xiàn)在原始圖像中,該網(wǎng)絡(luò)都會產(chǎn)生一個理想的背景圖像.

卷積網(wǎng)絡(luò)(圖1)是由U-net 架構(gòu)[22]的自動編碼器和解碼器組成的,網(wǎng)絡(luò)中的跳層連接結(jié)構(gòu)可以保證網(wǎng)絡(luò)在很深的情況下也能夠使得網(wǎng)絡(luò)前端卷積核權(quán)重進行有效的更新.通過實驗發(fā)現(xiàn),加入注意力機制[7]會使得在此任務(wù)上的表現(xiàn)更加好.網(wǎng)絡(luò)的輸入是原子吸收成像的圖像,從中掩蔽了寬度為120 個像素的條形區(qū)域,用于掩蔽原子本身的散射光信號.訓練中使用一個生成器來打亂圖像,在輸入上應用掩碼,并以8 張圖片為一批輸入DNN進行訓練.

圖1 U-net 神經(jīng)網(wǎng)絡(luò)架構(gòu)Fig.1.U-net neural network architecture.

神經(jīng)網(wǎng)絡(luò)將整體數(shù)據(jù)集分為三類數(shù)據(jù)集,訓練集、驗證集和測試集,其中訓練集占據(jù)60%,驗證集占據(jù)20%,測試集占據(jù)20%,訓練集用于訓練和生成神經(jīng)網(wǎng)絡(luò),驗證集用于驗證神經(jīng)網(wǎng)絡(luò)的準確性,測試集用于最終的實際測試.該網(wǎng)絡(luò)由74457081 個參數(shù)組成,共34 層,其中可訓練參數(shù)74433017 個,不可訓練參數(shù)24064 個(主要是批歸一化層中的平均值與方差).將網(wǎng)絡(luò)輸出與每張圖像的原始中心部分進行比較,并最小化均方誤差損失函數(shù).使用ADAM 優(yōu)化器[23]和He 初始化[24]進行參數(shù)優(yōu)化,學習率為5×10—3[25],應用0.99 批歸一化[26]和通道注意力機制[10,11]可以進一步使得模型的表現(xiàn)能力更加好.

3 實驗結(jié)果分析

3.1 構(gòu)建神經(jīng)網(wǎng)絡(luò)

本文中僅有訓練集參與神經(jīng)網(wǎng)絡(luò)的訓練,在多次訓練后形成相應的神經(jīng)網(wǎng)絡(luò).基于相應次數(shù)的神經(jīng)網(wǎng)絡(luò),將利用驗證集中的圖庫數(shù)據(jù)來驗證其準確性.圖像的損失函數(shù)通常用于衡量模型的輸出圖像與真實圖像之間的差異.圖像損失函數(shù)表示為

其中,I表示真實圖像,表示生成的圖像,W,H和C分別表示圖像的寬度、高度和通道數(shù).該損失函數(shù)計算了兩個圖像之間每個像素通道的均方誤差(MSE),然后對所有像素的誤差求平均,最終的損失值越小,表示生成的圖像與真實圖像之間的差異越小.因此,原則上訓練集和驗證集的損失函數(shù)都應該同步降低并最終收斂.

模型的收斂性如圖2 所示,其中訓練(藍色)和驗證(黃色)數(shù)據(jù)集的損失函數(shù)隨著訓練次數(shù)的增加而減少,兩個數(shù)據(jù)集(訓練集和驗證集)的衰減具有相似的趨勢.而當訓練次數(shù)超過200 次之后,盡管訓練集的損失函數(shù)仍在不停降低,驗證集的損失函數(shù)不再振蕩并保持收斂趨勢,甚至存在非常微弱的上升趨勢,這就意味著模型的進一步訓練和優(yōu)化將不再對真實的背景校正有本質(zhì)影響,為了防止過擬合,最后取驗證集表現(xiàn)最好時候的模型作為預測模型.

圖2 損失函數(shù)隨著DNN 訓練次數(shù)的變化過程,即模型預測和真實值之間的均方誤差作為訓練迭代次數(shù)的函數(shù),較低的值意味著更好的預測性能.黃色的曲線表示驗證集的損失函數(shù),藍色曲線是訓練集上的損失函數(shù)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.

而原則上,只要驗證集損失函數(shù)不再發(fā)生變化,對于神經(jīng)網(wǎng)絡(luò)的訓練就可以停止,推測超過200 次訓練后訓練集的損失函數(shù)仍在降低的現(xiàn)象來自于較小的數(shù)據(jù)庫引起的局域化分布.當訓練次數(shù)達到1000 次之后,我們停止了訓練,基于此構(gòu)建了整體的神經(jīng)網(wǎng)絡(luò).

3.2 成像質(zhì)量分析

訓練生成的神經(jīng)網(wǎng)絡(luò)可以應用于原子信息重構(gòu),首先進行的是背景圖片重構(gòu).選取測試集中無原子信息存在的區(qū)域進行對比研究,由于這部分特定選取的區(qū)域不存在原子分布,可以用于對比研究生成的背景圖片與實際拍照獲取的背景圖片的噪聲分布.為了比較兩者之間的差異,同時采用冷原子光學密度(optical density,OD)的表征方法來進行描述,O D=log[Iref/Isig] ,其中Isig為實際拍攝的信號圖片,Iref為相應的背景圖片,可以為預測生成或者實際拍攝.

圖3 展示了一組無原子信息存在時的實驗結(jié)果,其中圖3(a)和圖3(b) 分別展示的是預測生成和實際拍攝的光學密度分布,可以發(fā)現(xiàn)背景圖片扣除后圖3(a)中的噪聲明顯小于圖3(b)中的噪聲,且?guī)缀蹩床坏酱嬖诳臻g關(guān)聯(lián)的條紋噪聲.這說明利用已訓練的神經(jīng)網(wǎng)絡(luò)進行預測生成的背景圖片更加貼近于真實噪聲分布,而雙次拍照獲得的圖像信息由于殘差干涉噪聲的存在信噪相對較差.

圖3 神經(jīng)網(wǎng)絡(luò)預測的成像信息和實際拍攝的結(jié)果對比圖(無原子信息部分) (a),(c) 神經(jīng)網(wǎng)絡(luò)預測的光學密度圖及其傅里葉空間分布;(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.

將殘余噪聲轉(zhuǎn)換為傅里葉空間的分布上區(qū)分則更加清楚,如圖3(c)和圖3(d)所示,在相同的幅度顯示下,雙拍成像帶來了更多的空間條紋噪聲,對其進行角平均后可以發(fā)現(xiàn)其空間頻率主要分布于0.05 μm—1附近(圖3(e)),而采用神經(jīng)網(wǎng)絡(luò)生成的背景圖片進行背景扣除后OD 的測量結(jié)果則可以大幅度抑制空間條紋噪聲,對于某些特定的空間頻率,其噪聲強度甚至可以抑制到雙拍成像的1/20 左右.

該神經(jīng)網(wǎng)絡(luò)構(gòu)建的吸收成像優(yōu)化模型不僅可以優(yōu)化背景圖片中存在的空間條紋噪聲,同樣也可以用于優(yōu)化原子信息處的噪聲分布.本次測試中原子數(shù)目總計大約 20 萬個,但是在一般雙拍吸收成像方式中由于殘差噪聲的存在導致原子處能看到清晰的條紋結(jié)構(gòu),如圖4(b)所示,表現(xiàn)出原子分布被空間噪聲所調(diào)制,一般只能通過多次平均來克服這種噪聲干擾,但是往往需要耗費大量的時間成本的同時還會引入隨機噪聲.圖4(a)展示了基于神經(jīng)網(wǎng)絡(luò)構(gòu)建背景圖片優(yōu)化后的原子密度分布,可以看到原子的密度分布變得更加的平滑,空間條紋噪聲被大幅度的抑制,成像得到了很好的優(yōu)化.

圖4 神經(jīng)網(wǎng)絡(luò)預測的成像信息和實際拍攝的結(jié)果對比圖(原子信息部分) (a) 經(jīng)網(wǎng)絡(luò)預測生成的無條紋的原子密度分布圖;(b) 使用傳統(tǒng)雙拍成像得到的原子密度分布;(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方向進行積分得到結(jié)果,如圖4(c)所示,經(jīng)算法優(yōu)化后的原子密度分布更加平滑,信噪比大幅度提高.同時值得注意的是,傳統(tǒng)雙拍成像得到的密度分布比算法校正后的原子密度分布更高,意味著部分條紋噪聲被誤計入原子數(shù)目中,影響實驗中對原子實際參數(shù)的評估.經(jīng)計算實際大約13%左右的原子應該被移出總原子基數(shù),這將最終導致實際費米面縮小2%左右.

4 總結(jié)與展望

本文展示了一個基于深度卷積神經(jīng)網(wǎng)絡(luò)背景還原的單拍冷原子成像技術(shù),證實了這種方案可以大幅度抑制由于空間干涉帶來的空間條紋噪聲干擾,準確地重建比傳統(tǒng)的雙拍成像技術(shù)質(zhì)量更好的冷原子成像.同時利用該單拍成像方案可以降低相機高幀轉(zhuǎn)移速率的要求,結(jié)合偏振成像技術(shù)、相襯成像技術(shù)等無損探測技術(shù)可應用于冷原子單循環(huán)多次成像,便于更快地還原原子的動力學行為.該機器學習方案具有很好的魯棒性,將來可以通過在線學習方案進行進一步的優(yōu)化,通過把圖像定期添加到數(shù)據(jù)集,使得模型可以一邊預測一邊不斷更新,從而獲得更高的預測精度[27].