【NeuIPS 2021】《Meta-learning sparse implicit neural representations》

---

Code： https://github.com/jaeho-lee/MetaSparseINR
Author： KAIST

---

摘要

目前的INRs方法很难被扩展到用于大量的信号或数据集（也就是无法泛化？），因为每个INR的参数都很heavy，且需要很大的内存和计算。为此，这篇文章提出使用元学习的方法，结合稀疏约束下的网络压缩，这样它呈现一个初始化良好的稀疏参数化，在随后的训练中可以快速演化为一组未见过的信号。这篇文章证明，当使用相同数量的优化步骤进行训练以适合每个信号时，meta-learned sparse neural representations比具有相同数量参数的dense meta-learned models获得的损失要小得多。

一、动机

\quad 最近关于INR，也被称为coordinate-based neural representations，的研究逐渐增多。这种方式其实就是信号表征的一种新范式，它用一个具有连续激活函数的神经网络来近似coordinate-to-value的映射。因为激活函数是连续的，所以INR能够为信号带来连续的表征。INRs的连续性带来了很多好处。
\quad 尽管有很多优点，但这种network-as-a-representation的方法很难扩展到处理大量的信号集，因为为每个信号训练一个parameter-heavy的神经网络需要很多内存和计算。为了解决这样的问题，现在的方法可以分为三类。

• 第一类方法，利用一个跨信号共享的神经网络结构，它以一个latent code vector作为输入，并相应地修改INR输出，如《Learning continuous image representation with local implicit image function》、《DeepSDF》等。虽然这些方法只需要为每个信号存储一个latent vector（以及共享的模型参数），但它们在适应learned latent space之外的信号方面的能力有限，这可能会导致对新信号的表示能力下降（这样的问题其实与目前泛化性NeRF的condition NeRF方式一样，到时候可以在我们文章的related work中分析一下）。
• 第二类方法，利用Meta-learning的方法训练一个initial INR，使每个信号可以在少量优化步骤下被训练拟合，如《MetaSDF》、《Learned initializations for optimizing coordinate-based neural representations》等。但不幸的是，这些方法没有减少存储每个signal的内存需求。
• 第三类方法，减少INR模型参数需要的内存，通过统一量化每个INR的权重，如《COIN: Compression with implicit neural representations》。但是，这个方法并没有减少为每个信号进行训练所需的优化步骤数，也不利用来自其他信号的任何信息来改进量化。

\quad 从这个意义上说，这三种方法在提供memory- and compute-effificient的方法来训练可见和看不见信号的INR方面都有自己的局限性。
\quad 为了解决这些限制，这篇文章基于neural network pruning建立了一个新的框架，即，通过将参数固定为零，从神经表示中删除参数（比如《Skeletonization: A technique for trimming the fat from a network via relevance assessment》）。在这个框架下，目标是训练一组 sparse INRs，以 compute-efficient manner 表示一组不同的信号（以前看到或看不到的）。贡献如下：

• 将上述问题描述为 finding a sparse initial INR，它可以在少量的优化步骤内被训练来拟合每个信号。这样的方式消除了对每个信号单独修剪INR的需求，通过考虑这样的公式，我们消除了对每个信号单独修剪inr的需要，每个信号通常需要一个计算量重的修剪-再训练周期（如图1b）。
• 为了解决这个新定义的问题，作者提出了一种名为Meta-SparseINR的裁剪算法。Meta-SparseINR交替应用（1）a meta-learning step来学习一个INR，和（2）a pruning step使用元训练的权重大小作为显著性得分/saliency score，从INR中删除部分幸存的参数（如图1c）。

这篇文章使用了一个被广泛使用的INR架构，SEREN（《Implicit neural representations with periodic activation functions.》）。此外，作者还提供了一个探索性实验的结果，该实验阐明了在INR上下文中拟合多个信号的初始化时修剪方案的潜在局限性（到时候我们使用时，可以对此进行轻微的改进）。

二、相关工作

\quad Fitting implicit neural representations for multiple signals：目前，至少有三种不同的方法来有效地训练INR用于大量信号集。

