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

應大衛 張思慧 鄧書金 武海斌

(華東師范大學,精密光譜科學與技術國家重點實驗室,上海 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].