Equivariant Architectures for Learning in Deep Weight Spaces
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, Haggai Maron
ICML 2023
cs.LG
2023-01-30
DWSNet 对神经元换序等变,直接读 MLP 权重。MNIST INR 分类 85.71%,对齐基线 58.98%。
两件事都要直接读原始权重。INR 用小 MLP 拟合一张图或一个三维形状,样本就是这份权重。另一件是已有分类器,要改权重跟上新分布,而不是在新数据上重训。
中间层神经元没有顺序。同一个置换去换前一层权重的行、后一层的列,函数不变。这是 Hecht-Nielsen 1990 年的观察。参数可以差很远,映射却相同。拉直之后,普通网络会把这些写法当成不同样本。
此前的模型多半用全连接或 Transformer 吃拉直参数,去猜超参或泛化。INR 上,functa 用 meta-learning 学出低维 modulation,和训练流程绑在一起。INR2Vec 把各层的行收成一个全局集合,跨层换行通常会改函数。置换增强和对齐,是在样本侧补这套对称。
DWSNet 吃进一份 MLP 的全部权重和偏置。对称群是中间各层宽度上对称群的直积,输入维和输出维不参加,因为它们对齐的是数据坐标。第一层权重左乘置换,最后一层右乘,中间层两侧各乘一次。偏置跟着本层神经元走,输出偏置不动。
等变就是:权重重排后,输出做同一组重排。只要一个与顺序无关的标量时,后面接不变读出,再过小 MLP。
定理 5.1 把权重空间到自身的线性等变层全部写成了分块矩阵。各层权重、偏置先看成直和,层就拆成它们两两之间的映射。每个小块只用求和、复制、全连接三种运算,参数各自独立,比拉直后的全连接小得多。
中间层权重由四项相加得到。本层用 Hartford 等人的双集合等变层,行和列各看成一个集合。相邻层用 DeepSets,单项变换加上全体求和。更远的层先求和再乘标量。首尾层有一维不被置换,用全连接处理。不同中间层偏置之间只做一个标量的求和再广播。
点态非线性夹在线性层之间。求和可以换成 max 或注意力,实验用的是线性和。
命题 6.1:带 ReLU 的 DWSNet 能在紧集上一致逼近一次前向,也就是用这份权重去算输入 x。命题 6.2 在 Lipschitz 等假设下说明,定义在这些 MLP 所表示的函数上的映射也能被逼近,因此同一函数的不同参数写法可以落到同一个输出。
输入网络独立训练、种子不同。默认每个函数一份权重。INR2Vec 只比架构,不用其预训练。
图像 INR 的类别就是它拟合的那张图。测试准确率:
| 方法 | MNIST INR | Fashion-MNIST INR |
| MLP | 17.55 ± 0.01 | 19.91 ± 0.47 |
| MLP + 置换增强 | 29.26 ± 0.18 | 22.76 ± 0.13 |
| MLP + 权重对齐 | 58.98 ± 0.52 | 47.79 ± 1.03 |
| INR2Vec 架构 | 23.69 ± 0.10 | 22.33 ± 0.41 |
| Transformer | 26.57 ± 0.18 | 26.97 ± 0.33 |
| DWSNets | 85.71 ± 0.57 | 67.06 ± 0.29 |
MNIST 上拉直的 MLP 是 17.55%,置换增强 29.26%,对齐 58.98%,DWSNet 85.71%。Fashion-MNIST 上对齐 47.79%,DWSNet 67.06%。
频率回归的正文没有逐点 MSE。正弦定义在 [-π, π],频率取自 U(0.5, 10),图 3 显示样本很少时 DWSNet 的测试误差已经低于各基线。
对比学习把 a sin(bx) 嵌到低维,再线性预测 a 和 b。a、b 取自 U(0, 10),视图是标准差 0.2 的高斯噪声,加上概率 0.5 的遮挡。
| 方法 | 预测 a、b 的 MSE |
| MLP | 7.39 ± 0.19 |
| MLP + 置换增强 | 5.65 ± 0.01 |
| MLP + 权重对齐 | 4.47 ± 0.15 |
| INR2Vec 架构 | 3.86 ± 0.32 |
| Transformer | 5.11 ± 0.12 |
| DWSNets | 1.39 ± 0.06 |
域适应里,CIFAR-10 上的输入网络每次只区分随机两类。目标域加上旋转、翻转、高斯噪声和颜色抖动。模型输出残差 Δv,新权重是 v - Δv,一次前向,分类器和图像都没见过。
| 方法 | 测试准确率 |
| 不适应 | 60.92 ± 0.41 |
| MLP | 64.33 ± 0.36 |
| MLP + 置换增强 | 64.69 ± 0.56 |
| MLP + 权重对齐 | 67.66 ± 0.90 |
| INR2Vec 架构 | 65.69 ± 0.41 |
| Transformer | 61.37 ± 0.13 |
| DWSNets | 71.36 ± 0.38 |
不适应已经有 60.92%,Transformer 停在 61.37%。对齐 67.66%,DWSNet 71.36%。分类上那组差距是 85.71% 对 58.98%,这里是 71.36% 对 67.66%。
同一张 Fashion-MNIST 多拟合几份 INR:1 份 67.06%,2 份 70.22%,4 份 70.31%,6 份 73.32%,8 份 74.87%,10 份 75.12%。从 1 份到 10 份大约 8 个点,2 份到 4 份几乎没涨。
神经元顺序一旦被当成特征,每种初始化都像新样本。MNIST 上置换增强只从 17.55% 走到 29.26%,对称没被数据增强补完。对齐更强,但匹配本身难解,网络和数据变大以后不适合先做完再学习。等变层把这步折进前向,参数只是标量和小全连接。
INR 可以直接分类、回归,不用再采样成像素。分类器可以一次前向吐出域适应残差。源模型在实验里是 CIFAR-10 上随机两类的分类器。
MNIST 上 85.71% 对对齐的 58.98%,对称先验在样本就是函数时很值钱。域适应 71.36% 对 67.66%,用这套层去改权重才刚开始。
结论写了三条。层结构绑在一种 MLP 的深度和宽度上,换网络就要重搭。内部块共享参数留到以后,没有验证。一些任务很难训,正文归因于初始化,细节在附录 K.5,没给出失败任务的数字。实现偏繁。
缩放对称被明确排除。ReLU 网络可以一层乘常数、下一层除掉,函数仍然不变,这套群没包含它。命题 6.2 只覆盖函数空间上的映射,任意等变函数能否逼近没有证明。CNN 的通道置换和 Transformer 里 query、key 的配套置换写了,没有实验。
INR 只测了 MNIST 和 Fashion-MNIST。域适应的源模型是随机两类的二分类器,几种腐蚀捆在一起,没有分项结果,也没有大网络。附录里用权重预测泛化,数字没进正文。拿掉预训练之后,INR2Vec 那一栏不能单独用来否定原文。同一张图的 INR 从 1 份加到 10 份涨约 8 个点,份数本身就会改排名。