Meta-Learning Sparse Implicit Neural Representations

1 Introduction . An explosively growing line of research on implicit neural representations ( INRs ) —also known as coordinate-based neural representations—studies a new paradigm of signal representation : Instead of storing the signal values corresponding to the coordinate grid ( e.g. , pixels or voxels ) , we train a neural network with continuous activation functions ( e.g. , ReLU , sinusoids ) to approximate the coordinate-to-value mapping [ 31 , 26 , 3 ] . As the activation functions are continuous , the trained INRs give a continuous representation of the signal . The continuity of INRs brings several practical benefits over the conventional discrete representations , such as providing an out-of-the-box method for superresolution or inpainting tasks [ 2 , 37 , 24 , 36 ] , or having its number of parameters not strictly scaling with the spatial dimension and/or spatial resolution of the signal [ 31 , 26 , 3 , 8 , 1 , 34 , 32 ] . Furthermore , as INRs take a form of trainable models , they are readily amenable to the idea of “ learning a prior ” from a set of signals , which can be utilized for various purposes including image generation [ 2 , 37 , 36 ] or view synthesis [ 40 , 36 , 35 , 27 ] . Despite many advantages , this network-as-a-representation approach is difficult to scale to handle a large set of signals , as having a parameter-heavy neural network trained for each signal requires a lot ⇤Equal contributions Code : https : //github.com/jaeho-lee/MetaSparseINR 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of memory and computations . Existing approaches to address this problem can be roughly divided into three categories . The first category of works utilizes a neural network structure shared across the signals , which takes a latent code vector2 as an input and modifies the INR output accordingly ( e.g. , [ 36 , 2 , 31 ] ) . While these methods only need to store a latent vector for each signal ( along with shared model parameters ) , they have a limited ability to adapt to signals falling outside the learned latent space ; this may lead to a degraded representation capability on new signals . The second line of works adopt meta-learning approach to train an initial INR from which each signal can be trained to fit within a small number of optimization steps [ 35 , 40 ] . Unfortunately , these methods do not reduce the memory required to store the neural representations for each signal . The work belonging to the third category [ 4 ] reduces the memory required to store the INR model parameters by uniformly quantizing the weights of each INR . The method , however , does not reduce the number of optimization steps needed to train for each signal , or utilize any information from other signals to improve the quantization . In this sense , all three approaches have their own limitations in providing both memory- and compute-efficient method to train INRs for both seen and unseen signals . To address these limitations , we establish a new framework based on the neural network pruning , i.e. , removing parameters from the neural representation by fixing them to zero ( e.g. , [ 29 ] ) . Under this framework , the objective is to train a set of sparse INRs that represent a diverse set of signals ( either previously seen or unseen ) in a compute-efficient manner . To achieve this goal , we make the following contributions : • We formulate this problem as finding a sparse initial INR , which can be trained to fit each signal within a small number of optimization steps . By considering such formulation , we eliminate the need to prune the INRs separately for every single signal , each of which typically requires a computationally heavy prune-retrain cycle ( Fig . 1b ) . We give a detailed formulation in Section 3 . • To solve this newly formulated problem , we propose a pruning algorithm called Meta-SparseINR , which is the first pruning algorithm designed for the INR setups ( up to our knowledge ) . MetaSparseINR alternately applies ( 1 ) a meta-learning step to learn the INR that can be efficiently trained to fit each signal , and ( 2 ) a pruning step to removes a fraction of surviving parameters from the INR using the meta-trained weight magnitudes as a saliency score ( Fig . 1c ) . See Section 4 for a complete description of the algorithm . In Section 5 , we validate the performance of the proposed Meta-SparseINR scheme on three natural/synthetic image datasets ( CelebA [ 21 ] , Imagenette [ 14 ] , SDF [ 40 ] ) using a widely used INR 2a low dimensional vector generated by forwarding each signal through some encoder , architecture ( SIREN [ 36 ] ) 3 . In particular , we show that , when trained to fit seen/unseen signals with a small number of optimization steps , Meta-SparseINR consistently achieves a better PSNR than the baseline methods using similar number of parameters , which include the meta-trained dense INRs with narrower width [ 40 ] and random pruning . Furthermore , in Section 5.3 , we provide results on an exploratory experiment which sheds a light on the potential limitations of pruning-at-initialization schemes for fitting multiple signals under the INR context . 2 Related works . Instead of giving a general overview on the implicit neural representations and neural network pruning literature , we focus on describing the works that are more directly related to the main ideas of this manuscript . For an overview of two disciplines , we refer the readers to excellent surveys of Tewari et al . [ 41 ] and Hoefler et al . [ 12 ] on implicit neural representations and the neural network pruning , respectively . Fitting implicit neural representations for multiple signals . As briefly discussed in Section 1 , there are at least three different approaches to efficiently train the INRs for a large set of signals . The first approach is to have a neural network component shared across the signals , which uses a latent code ( i.e. , a low-dimensional vector generated from each signal ) to modify the INR output . [ 36 ] trains a hypernetwork to predict INR parameters from the latent code , to utilize the knowledge learned from the signal set to fit each signal better . Other works [ 2 , 26 , 31 ] constructs an INR architecture that takes a concatenation of latent code and coordinate vector as an input , to predict the signal values corresponding to the each signal . This approach has advantages in terms of both memory and computation : in terms of parameters , the per-signal memory cost equal to the size of the latent vector , and the per-signal training cost is total training cost of the shared network , divided by the number of samples . These methods , however , have limitations in expressing the signals that has not been observed during the training phase , as pointed out in [ 40 ] . The second approach is to meta-learn the initial INR from which each signal can be efficiently fit using a small number of gradient descent steps . [ 35 ] proposes an algorithm based on this approach to learn the signed distance function ( SDF ) corresponding to 2D/3D shapes , and shows that such pre-training by metalearning significantly reduces the number of optimization steps required to encode new signals . [ 40 ] proposes a similar algorithm that works well on a wider range of representation tasks , including 2D image regression . We note that , despite its computational benefits , this approach does not provide any benefit in terms of the parameter-efficiency . The third approach is compress each INR after training . [ 4 ] propose one such method based on the weight quantization : They first fit INRs for each signal via extensive training . Then , they quantize the weights from 32-bit precision to 16-bit precision . This method does not provide any advantage in terms of computations , but can potentially be combined with other approaches . The approach taken in the present paper is most related to the second line of works [ 35 , 40 ] ; indeed , one of our main baselines will be the method of [ 40 ] used to train dense but narrower INRs . Training from sparse initial models . Recent studies [ 6 , 18 , 23 ] have shown that one can find a sparse subnetwork ( i.e. , a pruned network ) of a randomly initialized model that can closely achieve the performance level of the original dense model when trained from scratch . This finding has brought a surge of interest toward developing algorithms that can find such “ winning ticket ” subnetworks by masking the original models using various notions of saliency scores ( e.g. , [ 44 , 38 ] ) . Although the research direction has seen a great progress overall , it has been less studied how one can find the best subnetwork when there are multiple datasets at hand . A closely related work is the paper by Morcos et al . [ 28 ] , which shows that the winning ticket discovered from a dataset can be transferred to a highly relevant dataset . In this paper , the aim is to find a sparse initial model that performs well on a set of different tasks ( which may be fitting a set of images/3D shapes in this context ) when trained further on each task . We note , however , that we consider a slightly different setup than the previous works : We do not restrict ourselves to the models generated by pruning a randomly initialized model , and instead prune a meta-trained model . Indeed , our observation in Section 5.3 suggests that the models generated by pruning a randomly initialized model may not be able to significantly outperform dense and narrow baselines , in the INR context . 3In Appendix B , we report additional experimental results another INR architecture called Fourier feature networks ( FFN [ 39 ] ) , consisting of a Fourier encoding layer and fully-connected ReLU layers thereafter . Meta-learning & network pruning . One of the earliest papers considering the combination of meta-learning with neural network pruning is the work by Liu et al . [ 22 ] . In [ 22 ] , the goal is to train a PruningNet which generates a set of channel-pruned architecture and corresponding weights , satisfying the layerwise sparsity requirements given as an input . A more recent paper by Tian et al . [ 43 ] is slightly more related to our setting : In [ 43 ] , the goal is to improve the generalization property of few-shot meta-learning algorithms , without any considerations on the memory or computation limitations ; the proposed algorithm alternatively removes and restores the connections from the meta-learned model , as in DSD [ 10 ] or SDNN [ 15 ] . Unlike these works , the primary aim of this paper is to provide a method to compress a set of neural networks ; we use meta-learning as a tool to generate a shared model that can be compressed and adapted to each signal . Despite its practical importance , up to our knowledge , we are the first to address this problem . 3 A framework for efficiently learning sparse neural representations . The goal of this paper is to develop a memory- and compute-efficient framework for learning and storing the implicit neural representations for a large number signals . We formulate this problem as finding a well-initialized sparse subnetwork structure ; the sparse neural representation for each signal can be directly trained from this sparse subnetwork , without requiring a further pruning by computation-heavy prune-retrain cycle ( see , e.g. , [ 9 ] ) for each individual signal . Formalisms . For a formal description , we introduce some notations : Let T = { T1 , T2 , . . . , TN } be a set of N target signals that we want to approximate . Each signal Tj : Rk ! Rm is a function that maps a k-dimensional coordinate vector to an m-dimensional signal value . For example , a two-dimensional RGB image can be represented by the signal function T : R2 ! R3 with T ( x , y ) = ( R , G , B ) , and the occupancy of a 3D object ( i.e. , whether a point in 3D space in occupied by the object ) can be represented by the coordinate-to-occupancy function T : R3 ! { 0 , 1 } ⇢ R. We assume that we have an access to data pairs { xi , Tj ( xi ) } i2I for each signal j 2 { 1 , . . . , N } , where { xi } i2I denotes the finite set of predefined input coordinate vectors indexed by I . For twodimensional 256⇥256 images , for instance , the input coordinate vectors may be 2D vectors belonging to the grid { 0 , 1/255 , 2/255 , . . . , 254/255 , 1 } 2 . We use the implicit neural representation f ( · ; ✓ ) : Rk ! Rm with a trainable parameter vector ✓ 2 Rd to approximate each signal . For SIRENs [ 36 ] , f ( · ; ✓ ) will be a multi-layer perceptron ( MLP ) with sinusoidal activation functions ; for Fourier feature networks ( FFN [ 39 ] ) , f ( · ; ✓ ) will be an MLP with random Fourier feature encoders and ReLU activation functions . Now , the representation risk of approximating j-th signal with a parameter ✓ can be written as Lj ( ✓ ) = X i2I kf ( xi ; ✓ ) Tj ( xi ) k22 . ( 1 ) Note that , while we are using the squared loss for a simple presentation , any other loss can be used ( e.g. , ` 1 loss ) . Under typical scenarios without any sparsity considerations , one may optimize the parameters independently for each signal . In other words , we solve the optimization problem min ✓1 , ... , ✓N2Rd 1 N NX j=1 Lj ( ✓j ) = 1 N NX j=1 min ✓j2Rd Lj ( ✓j ) ! , ( 2 ) which can be decomposed into N subproblems . In this case , the total number of parameters is dN . Sparse subnetwork and initialization . To reduce the number of parameters , we consider finding a shared sparse subnetwork structure having at most < d connections , so that the total number of parameters is reduced to N . More formally , the subnetwork structure can be expressed as a binary mask M 2 { 0 , 1 } d with kMk0 , which is applied to the neural representation by taking an element-wise product with the parameter , i.e. , f ( x ; M ✓ ) . We hope to find a zero-one mask vector M that solves the optimization min M2 { 0,1 } d kMk0 min ✓1 , ... , ✓N2Rd 1 N NX j=1 Lj ( M ✓j ) = min M2 { 0,1 } d kMk0 1 N NX j=1 min ✓j2Rd Lj ( M ✓j ) ! . ( 3 ) The problem ( 3 ) , however , needs to be refined for practical considerations . To solve each subproblem in Eq . ( 3 ) without too many gradient computations , we also require a good initialization . Indeed , it has been known that without a well-chosen initialization , the optimization of sparse networks often takes significantly longer [ 20 ] , or even fail to achieve the performance level that can be achieved with a well-chosen one [ 6 , 19 ] . To this end , we augment the problem ( 3 ) with the notion of optimized initialization as follows . Let ✓ ( t ) j ( M ; ✓ ( 0 ) ) 2 Rd denote the parameter that has been attained by taking t steps of gradient-based optimization ( e.g. , SGD or Adam ) to fit the target signal Tj , on the subnetwork with an initial value ✓ ( 0 ) and masked by M . Given some steps-to- budget t , we want to find a pair of mask M 2 { 0 , 1 } d and the initial value ✓ ( 0 ) 2 Rd that jointly solves min ✓ ( 0 ) 2Rd min M2 { 0,1 } d kMk0 1 N NX j=1 Lj ⇣ M ✓ ( t ) ( Tj , ✓ ( 0 ) ) ⌘ . ( 4 ) We propose a meta-learning approach to solve the problem ( 4 ) in Section 4 . Before we move on , we note that the problem can not be directly solved—without an architectural change—by simply pooling the data pairs from N signals and applying standard pruning techniques ( e.g. , [ 9 ] ) . This is because we are considering the task of signal representation , instead of a classification ( or similar ) . Under the signal representation setup , the data pairs from different signals share the coordinate input space { xi } i2I . If data pairs from different signals are pooled , this leads to conflicting pairs of samples that have different output values for the same input .

Summary: This paper presents a method of alternately meta-learning the Implicit Neural Representations (INR) for a set of signals using the MAML method, and pruning the INR using the global magnitude-based pruning method for sparsifying the INR. The proposed method is tested on the 2D image regression task in three datasets (CeleA, Imagenette and 2D SDF). In comparisons, three baselines are used (MAML+w/ random pruning,  MAML+w/ dense-narrow, and w/o MAML + dense-narrow). The proposed method obtains better PSNR across different number of optimization steps, thus showing more efficient learning in terms of computation and memory footprint.