下面的兩組動圖展示了計算機(jī)對于宇宙形成的模擬：

其中一行是根據(jù)物理定律計算的，而另一行是由人工智能學(xué)習(xí)后生成的。

你能看出來哪一個結(jié)果出自AI之手嗎？

這就是來自馬克斯·普朗克研究所等機(jī)構(gòu)發(fā)表的工作：利用宇宙學(xué)和紅移依賴性對宇宙結(jié)構(gòu)形成進(jìn)行場級仿真。

AI終于開始模擬宇宙了！LeCun也在第一時間轉(zhuǎn)發(fā)和推薦：

作者表示：我們現(xiàn)在可以使用人工智能在多個時間步上對大多數(shù)宇宙學(xué)特性生成宇宙的冷暗物質(zhì)模擬。

并且，物理學(xué)是這個神經(jīng)網(wǎng)絡(luò)設(shè)計的核心，它可以看成是PINN（內(nèi)嵌物理知識神經(jīng)網(wǎng)絡(luò)）的一種實現(xiàn)，其損失函數(shù)建模了時間相關(guān)的粒子坐標(biāo)和速度之間的特定關(guān)系。

在這項工作中，研究人員提出了一個用于大規(guī)模結(jié)構(gòu)的場級模擬器，捕獲宇宙學(xué)依賴性和宇宙結(jié)構(gòu)形成的時間演化。

模擬器將線性位移場映射到特定紅移處對應(yīng)的非線性位移。

模擬器是一個神經(jīng)網(wǎng)絡(luò)，包含對Ω和紅移z處的線性增長因子D(z)的依賴性進(jìn)行編碼的樣式參數(shù)。

研究人員在六維N-body相空間上訓(xùn)練模型，將粒子速度預(yù)測為模型位移輸出的時間導(dǎo)數(shù)，顯著提高了訓(xùn)練效率和模型準(zhǔn)確性。

最終，模擬器在測試數(shù)據(jù)（訓(xùn)練期間未見過的各種宇宙學(xué)和紅移）上實現(xiàn)了良好的精度和性能，在z = 0，k ～ 1 Mpc/h的尺度上達(dá)到了百分比級精度，并在較高紅移下提高了性能。

通過合并樹將預(yù)測的結(jié)構(gòu)形成歷史與N體模擬進(jìn)行比較，可以找到一致的合并事件序列和統(tǒng)計特性。

并且，該模擬器速度極快，在單個GPU上半秒內(nèi)就能夠預(yù)測128的立方個粒子的非線性位移和速度場。

同時又可以通過多GPU并行處理進(jìn)行良好的擴(kuò)展，支持任意大尺寸的實現(xiàn)。

模擬宇宙的AI

隨著宇宙學(xué)數(shù)據(jù)分析推向更小的尺度，利用高階統(tǒng)計數(shù)據(jù)，并實現(xiàn)場級分析和基于模擬的推理方案，為預(yù)測非線性宇宙結(jié)構(gòu)形成提供了更準(zhǔn)確的方法?！?dāng)然也對算力提出了超高要求。

比如利用N點(diǎn)統(tǒng)計的傳統(tǒng)分析方法，需要大量模擬數(shù)據(jù)集來進(jìn)行準(zhǔn)確的協(xié)方差估計。

而基于模擬的推理方法和場級分析，則需要生成許多后期密度場的準(zhǔn)確實現(xiàn)，以約束模型參數(shù)和初始條件重建。

DESI、Euclid、Vera C. Rubin天文臺、SPHEREx和Subaru Prime Focus Spectrograph可以為研究者提供大量最新的星系巡天數(shù)據(jù)。

為了探明宇宙學(xué)參數(shù)和初始條件的最佳約束，需要對巡天觀測值進(jìn)行快速、高度準(zhǔn)確的預(yù)測。

在這項工作中，作者通過添加紅移依賴性和對多個紅移模擬快照的訓(xùn)練來擴(kuò)展場級N體模擬器。

由于本模型的時間依賴性和自可微性，研究者可以有效地獲得N體粒子速度作為輸出粒子位移的時間導(dǎo)數(shù)。

可以在訓(xùn)練期間動態(tài)評估這些速度，由此定義一個取決于粒子位置和速度的損失函數(shù)，在六維N體相空間上進(jìn)行訓(xùn)練。

強(qiáng)制執(zhí)行「速度必須等于位移時間導(dǎo)數(shù)」的物理約束，可以提高訓(xùn)練效率并提高模型的準(zhǔn)確性，特別是對于速度場。

模型結(jié)構(gòu)

作者通過周期模擬框中的坐標(biāo)x來描述N體粒子。每個粒子都與規(guī)則立方晶格上的一個位點(diǎn)q相關(guān)聯(lián)，因此它在紅移z處的位置定義為：

這里Ψ是位移場，q是粒子的拉格朗日坐標(biāo)。在線性Zeldovich近似 (ZA) 中，位移場演變?yōu)椋?/span>

其中 D(z) 是線性增長因子，zi是早期選擇的紅移，以便線性理論可以很好地描述位移場。

隨著引力簇的非線性在后期變得重要，這種線性近似變得不準(zhǔn)確，而模擬宇宙結(jié)構(gòu)形成的非微擾方法（如N體模擬）變得必要：

