[论文解读] Equivariant Architectures for Learning in Deep Weight Spaces

Aviv Navon, Aviv Shamsian+4 more|arXiv (Cornell University)|Jan 30, 2023

Machine Learning in Materials Science被引用 5

一句话总结

本文提出了深度权重空间网络（DWSNets），一种新颖的神经网络架构，通过在权重和偏置矩阵上强制执行自然排列对称性，直接从预训练的MLP中处理其权重。该方法通过池化、广播和全连接操作表征仿射等变层，实现参数高效的训练，并在泛化误差预测和INR分类等任务中达到最先进性能。

ABSTRACT

Designing machine learning architectures for processing neural networks in their raw weight matrix form is a newly introduced research direction. Unfortunately, the unique symmetry structure of deep weight spaces makes this design very challenging. If successful, such architectures would be capable of performing a wide range of intriguing tasks, from adapting a pre-trained network to a new domain to editing objects represented as functions (INRs or NeRFs). As a first step towards this goal, we present here a novel network architecture for learning in deep weight spaces. It takes as input a concatenation of weights and biases of a pre-trained MLP and processes it using a composition of layers that are equivariant to the natural permutation symmetry of the MLP's weights: Changing the order of neurons in intermediate layers of the MLP does not affect the function it represents. We provide a full characterization of all affine equivariant and invariant layers for these symmetries and show how these layers can be implemented using three basic operations: pooling, broadcasting, and fully connected layers applied to the input in an appropriate manner. We demonstrate the effectiveness of our architecture and its advantages over natural baselines in a variety of learning tasks.

研究动机与目标

• 设计一种能够直接以原始矩阵形式处理神经网络权重的深度学习架构。
• 解决在具有复杂对称性结构的深度权重空间中学习的挑战，这些结构会阻碍标准架构的性能。
• 开发一个尊重MLP权重矩阵中固有排列对称性的系统性框架。
• 实现在权重空间上的下游任务，如模型适配、剪枝和表征学习。

提出的方法

• 该架构基于DWS层——一种尊重MLP权重排列对称性的仿射等变线性层。
• 这些层被表征为在输入网络特定权重和偏置子空间之间映射的分块矩阵。
• 这些层通过三种核心操作实现：池化、广播和标准全连接层。
• 网络由一系列DWS层堆叠而成，配合逐点非线性激活，形成用于权重空间处理的深层架构。
• 该方法确保等变性：对输入MLP中的神经元进行排列，将导致输出产生相应且可预测的变换。
• 该架构具有参数效率，由于结构化的权重共享，其参数量远少于全连接基线模型。

实验结果

研究问题

• RQ1 对MLP权重矩阵排列对称性保持等变的线性层，其数学约束是什么？
• RQ2 如何利用基本深度学习操作高效实现此类等变层？
• RQ3 DWSNet能否有效学习预训练MLP权重空间上的函数，例如预测泛化误差或学习密集表征？
• RQ4 在涉及原始权重矩阵的任务中，该架构与标准基线相比表现如何？
• RQ5 DWSNets在其他架构（如INRs或NeRFs）上的泛化能力如何？

主要发现

• DWSNets在从权重预测模型性能方面实现了最低的泛化误差，在CIFAR-10和MNIST数据集上均优于所有基线模型。
• 该模型学习到了INRs的解耦2D密集表征，能按频率和振幅对模型进行分组，展现出良好的可解释性。
• 同时使用权重到权重和偏置到偏置的模块，性能优于仅使用其一的情况，其中对角线变体达到了接近最优的结果。
• 数据增强和批量归一化显著提升了性能，在MNIST INRs分类任务中，测试准确率从77.20%提升至85.71%。
• 该架构表现出强大的表达能力，能够通过其学习层近似输入MLP的前向传播过程。
• 尽管在剪枝和INR处理方面面临挑战，但通过适当的初始化和领域感知正则化，该方法展现出巨大潜力。