DNN 已經可以這么玩了？

不論計算機視覺還是 NLP，深度神經網絡（DNN）是如今我們完成機器學習任務的首選方法。在基于此構建的模型中，我們都需要對模型權重執行某種變換，但執行該過程的最佳方法是什么？

最近，英偉達對其 ICML 2023 研究 DWSNet 進行了解讀。DWSNet 展示了一種全新架構，它可以有效處理神經網絡的權重，同時保持其排列對稱性的等方差。

根據這種方法，我們可以根據其他網絡的權重來訓練一個網絡，這也許是一個用來做持續學習的好方法。同樣有趣的是，基于 DWSNet 的探索，我們發現網絡權重具有排列對稱性 —— 這意味著可以更改神經元的順序而不更改輸出。

論文鏈接：https://arxiv.org/abs/2301.12780

官方 GitHub：https://github.com/AvivNavon/DWSNets

這種方法被認為具有廣泛潛力，可以實現各種有趣的任務，例如使預先訓練的網絡適應新的領域?；蛟S我們可以訓練一個從另一個網絡提取、編輯或刪除信息的網絡。

而在生成模型上，我們或許可以在很多針對各種任務進行訓練的網絡上進行訓練，然后在運行時為特定任務生成一個網絡 —— 就像現代版本的快速權重網絡一樣。

讓我們看看 DWSNet 是怎么做到的：

在使用隱式神經表征（Implicit Neural Representations，INR）或神經輻射場（Neural Radiance Fields，NeRF）表征的 3D 對象數據集時，我們經常需要「編輯」對象以更改其幾何形狀或修復錯誤，例如移除杯子的把手、使車輪更加對稱。然而，使用 INR 和 NeRF 的一個主要挑戰是它們必須在編輯之前先進行渲染，編輯工具依賴于渲染和微調 INR 或 NeRF 參數。

圖 1. 數據專用架構示例。

來自英偉達的研究團隊試圖把神經網絡用作一種處理器，來處理其他神經網絡的權重。

表征深度網絡參數最簡單的方法是將所有權重（和偏置）矢量化為簡單的平面向量，然后應用全連接網絡（多層感知機（MLP））。這種方法可以預測神經網絡的性能。

但這種方法有一個缺點。神經網絡權重空間具有復雜的結構，將 MLP 應用于所有參數的矢量化版本會忽略該結構，進而損害泛化能力。

圖 2. 具有兩個隱藏層（下）的多層感知機（MLP）的權重對稱性（上）。

幾何深度學習（GDL）領域已經針對 MLP 的這個問題進行了廣泛的研究。

在許多情況下，學習任務對于一些變換是不變的。例如，查找點云類別與給網絡提供點的順序無關。但在有些情況下，例如點云分割（point cloud segmentation），點云中的每個點都被分配一個類，那么輸出就會隨著輸入順序的改變而改變。

這種輸出隨著輸入的變換而變換的函數稱為等變函數。對稱感知架構因其有意義的歸納偏置而具有多種優勢，例如它們通常具有更好的樣本復雜性和更少的參數，這些因素可以顯著提高泛化能力。

權重空間的對稱性

那么，哪些變換可以應用于 MLP 的權重，使得 MLP 所表征的底層函數不會改變？

這就涉及到一種特定類型的變換 —— 神經元排列。如圖 2 所示，直觀地講，更改 MLP 某個中間層神經元的順序，函數不會發生改變。此外，每個內部層的重新排序過程可以獨立完成。

MLP 可以使用如下方程組表示：

該架構的權重空間被定義為包含矢量化權重和偏差的所有串聯的（線性）空間。

重要的是，這樣的話，權重空間是（即將定義的）神經網絡的輸入空間。

那么，權重空間的對稱性是什么？對神經元重新排序可以正式建模為將置換矩陣應用于一層的輸出以及將相同的置換矩陣應用于下一層。形式上，可以通過以下等式定義一組新參數：

新的參數集有所不同，但很容易看出這種變換不會改變 MLP 表示的函數。這是因為兩個置換矩陣 P 和 P^t 相互抵消（假設有像 ReLU 這樣的元素激活函數）。

更普遍的，如前所述，不同的排列可以獨立地應用于 MLP 的每一層。這意味著以下更通用的變換集不會改變底層函數。我們將它們視為權重空間的「對稱性」。

在這里，Pi 表示置換矩陣。這一觀察是由 Hecht-Nielsen 于 30 多年前在論文《ON THE ALGEBRAIC STRUCTURE OF FEEDFORWARD NETWORK WEIGHT SPACES》中提出的。類似的變換可以應用于 MLP 的偏差。

構建深度權重空間網絡

文獻中的大多數等變架構都遵循相同的方法：定義一個簡單的等變層，并將架構定義為此類簡單層的組合，它們之間可能具有逐點非線性。

CNN 架構就是這種結構的一個很好的例子。在這種情況下，簡單的等變層執行卷積運算，CNN 被定義為多個卷積的組合。DeepSets 和許多 GNN 架構都遵循類似的方法。有關更多信息，請參閱論文《Weisfeiler and Leman Go Neural： Higher-Order Graph Neural Networks》和《Invariant and Equivariant Graph Networks》。

當目標任務不變時，可以使用 MLP 在等變層之上添加一個不變層，如圖 3 所示。

圖 3：典型的等變架構由幾個簡單的等變層組成，后面是不變層和全連接層。