這里f(z)為線性增長率，H(z) 是哈勃率（Hubble rate），使用粒子速度來模擬星系探測中的紅移空間扭曲。

研究人員設(shè)計了場級模擬器，根據(jù)目標(biāo)紅移處的ZA位移場來預(yù)測z = 3-0范圍內(nèi)任何紅移處的非線性粒子位移和速度。

場級模擬器采用U-Net/V-Net設(shè)計，使用PyTorch的map2map庫實現(xiàn)和訓(xùn)練模型。

如上圖所示，模型的輸入具有三個通道，對應(yīng)于所需紅移處ZA位移的笛卡爾分量，排列在3D網(wǎng)格中。

輸入經(jīng)過四個ResNet 3×3×3卷積，第一個卷積運(yùn)算將3個輸入通道轉(zhuǎn)換為64個內(nèi)部通道。

在四次卷積操作之后，結(jié)果的副本被存儲以供網(wǎng)絡(luò)的上采樣端使用，然后使用2×2×2卷積核對結(jié)果進(jìn)行下采樣。

該架構(gòu)的感受野對應(yīng)于給定焦點(diǎn)單元兩側(cè)的48個網(wǎng)格點(diǎn)。預(yù)測單個粒子的位移時，以焦點(diǎn)粒子為中心的大小為97的區(qū)域需要通過網(wǎng)絡(luò)，對應(yīng)于拉格朗日體積189.45 Mpc/h。

不過，由于網(wǎng)絡(luò)缺少填充區(qū)域中所有單元的信息，因此這些區(qū)域的粒子位移并不準(zhǔn)確，需要從輸出中移除。

網(wǎng)絡(luò)有限的感受野也有一個優(yōu)點(diǎn)：它在線性理論準(zhǔn)確的大尺度上保留了ZA場。

以上的操作可以針對固定宇宙學(xué)的單個紅移來訓(xùn)練來自模擬快照的數(shù)據(jù)。

為了擴(kuò)展網(wǎng)絡(luò)功能，允許網(wǎng)絡(luò)學(xué)習(xí)N體映射作為Ω和紅移的函數(shù)，作者對其進(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ī)采樣快照，同時訓(xùn)練樣式參數(shù)和網(wǎng)絡(luò)參數(shù)，使用Quijote Latin超立方體模擬，在邊長1 Gpc/h的空間中使用512個粒子運(yùn)行。

所有這些模擬的拉格朗日空間分辨率均為1.95 Mpc/h，整個數(shù)據(jù)集包含2000個模擬，每個模擬都有一組獨(dú)特的五個ΛCDM宇宙學(xué)參數(shù)Ω。

研究人員將2000個模擬分為三組：1874個用于訓(xùn)練，122個用于驗證，4個用于測試。為了鼓勵各向同性，這里使用數(shù)據(jù)增強(qiáng)，通過立方體的對稱性隨機(jī)變換輸入和目標(biāo)數(shù)據(jù)。

損失函數(shù)

模型訓(xùn)練使用的損失函數(shù)包含四個項。第一個是粒子位移的平均平方誤差（MSE），比較粒子的模擬器位移預(yù)測和真實的N體位移：

第二項是歐拉密度的MSE：

第三項是采用粒子速度的MSE：

最后一項對應(yīng)于歐拉動量場的MSE損失。這里將粒子速度分布到與初始拉格朗日網(wǎng)格具有相同分辨率的歐拉矢量場網(wǎng)格并計算，p是網(wǎng)格單元中每粒子質(zhì)量的歐拉動量。

于是，紅移z處快照的最終損失函數(shù)為：

實驗結(jié)果

研究人員通過構(gòu)建模擬器輸出和N體模擬真值的歐拉密度、和動量自功率譜和互功率譜來評估模擬器的準(zhǔn)確性。

對于密度場，使用CIC插值方案將粒子分布到512網(wǎng)格并估計功率譜，下表列出了用于測試本文模擬器的五個模擬宇宙學(xué)參數(shù)：

模擬器在訓(xùn)練過程中從未遇到過以上5種測試模擬中的任何一種。

下面測試在訓(xùn)練數(shù)據(jù)中的五個固定紅移之間進(jìn)行插值時的模型性能：

上圖顯示了SNN模擬器的功率譜誤差（比例函數(shù)），每條曲線的顏色表示紅移。

歐拉密度誤差僅源于粒子位移的誤差，當(dāng)紅移z = 0時，新的瞬態(tài)模型的隨機(jī)性與原始模型的隨機(jī)性相當(dāng)，并且傳遞函數(shù)誤差通常比原始模型有所改善。

在上圖的最右列中，可以看到由于模擬器無法完美預(yù)測BAO幅度而導(dǎo)致的振蕩誤差。不過模擬器的這些錯誤特征低于1%，并且可能會隨著更多的訓(xùn)練數(shù)據(jù)而得到改善。

紅移相關(guān)模型的性能與z = 0時的真實空間密度統(tǒng)計數(shù)據(jù)相當(dāng)，并且在較高紅移時對于紅移空間和真實空間統(tǒng)計數(shù)據(jù)的性能明顯更好。

隨著紅移的減小，誤差平滑且單調(diào)地增加。這表明模擬器可以在其訓(xùn)練數(shù)據(jù)中的少量固定紅移快照之間有效地進(jìn)行插值，而不會過度擬合，否則我們會在中間看到錯誤的振蕩特征。