編輯丨&
擴充套件一直是提高機器學習各個領域的模型效能和泛化的關鍵因素。儘管在擴充套件其他型別的機器學習模型方面取得了成功，但對神經網路原子間電位 （NNIP） 擴充套件的研究仍然有限。
該領域的主要正規化是將許多物理域約束合併到模型中，例如旋轉等方差等對稱性約束。來自美國加州大學伯克利分校和勞倫斯伯克利國家實驗室 （Berkeley Lab） 的研究團隊認為，這些日益複雜的領域約束抑制了 NNIP 的擴充套件能力，從長遠來看，這種策略可能會導致模型效能停滯不前。
為了系統地研究 NNIP 擴充套件屬性和策略，團隊提出了一種全新的專為可擴充套件性而設計的 NNIP 架構：高效縮放注意力原子間勢 （EScAIP），意在透過注意力機制擴充套件模型有效提高模型表達能力。
該方法以「The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains」為題，於 2024 年 10 月 31 日釋出於 Arxiv 預印平臺。
與現有的 NNIP 模型相比，EScAIP 因使用高度最佳化的注意力 GPU 核心實現，效率大幅提高，推理時間至少加快 10 倍，記憶體使用量減少 5 倍。
團隊強調，他們的方法應被視為一種哲學而非特定模型。它代表了開發通用 NNIP 的概念驗證，這些 NNIP 透過擴充套件實現更好的表達性，並透過增加計算資源和訓練資料繼續有效擴充套件。
關於神經網路通用進化的思考
近年來，擴充套件模型大小、資料和計算的原理已成為提高機器學習 （ML） 效能和泛化的關鍵因素，涵蓋從自然語言處理 （NLP）到計算機視覺 （CV）。ML 中的擴充套件在很大程度上取決於是否能最好地利用 GPU 計算能力。這通常涉及如何有效地將模型大小增加到較大的引數量級，以及最佳化模型訓練和推理以實現最佳計算效率。
與這類發展並行的方向，涵蓋了原子模擬，解決藥物設計、催化、材料等方面的問題。其中，機器學習原子間勢，尤其是神經網路原子間勢 （NNIP），作為密度泛函論等計算密集型量子力學計算的替代模型而廣受歡迎。
NNIP 旨在高效、準確地預測分子系統的能量和力，允許在難以用密度泛函理論直接模擬的系統上執行幾何弛豫或分子動力學等下游任務。
當前的 NNIP 主要基於圖神經網路 （GNN）。該領域的許多有效模型越來越多地嘗試將受物理啟發的約束嵌入到模型中，這些約束包括將預定義的對稱性（例如旋轉等方差）合併到 NN 架構中，以及使用複雜的輸入特徵集。
團隊認為，這些日益複雜的領域約束抑制了 ML 模型的擴充套件能力，並且可能會隨著時間的推移在模型效能方面趨於平穩。隨著模型規模的增加，可以假設施加這些約束會阻礙有效表示的學習，限制模型的泛化能力，並阻礙有效的最佳化。其中許多功能工程方法並未針對 GPU 上的高效並行化進行最佳化，進一步限制了它們的可擴充套件性和效率。
基於這些理解，他們開發了高效縮放注意力原子電位 （EScAIP）。該模型在各種化學應用中實現了最佳效能，包括在 Open Catalyst 2020、Open Catalyst 2022、SPICE 分子和材料專案 （MPTrj） 資料集上的最佳效能。
EScAIP 可以很好地與計算一起擴充套件，並且其設計方式將隨著 GPU 計算的不斷進步而進一步提高效率。
模型的種子與未來的枝芽
這些模型通常經過訓練，可以根據系統屬性（包括原子序數和位置）預測系統能量和每原子力。模型分為兩類：基於組表示節點特徵的模型，以及基於笛卡爾座標表示的節點特徵的模型。
NNIP 領域也越來越關注使用量子力學模擬生成更大的資料集，並使用它來訓練模型。有一種趨勢是將物理啟發的約束納入 NNIP 模型架構，例如所有將對稱約束合併到模型中的組。然而，還有其他工作路線並沒有試圖直接在 NNIP 中構建對稱性，而是嘗試“近似”對稱性。
透過消融研究，該團隊系統地研究了縮放神經網路原子間勢 （NNIP） 模型的策略。在確認了高階對稱性（旋轉階數 𝐿）對擴充套件效率的影響後，他們也得出了增加模型引數的最佳方法。
提高具有組表示特徵的 NNIP 模型能力的一種流行方法是增加表示的順序。如果沒有控制模型中可訓練引數的總數，會在模擬中引入差異，這可能混淆 𝐿 對模型的效能影響。
為了明確增加的影響𝐿在模型效能上，並確定在 NNIP 模型中增加引數的最有效策略，團隊將不同值的可訓練引數數量標準化，並系統地新增進模型。
比較分析揭示了不同引數擴充套件策略的效能增益具有明顯的層次結構。一旦跨模型的引數數量得到控制，許多模型與原始模型就會產生相當的誤差。增加注意力機制的引數是最有利的，並且比簡單地在所有元件中新增更多引數有著更實質性的突破。
EScAIP
為了避免冗雜的張量積，團隊對旋轉和平移不變的標量特徵進行調整，以利用自然語言處理中最佳化的自我注意機制。這種處理使該模型比 EquiformerV2 等等變群表示模型的時間和記憶體效率更高。
透過引導模型進行精細的力大小預測，模型可以學習更好的系統表示，從而幫助它更準確地預測系統能量。力的大小是原子的區域性屬性，而能量預測是分子系統的全域性屬性。這可能解釋了為何節點讀出特徵預測力的大小對於能量預測很有幫助。
甜蜜的教訓與輕快的步伐
團隊注意到，他們的調查路線遵循了一些苦澀教訓的原則。專注於擴充套件和計算的策略往往優於那些試圖將領域知識嵌入模型的策略。但是他們卻相信這個教訓是甜蜜的，因其使大規模培訓大眾化，並使其可供更廣泛的社群使用。
原子系統比嵌入到模型中的特定領域資訊要複雜得多。預定義的對稱約束和手工製作的特徵僅提供了這種複雜性的簡單表示。後續改進最好伴隨著相關的評估指標，使 NNIP 能夠透過縮放獲得表達能力來學習其餘資訊。
隨著資料集的不斷增長，在小型資料集上從頭開始訓練模型可能變得沒有必要。雖然約束在非常小的資料制度中可能會提供一些幫助，但也可以利用預先訓練的大型模型的表示作為在較小資料集上進行微調的起點。
除了專注於資料生成之外，其他技術在 NNIP 領域可能會越來越重要。其中包括模型蒸餾、通用訓練和推理策略，這些策略與模型無關，可以應用於任何 NNIP，以及更好地與實驗結果聯絡起來的方法。這表明更全面的策略對於 NNIP 的準確性與效用極為重要。
原文連結：https://arxiv.org/abs/2410.24169
相關程式碼：https://github.com/ASK-Berkeley/EScAIP