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: Authors propose a method for learning sparse implicit functional representations for signal data. `Their method combines ideas from MAML to learn a good initialization from which a sparse neural network representing the signal can be efficiently estimated. Experiments on a few datasets suggest that the proposed method achieves better PSNR than more naive approaches for the same levels of sparsity. Overall, the paper is in a good shape but still could benefit from a few additions.