• 第一种方法是在信号之间共享一个神经网络组件，它使用一个潜在代码（即从每个信号生成的低维向量）来修改INR输出。《Implicit neural representations with periodic activation functions》训练一个超网络从latent code中预测INR参数，利用从信号集学习到的知识来更好地拟合每个信号。《Learning continuous image representation with local implicit image function》、《Occupancy networks》、《DeepSDF》构建了一个INR架构，以latent code和坐标向量的串联作为输入，以预测每个信号对应的信号值。该方法在内存和计算方面都有优点：在参数方面，每个信号的内存代价等于潜在向量的大小，每个信号的训练代价是共享网络的总训练代价除以样本数。然而，正如《Learn Init.》中所指出的那样，这些方法在表达训练阶段没有观察到的信号时有局限性。（这一类方法与现在的条件性泛化性NeRF相同，所以这个局限性对条件性泛化性NeRF是同样存在的）
• 第二种方法是meta-learn初始INR，从中每个信号可以通过少量的优化步骤被有效地拟合。《MetaSDF》提出了一种基于该方法的学习二维/三维形状对应SDF的算法，结果表明，这种基于元学习的预训练显著减少了编码新信号时所需的优化步骤数。《Learn Init.》提出了一种类似的算法，可以很好地应用于更广泛的表示任务，包括二维图像回归。我们注意到，尽管它在计算上有好处，但这种方法在参数效率方面并没有提供任何好处。
• 第三种方法是在训练后压缩每个INR。《COIN》提出了一种基于权重量化的方法：他们首先通过广泛的训练拟合每个信号。然后，他们将权重从32位精度量化到16位精度。这种方法在计算方面没有任何优势，但有可能与其他方法相结合。
  本文所采用的方法与第二类工作《MetaSDF》、《Learn Init.》最相关。事实上，我们的主要基线之一是《Learn Init.》，但它被用于dense but narrower INRs。

\quad Training from sparse initial models：最近的工作《SNIP: Single-shot network pruning based on connection sensitivity》、《The lottery ticket hypothesis: finding sparse, trainable neural net works》、《Rethinking the value of network pruning》表明，可以找到一个随机初始化模型的稀疏子网络（即被修剪的网络），当从头开始训练时，可以接近达到原始密集模型的性能水平。这一发现引起了人们的兴趣，来开发算法可以通过使用各种概念的显著性分数（如《Pruning neural networks without any data by iterative conserving synaptic flow》、《Picking winning tickets before training by preserving gradient flow》）来mask原始模型来找到这样的“winning ticket” subnetwork。虽然研究方向总体上取得了很大的进展，但在多个数据集时如何找到最佳子网的研究还较少。一个密切相关的工作是《One ticket to win them all: generalizing lottery ticket initializations across datasets and optimizers》，该论文表明，从一个数据集发现的winning ticket可以转移到一个高度相关的数据集。这篇文章发现，在INR背景下，这种通过剪枝随机初始化的模型生成的模型可能不能显著优于密集和狭窄的基线。所以这篇文章与先前工作关注的区别在于：作者不将自己限制在通过修剪一个随机初始化的模型而生成的模型上，而是修剪一个经过元训练的模型。

\quad Meta-learning & network pruning：最早的工作考虑结合元学习和神经网络裁剪的是《MetaPruning: Meta-learning for automatic neural network channel pruning》。在这篇论文中，目标是训练一个PruningNet，他生成一组channel-pruned architecture和相应的权重，以对给定的输入满足逐层稀疏性的要求。最近的一个工作是《Meta-learning with network pruning》，与这篇文章有点像。《Meta-learning with network pruning》的目标是提升few-shot learning的泛化属性，没有考虑内存和计算的限制；提出的算法也可以选择性地从元学习的模型中去除和恢复连接，像在DSD或SDNN那样。

三、A framework for efficiently learning sparse neural representations

\quad 本文的目的是 finding a well-initialized sparse subnetwork structure。这样每个信号可以直接从这个稀疏子网络被训练，无需额外的修剪，像《Learning both weights and connections for effificient neural networks.》那样对每个信号都进行一个 computation-heavy prune-retrain cycle。
\quad Formalisms。
\quad 对于传统的INR来讲，可以用一个MLP：$f(\cdot ;\theta )$表示，对单个信号而言，损失函数如下：

\quad 如果不考虑稀疏性的话，可以为每个信号单独优化，也就是如下，被分为N个子问题。在这种情况下，总参数量是 dN。

\quad Sparse subnetwork and initialization。
\quad 为了减少参数的数量，作者考虑寻找一个最多具有 k<d 连接的 共享稀疏子网结构，从而将参数的总数减少到kN。更正式一点，自网络结构可以被表示为一个二进制掩码 M∈{0，1}${}^d$，其中$||M|{|}_0<=k$，即M中为1的数量小于等于K。M被应用于神经表征中，通过逐元素地与参数进行点积，目的就是把某些参数屏蔽掉。其实这样的一个优化过程，就是找到一个 0-1 mask vector M，执行下面的优化：

