AI探索宇宙結(jié)構(gòu)新突破!超精準(zhǔn)場(chǎng)級(jí)模擬,半秒完成冷暗物質(zhì)仿真
下面的兩組動(dòng)圖展示了計(jì)算機(jī)對(duì)于宇宙形成的模擬:
其中一行是根據(jù)物理定律計(jì)算的,而另一行是由人工智能學(xué)習(xí)后生成的。
你能看出來(lái)哪一個(gè)結(jié)果出自AI之手嗎?
這就是來(lái)自馬克斯·普朗克研究所等機(jī)構(gòu)發(fā)表的工作:利用宇宙學(xué)和紅移依賴性對(duì)宇宙結(jié)構(gòu)形成進(jìn)行場(chǎng)級(jí)仿真。
AI終于開(kāi)始模擬宇宙了!LeCun也在第一時(shí)間轉(zhuǎn)發(fā)和推薦:
作者表示:我們現(xiàn)在可以使用人工智能在多個(gè)時(shí)間步上對(duì)大多數(shù)宇宙學(xué)特性生成宇宙的冷暗物質(zhì)模擬。
并且,物理學(xué)是這個(gè)神經(jīng)網(wǎng)絡(luò)設(shè)計(jì)的核心,它可以看成是PINN(內(nèi)嵌物理知識(shí)神經(jīng)網(wǎng)絡(luò))的一種實(shí)現(xiàn),其損失函數(shù)建模了時(shí)間相關(guān)的粒子坐標(biāo)和速度之間的特定關(guān)系。
在這項(xiàng)工作中,研究人員提出了一個(gè)用于大規(guī)模結(jié)構(gòu)的場(chǎng)級(jí)模擬器,捕獲宇宙學(xué)依賴性和宇宙結(jié)構(gòu)形成的時(shí)間演化。
模擬器將線性位移場(chǎng)映射到特定紅移處對(duì)應(yīng)的非線性位移。
模擬器是一個(gè)神經(jīng)網(wǎng)絡(luò),包含對(duì)Ω和紅移z處的線性增長(zhǎng)因子D(z)的依賴性進(jìn)行編碼的樣式參數(shù)。
研究人員在六維N-body相空間上訓(xùn)練模型,將粒子速度預(yù)測(cè)為模型位移輸出的時(shí)間導(dǎo)數(shù),顯著提高了訓(xùn)練效率和模型準(zhǔn)確性。
最終,模擬器在測(cè)試數(shù)據(jù)(訓(xùn)練期間未見(jiàn)過(guò)的各種宇宙學(xué)和紅移)上實(shí)現(xiàn)了良好的精度和性能,在z = 0,k ~ 1 Mpc/h的尺度上達(dá)到了百分比級(jí)精度,并在較高紅移下提高了性能。
通過(guò)合并樹(shù)將預(yù)測(cè)的結(jié)構(gòu)形成歷史與N體模擬進(jìn)行比較,可以找到一致的合并事件序列和統(tǒng)計(jì)特性。
并且,該模擬器速度極快,在單個(gè)GPU上半秒內(nèi)就能夠預(yù)測(cè)128的立方個(gè)粒子的非線性位移和速度場(chǎng)。
同時(shí)又可以通過(guò)多GPU并行處理進(jìn)行良好的擴(kuò)展,支持任意大尺寸的實(shí)現(xiàn)。
模擬宇宙的AI
隨著宇宙學(xué)數(shù)據(jù)分析推向更小的尺度,利用高階統(tǒng)計(jì)數(shù)據(jù),并實(shí)現(xiàn)場(chǎng)級(jí)分析和基于模擬的推理方案,為預(yù)測(cè)非線性宇宙結(jié)構(gòu)形成提供了更準(zhǔn)確的方法?!?dāng)然也對(duì)算力提出了超高要求。
比如利用N點(diǎn)統(tǒng)計(jì)的傳統(tǒng)分析方法,需要大量模擬數(shù)據(jù)集來(lái)進(jìn)行準(zhǔn)確的協(xié)方差估計(jì)。
而基于模擬的推理方法和場(chǎng)級(jí)分析,則需要生成許多后期密度場(chǎng)的準(zhǔn)確實(shí)現(xiàn),以約束模型參數(shù)和初始條件重建。
DESI、Euclid、Vera C. Rubin天文臺(tái)、SPHEREx和Subaru Prime Focus Spectrograph可以為研究者提供大量最新的星系巡天數(shù)據(jù)。
為了探明宇宙學(xué)參數(shù)和初始條件的最佳約束,需要對(duì)巡天觀測(cè)值進(jìn)行快速、高度準(zhǔn)確的預(yù)測(cè)。
在這項(xiàng)工作中,作者通過(guò)添加紅移依賴性和對(duì)多個(gè)紅移模擬快照的訓(xùn)練來(lái)擴(kuò)展場(chǎng)級(jí)N體模擬器。
由于本模型的時(shí)間依賴性和自可微性,研究者可以有效地獲得N體粒子速度作為輸出粒子位移的時(shí)間導(dǎo)數(shù)。
可以在訓(xùn)練期間動(dòng)態(tài)評(píng)估這些速度,由此定義一個(gè)取決于粒子位置和速度的損失函數(shù),在六維N體相空間上進(jìn)行訓(xùn)練。
強(qiáng)制執(zhí)行「速度必須等于位移時(shí)間導(dǎo)數(shù)」的物理約束,可以提高訓(xùn)練效率并提高模型的準(zhǔn)確性,特別是對(duì)于速度場(chǎng)。
模型結(jié)構(gòu)
作者通過(guò)周期模擬框中的坐標(biāo)x來(lái)描述N體粒子。每個(gè)粒子都與規(guī)則立方晶格上的一個(gè)位點(diǎn)q相關(guān)聯(lián),因此它在紅移z處的位置定義為:
這里Ψ是位移場(chǎng),q是粒子的拉格朗日坐標(biāo)。在線性Zeldovich近似 (ZA) 中,位移場(chǎng)演變?yōu)椋?/span>
其中 D(z) 是線性增長(zhǎng)因子,zi是早期選擇的紅移,以便線性理論可以很好地描述位移場(chǎng)。
隨著引力簇的非線性在后期變得重要,這種線性近似變得不準(zhǔn)確,而模擬宇宙結(jié)構(gòu)形成的非微擾方法(如N體模擬)變得必要:
這里f(z)為線性增長(zhǎng)率,H(z) 是哈勃率(Hubble rate),使用粒子速度來(lái)模擬星系探測(cè)中的紅移空間扭曲。
研究人員設(shè)計(jì)了場(chǎng)級(jí)模擬器,根據(jù)目標(biāo)紅移處的ZA位移場(chǎng)來(lái)預(yù)測(cè)z = 3-0范圍內(nèi)任何紅移處的非線性粒子位移和速度。
場(chǎng)級(jí)模擬器采用U-Net/V-Net設(shè)計(jì),使用PyTorch的map2map庫(kù)實(shí)現(xiàn)和訓(xùn)練模型。
如上圖所示,模型的輸入具有三個(gè)通道,對(duì)應(yīng)于所需紅移處ZA位移的笛卡爾分量,排列在3D網(wǎng)格中。
輸入經(jīng)過(guò)四個(gè)ResNet 3×3×3卷積,第一個(gè)卷積運(yùn)算將3個(gè)輸入通道轉(zhuǎn)換為64個(gè)內(nèi)部通道。
在四次卷積操作之后,結(jié)果的副本被存儲(chǔ)以供網(wǎng)絡(luò)的上采樣端使用,然后使用2×2×2卷積核對(duì)結(jié)果進(jìn)行下采樣。
該架構(gòu)的感受野對(duì)應(yīng)于給定焦點(diǎn)單元兩側(cè)的48個(gè)網(wǎng)格點(diǎn)。預(yù)測(cè)單個(gè)粒子的位移時(shí),以焦點(diǎn)粒子為中心的大小為97的區(qū)域需要通過(guò)網(wǎng)絡(luò),對(duì)應(yīng)于拉格朗日體積189.45 Mpc/h。
不過(guò),由于網(wǎng)絡(luò)缺少填充區(qū)域中所有單元的信息,因此這些區(qū)域的粒子位移并不準(zhǔn)確,需要從輸出中移除。
網(wǎng)絡(luò)有限的感受野也有一個(gè)優(yōu)點(diǎn):它在線性理論準(zhǔn)確的大尺度上保留了ZA場(chǎng)。
以上的操作可以針對(duì)固定宇宙學(xué)的單個(gè)紅移來(lái)訓(xùn)練來(lái)自模擬快照的數(shù)據(jù)。
為了擴(kuò)展網(wǎng)絡(luò)功能,允許網(wǎng)絡(luò)學(xué)習(xí)N體映射作為Ω和紅移的函數(shù),作者對(duì)其進(jìn)行了增強(qiáng)以包含樣式參數(shù) 。
在執(zhí)行任何卷積(包括下采樣/上采樣操作)之前,快照的Ω和D(z)值將傳遞并映射到與卷積核尺寸匹配的內(nèi)部數(shù)組,然后使用這些參數(shù)調(diào)節(jié)網(wǎng)絡(luò)權(quán)重。
模型訓(xùn)練
訓(xùn)練數(shù)據(jù)
研究人員從一組具有不同宇宙學(xué)參數(shù)和一組固定快照紅移的模擬中隨機(jī)采樣快照,同時(shí)訓(xùn)練樣式參數(shù)和網(wǎng)絡(luò)參數(shù),使用Quijote Latin超立方體模擬,在邊長(zhǎng)1 Gpc/h的空間中使用512個(gè)粒子運(yùn)行。
所有這些模擬的拉格朗日空間分辨率均為1.95 Mpc/h,整個(gè)數(shù)據(jù)集包含2000個(gè)模擬,每個(gè)模擬都有一組獨(dú)特的五個(gè)ΛCDM宇宙學(xué)參數(shù)Ω。
研究人員將2000個(gè)模擬分為三組:1874個(gè)用于訓(xùn)練,122個(gè)用于驗(yàn)證,4個(gè)用于測(cè)試。為了鼓勵(lì)各向同性,這里使用數(shù)據(jù)增強(qiáng),通過(guò)立方體的對(duì)稱性隨機(jī)變換輸入和目標(biāo)數(shù)據(jù)。
損失函數(shù)
模型訓(xùn)練使用的損失函數(shù)包含四個(gè)項(xiàng)。第一個(gè)是粒子位移的平均平方誤差(MSE),比較粒子的模擬器位移預(yù)測(cè)和真實(shí)的N體位移:
第二項(xiàng)是歐拉密度的MSE:
第三項(xiàng)是采用粒子速度的MSE:
最后一項(xiàng)對(duì)應(yīng)于歐拉動(dòng)量場(chǎng)的MSE損失。這里將粒子速度分布到與初始拉格朗日網(wǎng)格具有相同分辨率的歐拉矢量場(chǎng)網(wǎng)格并計(jì)算,p是網(wǎng)格單元中每粒子質(zhì)量的歐拉動(dòng)量。
于是,紅移z處快照的最終損失函數(shù)為:
實(shí)驗(yàn)結(jié)果
研究人員通過(guò)構(gòu)建模擬器輸出和N體模擬真值的歐拉密度、和動(dòng)量自功率譜和互功率譜來(lái)評(píng)估模擬器的準(zhǔn)確性。
對(duì)于密度場(chǎng),使用CIC插值方案將粒子分布到512網(wǎng)格并估計(jì)功率譜,下表列出了用于測(cè)試本文模擬器的五個(gè)模擬宇宙學(xué)參數(shù):
模擬器在訓(xùn)練過(guò)程中從未遇到過(guò)以上5種測(cè)試模擬中的任何一種。
下面測(cè)試在訓(xùn)練數(shù)據(jù)中的五個(gè)固定紅移之間進(jìn)行插值時(shí)的模型性能:
上圖顯示了SNN模擬器的功率譜誤差(比例函數(shù)),每條曲線的顏色表示紅移。
歐拉密度誤差僅源于粒子位移的誤差,當(dāng)紅移z = 0時(shí),新的瞬態(tài)模型的隨機(jī)性與原始模型的隨機(jī)性相當(dāng),并且傳遞函數(shù)誤差通常比原始模型有所改善。
在上圖的最右列中,可以看到由于模擬器無(wú)法完美預(yù)測(cè)BAO幅度而導(dǎo)致的振蕩誤差。不過(guò)模擬器的這些錯(cuò)誤特征低于1%,并且可能會(huì)隨著更多的訓(xùn)練數(shù)據(jù)而得到改善。
紅移相關(guān)模型的性能與z = 0時(shí)的真實(shí)空間密度統(tǒng)計(jì)數(shù)據(jù)相當(dāng),并且在較高紅移時(shí)對(duì)于紅移空間和真實(shí)空間統(tǒng)計(jì)數(shù)據(jù)的性能明顯更好。
隨著紅移的減小,誤差平滑且單調(diào)地增加。這表明模擬器可以在其訓(xùn)練數(shù)據(jù)中的少量固定紅移快照之間有效地進(jìn)行插值,而不會(huì)過(guò)度擬合,否則我們會(huì)在中間看到錯(cuò)誤的振蕩特征。