自主設計實驗,速度提高2.5倍!馬普所AI方法拓展顯微鏡使用邊界
編輯丨&
在人類的聰明才智與創(chuàng)造力的推動之下,光學顯微鏡突破經(jīng)典光衍射極限,在質(zhì)量與數(shù)量上有了顯著的飛躍。然而,在如此廣袤的選擇空間與錯綜復雜的實驗配置組合中,一些高效的設計尚未被發(fā)現(xiàn),并且可能永遠不會采用人工驅(qū)動。
于此,德國馬克斯普朗克光科學研究所(Max Planck Institute for the Science of Light)的研究團隊介紹了一個基于JAX(Python 中的高性能計算庫)開發(fā)的開源計算框架 XLuminA 。其可以快速準確的探索實驗儀器的組配方式,在自己的虛擬光學設計環(huán)境中模擬出來。
通過 JAX 的加速線性代數(shù)編譯器、即時編譯及其無縫集成的自動矢量化、自動微分與 GPU 兼容性提供的計算速度,XLuminA 的加速相較于目前成熟的數(shù)值優(yōu)化方法提高了四個數(shù)量級。
該研究以「Automated discovery of experimental designs in super-resolution microscopy with XLuminA」為題,于2024年12月10日刊登于《Nature Communications》。
在過去的二十年中,幾項突破進一步拓寬了光學顯微鏡在生物科學的應用范圍,多功能且強大的超分辨率(SR)技術(shù)對不僅限于生物學領(lǐng)域也同樣包括化學與材料科學的技術(shù)范疇產(chǎn)生了相當大的影響。
這個研究的初始動機遠遠超過了對已知光學技術(shù)的小規(guī)模優(yōu)化,旨在發(fā)現(xiàn)目前完全未開發(fā)的先進光學顯微鏡的新穎、可用于實驗上的概念。
XLuminA
馬克斯普朗克光科學研究所的研究團隊提出了 XLuminA ,它可以將預設空間內(nèi)的某一處模擬位置轉(zhuǎn)換為我們需要的物理輸出,然后在目標函數(shù)里使用該輸出來描述所需的設計目標。
模擬器可以通過梯度的優(yōu)化技術(shù)直接調(diào)用,也可以用于為基于深度學習的代理模型生成訓練數(shù)據(jù)。用于自動設計與發(fā)現(xiàn)策略的模擬器必須高速可靠與通用,而 XLuminA 的光學模擬器很明顯滿足以上要求。
團隊使用的方法與以前利用 AI 去進行單個光學元件的數(shù)據(jù)驅(qū)動設計的方式截然不同。雖然這些技術(shù)很有影響力,但它們并不意味著要改變實驗方法的原理或光學布局本身。
相比之下,XLuminA 配備了從頭開始模擬、優(yōu)化和自動設計新光學設置和概念的工具。由于模擬中有不同的物理目標,XLuminA 可以在光束傳播中應用不同近似值。相較于受物理技術(shù)啟發(fā)的伴隨法,它極大地加快了仿真器的速度。
有趣的是,伴隨法可以看做是自動微分(團隊所采用)的一個特例,研究員在文章中表示道。
工作流程與性能
XLuminA 允許仿真經(jīng)典光學硬件配置,并支持優(yōu)化和自動發(fā)現(xiàn)未探索的設置設計。該軟件是使用 JAX 開發(fā)的,它通過加速線性代數(shù)編譯器 XLA 實現(xiàn)增速優(yōu)勢,同時無縫集成自動微分框架和自動 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。為了包含自動發(fā)現(xiàn)功能,XLuminA 的光學模擬器和優(yōu)化器由損失函數(shù)捆綁在一起。在這項工作中,團隊采用了一種基于梯度的策略,其中實驗裝置的參數(shù)在最陡的下降方向上迭代調(diào)整。
為了評估數(shù)值與分析方法計算一個梯度的評估時間以及他們在不同分辨率和設備上的收斂時間,團隊準備了梯度下降算法:Broyden-Fletcher-Goldfarb-Shanno (BFGS) 算法。
最后,利用 JAX 的內(nèi)置 autodiff 框架,解析計算損失函數(shù)的梯度。這種方法使優(yōu)化器能夠有效地構(gòu)建內(nèi)部梯度函數(shù),從而大大減少每次迭代的計算時間。
重新發(fā)現(xiàn)數(shù)據(jù)驅(qū)動
團隊的目標是使用 XLuminA 來發(fā)現(xiàn)新的顯微鏡概念。從本質(zhì)上講,發(fā)現(xiàn)新的實驗配置需要一個混合離散-連續(xù)搜索問題。離散方面源于配置光網(wǎng)絡拓撲,而連續(xù)部分與光學元件的設置有關(guān),例如激光功率和分束器反射率。
為了實現(xiàn)純連續(xù)優(yōu)化實驗設計拓撲的自動發(fā)現(xiàn)及其參數(shù)設置,他們采用了大型復雜的光學拓撲來初始化設置。從這里開始,XLuminA 應該能夠提取出人類可能還沒有考慮過的更復雜的解決方案。
令人驚訝的是,XLuminA 找到了一種替代方法,可以將相位奇點壓印到光束上,并在焦平面上產(chǎn)生明顯的縱向分量。在高參數(shù)化與復雜光學系統(tǒng)中的情況下,XLuminA 都成功地發(fā)現(xiàn)了替代光學解決方案,表現(xiàn)出與參考實驗相似的性能。
邁向大規(guī)模搜索
為了確保正確模擬真實世界的實驗條件,團隊在所有光學元件中都包含了噪聲源、對準誤差和缺陷。他們采取了均勻分布的隨機錯位,比實驗室中通常遇到的要大得多。
通過對比光相關(guān)器系統(tǒng),團隊進一步證明了方法的穩(wěn)健性。在復雜情況下,框架保持了良好的成像性能。
在許多其他顯微鏡和成像技術(shù)為 XLuminA 的高效和多功能光學模擬器奠定了基礎之后,它們自然而然地得到了擴展。尤其是鑒于 XLuminA 框架的模塊化特性可以輕松實現(xiàn)額外的光學元件和物理功能。
此外,XLuminA 已經(jīng)為擴展到復雜的量子光學顯微鏡技術(shù)或其他量子成像技術(shù)領(lǐng)域,因為光量子只不過是電磁場模式的激發(fā)。
展望未來,可以預期物質(zhì)波束可在同一框架中模擬,這可能是基于 AI 的顯微鏡技術(shù)設計。歸根結(jié)底,團隊希望他們的工作所邁出的這一步,是將迄今為止未探索的物理學不同領(lǐng)域的概念引入顯微鏡應用。
原文鏈接:https://www.nature.com/articles/s41467-024-54696-y
代碼鏈接:https://github.com/artificial-scientist-lab/XLuminA文中圖片涉及來自論文或網(wǎng)絡。- 免責聲明
- 本文所包含的觀點僅代表作者個人看法,不代表新火種的觀點。在新火種上獲取的所有信息均不應被視為投資建議。新火種對本文可能提及或鏈接的任何項目不表示認可。 交易和投資涉及高風險,讀者在采取與本文內(nèi)容相關(guān)的任何行動之前,請務必進行充分的盡職調(diào)查。最終的決策應該基于您自己的獨立判斷。新火種不對因依賴本文觀點而產(chǎn)生的任何金錢損失負任何責任。