英伟达神奇研究：用别的模型权重训练神经网络，改神经元不影响输出！

DNN 已经可以这么玩了？

不论计算机视觉还是 NLP，深度神经网络（DNN）是如今我们完成机器学习任务的首选方法。在基于此构建的模型中，我们都需要对模型权重执行某种变换，但执行该过程的最佳方法是什么？

最近，英伟达对其 ICML 2023 研究 DWSNet 进行了解读。DWSNet 展示了一种全新架构，它可以有效处理神经网络的权重，同时保持其排列对称性的等方差。

根据这种方法，我们可以根据其他网络的权重来训练一个网络，这也许是一个用来做持续学习的好方法。同样有趣的是，基于 DWSNet 的探索，我们发现网络权重具有排列对称性 —— 这意味着可以更改神经元的顺序而不更改输出。

论文链接：https://arxiv.org/abs/2301.12780

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

这种方法被认为具有广泛潜力，可以实现各种有趣的任务，例如使预先训练的网络适应新的领域。或许我们可以训练一个从另一个网络提取、编辑或删除信息的网络。

而在生成模型上，我们或许可以在很多针对各种任务进行训练的网络上进行训练，然后在运行时为特定任务生成一个网络 —— 就像现代版本的快速权重网络一样。

让我们看看 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 编辑和模型修剪将很有用。