在論文《Equivariant Architectures for Learning in Deep Weight Spaces》中，英偉達研究者遵循了這個思考。我們的主要目標是為上面定義的權重空間對稱性識別簡單而有效的等變層。不幸的是，表征一般等變函數的空間可能具有挑戰性。與之前的一些研究（例如跨集合交互的深度模型）一樣，我們的目標是表征所有線性等變層的空間。

因此，研究人員開發了一種新的方法來表征線性等變層，該方法基于如下觀察：權重空間 V 是表示每個權重矩陣 V=⊕Wi 的更簡單空間的串聯。（為簡潔起見，省略了偏差術語）。

這一觀察非常重要，因為它可以將任何線性層 L:V→V 寫入塊矩陣，其第 （i，j） 塊是 Wj 和 Wi Lij ： Wj→Wi 之間的線性等變層。塊結構如圖 4 所示。

但我們如何才能找到 Lij 的所有實例呢？論文中列出了所有可能的情況，并表明其中一些層已經在之前的工作中得到了表征。例如，內部層的 Lii 在跨集交互的深度模型中進行了表征。

值得注意的是：在這種情況下，最通用的等變線性層是常見的，僅使用四個參數的深度集層的泛化。對于其他層，新研究提出基于簡單等變操作的參數化，例如池化、broadcasting 和小型全連接層，并表明它們可以表示所有線性等變層。

圖 4 展示了 L 的結構，它是特定權重空間之間的分塊矩陣。每種顏色代表不同類型的圖層。Lii 是紅色的。每個塊將一個特定的權重矩陣映射到另一個權重矩陣。該映射以依賴于網絡中權重矩陣的位置的方式參數化。

圖 4：線性等變層的塊結構。

該層是通過獨立計算每個塊然后對每行的結果求和來實現的。英偉達在論文中涵蓋了一些額外的技術細節，例如處理偏差項和支持多個輸入和輸出功能。

我們將這些層稱為深度權重空間層（DWS 層），并將由它們構建的網絡稱為深度權重空間網絡（DWSNet）。我們在這里關注以 MLP 作為輸入的 DWSNet。

深度權重空間網絡的表達能力

如果我們把假設類限制為簡單等變函數的組合，可能會無意中損害等變網絡的表達能力，這在上面引用的圖神經網絡文獻中得到了廣泛的研究。英偉達的論文表明，DWSNet 可以近似輸入網絡上的前饋操作，這是理解其表達能力的一步。隨后，新研究證明 DWS 網絡可以近似 MLP 函數空間中定義的某些「表現良好」的函數。

實驗

DWSNet 在兩個任務系列中進行評估。首先采用代表數據的輸入網絡，例如 INR。其次，采用代表標準 I/O 映射（例如圖像分類）的輸入網絡。

實驗一：INR 分類

此配置根據 INR 所代表的圖像對 INR 進行分類。具體來說，它涉及訓練 INR 來表示來自 MNIST 和 Fashion-MNIST 的圖像。任務是讓 DWSNet 使用這些 INR 的權重作為輸入來識別圖像內容，例如 MNIST 中的數字。結果表明，英偉達提出的 DWSNet 架構大大優于其他基線。

表 1：INR 分類。INR 的類別由它所代表的圖像定義（平均測試準確度）。

重要的是，將 INR 分類到它們所代表的圖像類別比對底層圖像進行分類更具挑戰性。在 MNIST 圖像上訓練的 MLP 可以實現近乎完美的測試精度。然而，在 MNIST INR 上訓練的 MLP 卻取得了較差的結果。

實驗 2：INR 的自監督學習

這里的目標是將神經網絡（特別是 INR）嵌入到語義一致的低維空間中。這是一項很重要的任務，因為良好的低維表示對于許多下游任務至關重要。

在這里的數據由適合 asin （bx） 形式的正弦波的 INR 組成，其中 a、b 從區間 ［0，10］ 上的均勻分布中采樣。由于數據由這兩個參數控制，因此密集表示應該提取底層結構。

圖 5：使用自監督訓練獲得的輸入 MLP 的 TSNE 嵌入。

類似 SimCLR 的訓練過程和目標用于通過添加高斯噪聲和隨機掩碼來從每個 INR 生成隨機視圖。圖 4 展示了所得空間的 2D TSNE 圖。英偉達提出的 DWSNet 方法很好地捕捉了數據的潛在特征，而與之對比的方法則比較困難。

實驗 3：使預訓練網絡適應新領域

該實驗展示了如何在不重訓練的情況下使預訓練 MLP 適應新的數據分布（零樣本域適應）。給定圖像分類器的輸入權重，任務是將其權重變換為在新圖像分布（目標域）上表現良好的一組新權重。

在測試時，DWSnet 接收一個分類器，并在一次前向傳遞中使其適應新域。CIFAR10 數據集是源域，其損壞版本是目標域（圖 6）。

圖 6：使用 DWSNet 進行領域適應。

結果如表 2 所示。請注意：在測試時，模型應推廣到未見過的圖像分類器以及未見過的圖像。

表 2：使網絡適應新領域。

未來研究方向

英偉達認為，將學習技術應用于深度權重空間的能力提供了許多新的研究方向。首先，尋找有效的數據增強方案來訓練權重空間上的函數有可能會提高 DWSNet 的泛化能力。其次，研究如何將排列對稱性納入其他類型的輸入架構和層，如 skip 連接或歸一化層也是很自然的思考。

最后，將 DWSNet 擴展到現實世界的應用程序，如形變、NeRF 編輯和模型修剪將很有用。