自主設計實驗,速度提高2.5倍!馬普所AI方法擴充顯微鏡使用邊界

ScienceAI發表於2024-12-13

圖片

編輯丨&

在人類的聰明才智與創造力的推動之下,光學顯微鏡突破經典光衍射極限,在質量與數量上有了顯著的飛躍。然而,在如此廣袤的選擇空間與錯綜複雜的實驗配置組合中,一些高效的設計尚未被發現,並且可能永遠不會採用人工驅動。

於此,德國馬克斯普朗克光科學研究所(Max Planck Institute for the Science of Light)的研究團隊介紹了一個基於 JAX(Python 中的高效能運算庫)開發的開源計算框架 XLuminA 。其可以快速準確的探索實驗儀器的組配方式,在自己的虛擬光學設計環境中模擬出來。

透過 JAX 的加速線性代數編譯器、即時編譯及其無縫整合的自動向量化、自動微分與 GPU 相容性提供的計算速度,XLuminA 的加速相較於目前成熟的數值最佳化方法提高了四個數量級。

該研究以「Automated discovery of experimental designs in super-resolution microscopy with XLuminA」為題,於2024年12月10日刊登於《Nature Communications》。

圖片

在過去的二十年中,幾項突破進一步拓寬了光學顯微鏡在生物科學的應用範圍,多功能且強大的超解析度(SR)技術對不僅限於生物學領域也同樣包括化學與材料科學的技術範疇產生了相當大的影響。

這個研究的初始動機遠遠超過了對已知光學技術的小規模最佳化,旨在發現目前完全未開發的先進光學顯微鏡的新穎、可用於實驗上的概念。

XLuminA

馬克斯普朗克光科學研究所的研究團隊提出了 XLuminA ,它可以將預設空間內的某一處模擬位置轉換為我們需要的物理輸出,然後在目標函式里使用該輸出來描述所需的設計目標。

模擬器可以透過梯度的最佳化技術直接呼叫,也可以用於為基於深度學習的代理模型生成訓練資料。用於自動設計與發現策略的模擬器必須高速可靠與通用,而 XLuminA 的光學模擬器很明顯滿足以上要求。

圖片

圖 1 :XLuminA 的概述與效能。

團隊使用的方法與以前利用 AI 去進行單個光學元件的資料驅動設計的方式截然不同。雖然這些技術很有影響力,但它們並不意味著要改變實驗方法的原理或光學佈局本身。

相比之下,XLuminA 配備了從頭開始模擬、最佳化和自動設計新光學設定和概念的工具。由於模擬中有不同的物理目標,XLuminA 可以在光束傳播中應用不同近似值。相較於受物理技術啟發的伴隨法,它極大地加快了模擬器的速度。

有趣的是,伴隨法可以看做是自動微分(團隊所採用)的一個特例,研究員在文章中表示道。

工作流程與效能

XLuminA 允許模擬經典光學硬體配置,並支援最佳化和自動發現未探索的設定設計。該軟體是使用 JAX 開發的,它透過加速線性代數編譯器 XLA 實現增速優勢,同時無縫整合自動微分框架和自動 GPU 相容性。

需要注意的是,XLuminA 並不侷限於在 CPU 上執行。由於 JAX 整合,在預設情況下是執行在 GPU 上,否則會自動退回到 CPU 。

XLuminA 在模擬光衍射和傳播方面顯著提高了計算速度。例如,與 CPU 上的 RS(Rayleigh-Sommerfeld)和 CZT(Chirped z-transform)相比,速度提升了約 2 倍,對於 VRS 和 VCZT(RS 和 CZT 的向量化版本)分別提升了約2.5 倍。使用 GPU 時,這些提升因子分別為 64、76、80 和 78。

為了包含自動發現功能,XLuminA 的光學模擬器和最佳化器由損失函式捆綁在一起。在這項工作中,團隊採用了一種基於梯度的策略,其中實驗裝置的引數在最陡的下降方向上迭代調整。

為了評估數值與分析方法計算一個梯度的評估時間以及他們在不同解析度和裝置上的收斂時間,團隊準備了梯度下降演算法:Broyden-Fletcher-Goldfarb-Shanno (BFGS) 演算法。

最後,利用 JAX 的內建 autodiff 框架,解析計算損失函式的梯度。這種方法使最佳化器能夠有效地構建內部梯度函式,從而大大減少每次迭代的計算時間。

重新發現資料驅動

團隊的目標是使用 XLuminA 來發現新的顯微鏡概念。從本質上講,發現新的實驗配置需要一個混合離散-連續搜尋問題。離散方面源於配置光網路拓撲,而連續部分與光學元件的設定有關,例如鐳射功率和分束器反射率。

為了實現純連續最佳化實驗設計拓撲的自動發現及其引數設定,他們採用了大型複雜的光學拓撲來初始化設定。從這裡開始,XLuminA 應該能夠提取出人類可能還沒有考慮過的更復雜的解決方案。

圖片

圖 2 :在高度引數化的光學設定中發現未報導的實驗藍圖。

令人驚訝的是,XLuminA 找到了一種替代方法,可以將相位奇點壓印到光束上,並在焦平面上產生明顯的縱向分量。在高引數化與複雜光學系統中的情況下,XLuminA 都成功地發現了替代光學解決方案,表現出與參考實驗相似的效能。

邁向大規模搜尋

為了確保正確模擬真實世界的實驗條件,團隊在所有光學元件中都包含了噪聲源、對準誤差和缺陷。他們採取了均勻分佈的隨機錯位,比實驗室中通常遇到的要大得多。

透過對比光相關器系統,團隊進一步證明了方法的穩健性。在複雜情況下,框架保持了良好的成像效能。

圖片

圖 3 :用於大規模發現的通用虛擬光學設定。

在許多其他顯微鏡和成像技術為 XLuminA 的高效和多功能光學模擬器奠定了基礎之後,它們自然而然地得到了擴充套件。尤其是鑑於 XLuminA 框架的模組化特性可以輕鬆實現額外的光學元件和物理功能。

此外,XLuminA 已經為擴充套件到複雜的量子光學顯微鏡技術或其他量子成像技術領域,因為光量子只不過是電磁場模式的激發。

展望未來,可以預期物質波束可在同一框架中模擬,這可能是基於 AI 的顯微鏡技術設計。歸根結底,團隊希望他們的工作所邁出的這一步,是將迄今為止未探索的物理學不同領域的概念引入顯微鏡應用。

原文連結:https://www.nature.com/articles/s41467-024-54696-y

程式碼連結:https://github.com/artificial-scientist-lab/XLuminA
文中圖片涉及來自論文或網路。


相關文章