ICLR 2024 Oral｜用巧妙的「傳送」技巧，讓神經網路的訓練更加高效 机器之心 發表於2024-05-09
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com 本論文作者趙博是加州大學聖地亞哥分校的三年級在讀博士，其導師為 Rose Yu。她的主要研究方向為神經網路引數空間中的對稱性，及其對最佳化、泛化和損失函式地貌的影響。她曾獲 DeepMind 獎學金，並且是高通創新獎學金的決賽入圍者。郵箱：bozhao@ucsd.edu 眾多神經網路模型中都會有一個有趣的現象：不同的引數值可以得到相同的損失值。這種現象可以透過引數空間對稱性來解釋，即某些引數的變換不會影響損失函式的結果。基於這一發現，傳送演算法（teleportation）被設計出來，它利用這些對稱變換來加速尋找最優引數的過程。儘管傳送演算法在實踐中表現出了加速最佳化的潛力，但其背後的確切機制尚不清楚。 近日，來自加州大學聖地亞哥分校、Flatiron Institute、美國東北大學等機構的研究人員釋出的《Improving Convergence and Generalization Using Parameter Symmetries》論文中，展示了傳送不僅能在短期內加快最佳化過程，而且能總體上縮短收斂時間。此外，研究還發現，傳送到具有不同曲率的最小值可以改善模型的泛化能力，這意味著最小值的曲率與泛化能力之間存在聯絡。研究者們進一步將傳送技術與多種最佳化演算法以及基於最佳化的元學習相結合，這些結果充分展示了在最佳化過程中融入引數空間對稱性的潛力。 論文連結：https://openreview.net/forum?id=L0r0GphlIL 程式碼連結：https://github.com/Rose-STL-Lab/Teleportation-Optimization 作者主頁：https://b-zhao.github.io/ 引數空間對稱性（parameter space symmetry）是群 G 在引數空間（Param）上的一個作用，該作用使得損失函式 L 的值保持不變： 對稱傳送演算法（symmetry teleportation）在損失函式水平集中尋找更陡峭的點以加速梯度下降： 傳送後，損失值不會改變。然而，梯度和之後的訓練軌跡會有所不同。 透過傳送，隨機梯度下降（SGD）會收斂到一個包含多個靜止點的盆地，其中透過傳送可達的每個點都是靜止的。圖 1 中的定理提供了損失函式梯度期望值的上限。相比之下，普通 SGD 只能保證存在一個點 wt，使得梯度最終會很小。加入傳送後，對於群作用軌道上所有的點，梯度都將很小。 此外，在進行一次傳送後，梯度下降的方向與牛頓法（Newton’s method）方向相同。因此，收斂速度具有一個二次收縮項，這是二階方法的典型特徵。相反，如果在相同的假設下省略傳送步驟，演算法的收斂速度將僅呈現線性收斂。圖 2 中的定理展示了傳送導致的線性和二次收斂的混合。 圖 2： 該定理展示了傳送導致的線性和二次收斂的混合 在探索機器學習模型最佳化的過程中，「傳送」這一概念最初被提出是為了加速收斂並提高演算法的效率。然而，在該論文的第二部分，研究者們將視野擴充套件到了一個新的目標 —— 提升模型的泛化能力。 泛化能力通常與模型在訓練過程中達到的極小值的「銳度」(sharpness) 相關。為了深入理解這一點，研究者們引入了一個新的概念 ——「極小值曲率」(curvature)，並討論其對泛化的影響。透過觀察極小值的銳度、曲率與泛化能力之間的關聯，研究者們提出了一種新的方法，將銳度和曲率納入傳送的目標中來提升模型的泛化效能。 圖 3 透過視覺化的方式展示了一個梯度流 L (w) 和一條極小值上的曲線（γ），這兩條曲線的曲率對應著極小值的銳度和曲率。此外，表中還顯示了測試集上的損失與銳度或曲率之間的 Pearson 相關性。在三個資料集中，銳度與驗證損失呈強正相關，而極小值的曲率則與驗證損失呈負相關。這些發現表明，具有較小銳度或較大麴率的極小值，可能會帶來更好的泛化效果。 圖 3： 展示了一個梯度流 L (w) 和一條極小值上的曲線（γ） 這些相關性的一種解釋來源於損失地形（loss landscape）在不同資料分佈上的變化。當資料分佈發生變化導致損失地形變化時，尖銳的極小值損失增加較大（如圖 4 右側所示）。在圖 4 中，曲率較大的極小值與變化後的極小值距離更遠（如圖 4 左側所示）。 取得了極小值的屬性與泛化能力的相關性後，該研究人員將引數傳送到具有不同銳度和曲率的區域，以提高模型的泛化能力。圖 5 為在 CIFAR-10 上 SGD 的訓練曲線，其中在第 20 個 epoch 進行了一次傳送。實線代表平均測試損失，虛線代表平均訓練損失。 傳送到更平坦的點略微改善了驗證損失。透過傳送改變曲率對泛化能力有更顯著的影響。傳送到曲率較大的點有助於找到具有較低驗證損失的極小值，而傳送到曲率較小的點則產生相反的效果。這表明至少在區域性，曲率與泛化相關。 傳送不僅有助於 SGD 的收斂速度。為了展示傳送與其他標準最佳化器的良好相容性，研究者們使用不同的最佳化器在 MNIST 上訓練了一個三層神經網路，並進行了帶傳送和不帶傳送的訓練。如圖 6 所示，在使用 AdaGrad、帶動量的 SGD、RMSProp 和 Adam 時，傳送提高了這些演算法的收斂速度。 圖 6： 在使用 AdaGrad、帶動量的 SGD、RMSProp 和 Adam 時，傳送提高了這些演算法的收斂速度 受條件程式設計（conditional programming）和基於最佳化的元學習（meta-learning）的啟發，研究者們還提出了一個元最佳化器，用於學習在損失水平集中移動引數的位置。這種方法避免了在群流形上最佳化的計算成本，並改進了現有的僅限於區域性更新的元學習方法。 圖 7 總結了訓練演算法。研究人員使用了兩個 LSTM 元最佳化器 (m1, m2) 來學習區域性更新 ft 和傳送中使用的群元素 gt。在兩層 LeakyReLU 神經網路上的實驗中，他們使用了不同初始化建立的多個梯度下降軌跡上訓練元最佳化器，並在訓練中未見過的新初始值上測試收斂速度。 與基線相比，同時學習兩種更新型別（LSTM (update,tele)）比單獨學習它們（LSTM (update) 、LSTM (lr,tele)）能夠實現更好的收斂速率。 該論文的主要貢獻包括傳送加速 SGD 收斂的理論保證、極小值曲率的量化及其與泛化相關性的證據、基於傳送的改善泛化的演算法以及學習傳送位置的元學習演算法。 傳送廣泛適用於在損失水平集中搜尋具有優秀屬性的引數。對稱性與最佳化之間的密切關係開啟了許多激動人心的機會。探索傳送中的其他目標是一個有潛力的未來方向。其他可能的應用包括將傳送擴充套件到不同的架構，如卷積或圖神經網路，以及不同的演算法，如基於取樣的最佳化。