\quad 然而，为了便于实际考虑，上面这个问题需要进一步改进。来让这个等式中的每个子问题被解决时没有太多的梯度计算，我们也需要一个良好的初始化。实际上，众所周知，如果没有一个well-chosen初始化，稀疏网络的优化通常需要很长（《《Understanding the effects of data parallelism and sparsity on neural network training》》），甚至无法达到well-chosen参数可以达到的性能水平（《The lottery ticket hypothesis: finding sparse, trainable neural net works》、《A signal propagation perspective for pruning neural networks at initialization》）。
\quad 这里对上面的式子进行了一些扩充。首先，将

\quad 表示从初始化参数${\theta }^0$开始优化M个步骤后，被M进行mask的参数${\theta }^j$。在给定的t个优化步骤的预算下，我们想要找到一个mask M∈{0，1}${}^d$和一个初始化值${\theta }^0\in T^d$，来联合解决下面这个式子：

\quad Before we move on，我们注意到，如果没有架构更改——通过简单地汇集来自N个信号的数据对并应用标准的剪枝技术，这个问题是无法直接解决的。这是因为我们考虑的是信号表示的任务，而不是一个分类。在signal representation setup中，来自不同信号的data pairs共享coordinate input space {$x_i$}_{$i\in I$}。如果来自不同信号的数据对被合并，这将导致对相同输入（坐标）具有不同输出值（信号）的样本对发生冲突，即，训练数据中，相同的输入，具有不同的输出，跟泛化性NeRF一个道理。

四、Learning sparse neural representations with meta-learning

\quad τ是外循环优化步骤数量，t是内循环优化步骤数量，β是outer step size，大T是数据集，θ是初始参数，M是mask vector。

\quad 优化算法如下：

\quad 算法的流程是（1）首先对于一个signal set大T，运行τ个优化步骤，meta-learn一个INR，然后（2）对于学到的INR，按照下降排序的方式去除小的权重，直到达到k的稀疏性，（3）如果没达到，则继续运行（2）。
\quad 正如我们将通过下面实验来验证的那样，元学习是Meta-SparseINR中的一个重要组成部分。更具体地说，元学习步骤至少扮演了双重作用：首先，元学习给出的权重可以被训练，以更少的优化步骤拟合每个信号。这一事实可能并不令人惊讶，事实上，在没有考虑任何修剪的情况下，这一点已经被《Learn Init.》验证了。这篇文章观察到的是，这一优势对于稀疏的INRs也同样适用，这将在下面的实验中阐述。
\quad 第二点是，这篇文章发现：“通过元学习学习到的权值的大小可以作为修剪步骤的一个强大的显著性分数”。这正好印证了，对于一个 signal family，他们具有固有的结构规则（就像3D人脸的一个deformable模板一样），而meta INR的内部权重的各个部分其实attend着特定的component。所以，在对某个信号进行裁剪后进行t个步骤的重新训练后，该部分的权重仍然会很大。分析地，这是因为在一类signal的INR的权重空间中，可能不同的signals的INR的权重会稍微有一些不同，但是这些signals的INR的权重会有一个交集，这个交集中的权重无论在哪个signals INR中都是会占有一定比例的。换句话说，在这个交集中的权重，是这类signals的核心部分，如果我们meta-learn到的权重只保留这一部分，就可以即实现memory-efficient，又可以实现adaption-efficient。

五、实验

\quad 这篇文章关注2D image regression tasks。
\quad Models。使用带SIREN激活函数的MLP。
\quad Method & baselines：

• Meta-SparseINR (ours): 在每个裁剪阶段，裁剪20%的权重参数，即，γ=0.2。
• Random Pruning：Meta-SparseINR架构相同，但是在裁剪时使用均匀概率进行随机修剪，而不是使用magnitude criterion。
• Dense-Narrow：像在《Learn Init.》那样，meta-learn一个dense神经表示，它的宽度比原始的INR更窄。选择宽度是为了确保每个稀疏的INR都有一个dense-narrow counterpart，它的参数量与总参数量相同或更大。
• MAML+OneShot：对于meta-trained的模型，执行一次修剪，然后运行50次训练。
• MAML+IMP：执行t次修剪
• Scratch：与Dense-Narrow相同，但是没有用meta-learn权重，而是用标准初始化的权重。
  

六、讨论和限制

\quad 这篇文章采样神经网络裁剪的方法，关注(1)生成稀疏参数和(2)没有任何共享参数可用。虽然这些特性对某些应用程序可能是有利的（例如，更小的推理成本），但这些限制对于提高参数效率并不是必要的。例如，如果为了提升参数效率，可以考虑用a dense, shared reference model，然后在此基础上产生一个具有稀疏差异的INR来拟合信号；随着信号数量的增加，存储共享模型的内存成本将变得可以忽略不计。探索这种算法的潜力是一个很有前途的未来方向。