CoordX: Accelerating Implicit Neural Representation with a Split MLP Architecture

Implicit neural representations with multi-layer perceptrons ( MLPs ) have recently gained prominence for a wide variety of tasks such as novel view synthesis and 3D object representation and rendering . However , a significant challenge with these representations is that both training and inference with an MLP over a large number of input coordinates to learn and represent an image , video , or 3D object , require large amounts of computation and incur long processing times . In this work , we aim to accelerate inference and training of coordinate-based MLPs for implicit neural representations by proposing a new split MLP architecture , CoordX . With CoordX , the initial layers are split to learn each dimension of the input coordinates separately . The intermediate features are then fused by the last layers to generate the learned signal at the corresponding coordinate point . This significantly reduces the amount of computation required and leads to large speedups in training and inference , while achieving similar accuracy as the baseline MLP . This approach thus aims at first learning functions that are a decomposition of the original signal and then fusing them to generate the learned signal . Our proposed architecture can be generally used for many implicit neural representation tasks with no additional memory overheads . We demonstrate a speedup of up to 2.92x compared to the baseline model for image , video , and 3D shape representation and rendering tasks . 1 INTRODUCTION . Traditional discrete representations , including volumes or point clouds , have long been used for tasks such as 3D object representation and novel view synthesis . Until recently , implicit neural representations have emerged as a popular alternative approach to those tasks . These implicit neural representations leverage multilayer perceptrons ( MLPs ) to represent images , videos , 3D objects , etc ( Mildenhall et al. , 2020 ; Sitzmann et al. , 2020 ; Tancik et al. , 2020 ) . These MLPs ( coordinatebased MLPs , coordMLP for short ) take low dimensional coordinates ( e.g. , coordinate position or time index ) as input and predict the value of a learned signal ( e.g. , color , depth , shape , etc . ) of each coordinate point . The coordinate-based MLP learns an implicit , continuous , and differentiable function that maps input coordinates to various signals . There are several benefits to using an implicit neural representation instead of a traditional discrete representation . First , it is potentially more memory efficient . Unlike the discrete representation where the granularity of the represented object is limited by the grid resolution , the granularity of an implicit neural representation can be freely adjusted by choosing different resolutions of input coordinates while not exceeding the representation capacity of the learned MLP . Second , the fully differentiable learning process allows the MLP to learn complex signals from sparsely available data points to reconstruct high-quality images or objects . Recent research has demonstrated the effectiveness of coordMLPs for a wide range of signal fitting or learning tasks such as image super-resolution ( Chen et al. , 2021 ) , 3D shape representation ( Park et al. , 2019 ; Mescheder et al. , 2019 ; Gropp et al. , 2020 ) , novel view synthesis ( Mildenhall et al. , 2020 ; Martin-Brualla et al. , 2021 ; Schwarz et al. , 2020 ) and photo-realistic 3D scene editing ( Niemeyer & Geiger , 2021 ) . The search for more accurate and generalizable implicit neural network architectures and methodologies is an active area of research . However , a key challenge of implicit neural representation is the significant computation demand for both inference and training . For inference , the MLP has to be queried with each positional point to generate the learned signal at the corresponding position . Using one GTX1080 GPU , rendering a 1024 × 1024 sized image on 5-layer coordMLP takes 0.2 seconds , and rendering an image of size 400×400 from a novel view with the NeRF model takes 18 seconds . Both tasks perform millions of MLP queries . Similarly , training an MLP for neural representation is a very computation-intensive task as each MLP requires full training passes over a large number of sample points from the original signal . On a high-end RTX GPU , it takes more than 10 minutes to learn a 512× 512 image using a 5-layer coordMLP and more than 4 hours to reconstruct a low-resolution scene with NeRF . Several recent works aim to improve the speed of inference and/or training for coordinate-based MLPs for novel view synthesis tasks ( Liu et al. , 2020 ; Neff et al. , 2021 ; Garbin et al. , 2021 ; Yu et al. , 2021 ) or object representation ( Takikawa et al. , 2021 ) . These works leverage spatial sparsity inherent in images , videos , and objects to reduce the amount of computation required for inference or training . For example , NSVF ( Liu et al. , 2020 ) constructs a sparse discrete representation using a voxel octree . An implicit neural representation is then learned over the discrete set of voxels instead of the infinite set of coordinates . However , there are several drawbacks to these approaches . First , some of the mentioned approaches are only applicable to specific tasks , e.g. , novel view synthesis ( Neff et al. , 2021 ) . Second , the additional data structures that are used to explicitly save representations to avoid computation may lead to high memory overheads , e.g. , Garbin et al . ( 2021 ) . Third , the effectiveness of these approaches is limited by the sparsity of the signal being learned and they may be less effective for highly dense and feature-rich images/objects . Our goal , in this work , is to accelerate both training and inference for a wide range of coordMLPbased representation tasks . To this end , we leverage the following observation . The inputs to these MLPs are typically multi-dimensional coordinate points , i.e. , an ( x , y ) coordinate for an image , ( x , y , t ) for a video , etc. , and in existing MLPs , each coordinate is processed independently , as depicted in Fig . 1 ( d ) . However , this misses the implicit locality among points in an image that have similar coordinate values , for example , the neighboring points along the same axis . In this work , we propose a new architecture where the initial layers are split to independently process each dimension of the input . Thus , each dimension is treated independently rather than each coordinate point itself . The outputs of these initial layers are then fused to generate a final prediction based on all input dimensions , as depicted in Fig . 1 ( e ) . We refer to this split architecture as CoordX . The major benefits of our approach are : First , it enables significant speedups in both training and inference by essentially reducing the dimensions of the input . For example , the overall inputs of an H×W image fed into the MLP is reduced fromH×W toH+W . At the same time , by leveraging the locality between points with similar coordinate values , we are able to achieve accuracy close to the baseline architecture . Second , this technique can be applied orthogonally to previously-proposed optimization techniques that leverage sparsity without incurring much additional memory overheads . Third , the dimension-based splitting approach can be applied to a wide range of tasks that use coordinate-based MLPs as we demonstrate in this work . The key challenges in developing a split architecture are : ( i ) ensuring that the overall size of the model does not increase even when processing each dimension of the input with a separate layer ; and ( ii ) effectively fusing features from the split layers to retain the original model accuracy . In this work , we handle these challenges by ( i ) splitting only the first FC layer and then sharing the remaining FC layers across all input branches ; and ( ii ) fusing features from split branches using an outer product operation . We demonstrate that CoordX can significantly accelerate training and inference for various signal fitting tasks1 including image , video , and 3D shape fitting . Our analysis shows a speedup of∼2.5x using a 5-layer CoordX model for various tasks . We demonstrate a 2.35x speedup in inference time for a 1k resolution image and a 2.78x speedup for a 1k resolution video clip in our experiments . To summarize , in this work , we make the following contributions : • We develop a new architecture for coordinate-based MLPs that leverages locality between input coordinate points to reduce the amount of computation required and significantly speed up inference and training . 1CoordX however , does not improve the performance when fitting unidimensional signals ( e.g. , audio ) . • We design the split MLP architecture that achieves faster training and inference and retains the same level of accuracy as the baseline models with comparable or slightly larger parameter counts . • We demonstrate CoordX ’ s effectiveness in speeding up inference and training on various signal fitting tasks , including image , video and 3D shape rendering , and other advanced applications such as image super-resolution . 2 RELATED WORK . Coordinate-based MLPs for implicit neural representation . Recent works have demonstrated the potential of MLPs in implicitly representing signals such as images ( Stanley , 2007 ) , volume occupancy ( Peng et al. , 2020 ; Mescheder et al. , 2019 ) , signed distance ( SDF ) ( Park et al. , 2019 ; Xu et al. , 2019 ) , texture ( Saito et al. , 2019 ; Henzler et al. , 2020 ) , and radiance fields ( NeRF ) ( Mildenhall et al. , 2020 ) . Unlike MLPs used for other tasks , an MLP used for implicit representation needs to be able to capture the high frequency variations in the color or geometry of images and objects . To address this , Mildenhall et al . ( 2020 ) proposes a positional encoding ( PE ) method that maps the original input coordinates into a high dimensional space using a series of high frequency functions to better capture high frequency variations . Sitzmann et al . ( 2020 ) use sinusoidal activation functions ( SIREN ) to replace the ReLU activations in the MLP , and Tancik et al . ( 2020 ) use simple Fourier feature mappings of input coordinates to learn high frequency variations . Many recent works also aim to make the signals represented by coordinate-based MLPs editable or generalizable . A common way to enrich the representation power of these models is to add latent codes/features as additional inputs to the coordMLPs ( Schwarz et al. , 2020 ; Park et al. , 2021 ; Chen et al. , 2021 ; Mehta et al. , 2021 ) . These additional inputs are combined with positional coordinates and are essentially treated as additional dimensions in the input coordinates . CoordX is directly applicable to these works . Accelerating MLPs for implicit neural representation . Several recent works aim to accelerate inference and/or training for implicit neural representation using MLPs . There are broadly two directions among existing works : ( i ) reducing the number of points to be queried ; ( ii ) reducing the time spent in processing each point . Reducing the number of query points is a common optimization method for ray-casting-based 3D neural representations . Due to spatial sparsity in a 3D scene , a significant number of sampled points along cast rays are empty , without useful density or color information . The inference/training speed can be greatly improved by skipping the processing of these points . Neff et al . ( 2021 ) use an additional depth network to estimate object depth along each ray , hence only points around the object surface are sampled . Similarly , traditional explicit 3D representation methods can be leveraged to guide point sampling . Garbin et al . ( 2021 ) and Kellnhofer et al . ( 2021 ) explicitly extract the mesh of a rendered object after training to reduce the number of points that need to be processed by the MLP . Liu et al . ( 2020 ) ; Yu et al . ( 2021 ) ; Takikawa et al . ( 2021 ) construct an octree structure to represent the 3D object during training that enables skipping empty spaces . To reduce the computation done for each point during inference , Reiser et al . ( 2021 ) replace the larger MLP for the whole scene with many tiny MLPs that have fewer layers and parameters , each of which only represents a small region of the whole scene . Garbin et al . ( 2021 ) and Yu et al . ( 2021 ) also pre-compute points and store results in memory . Therefore , only lightweight interpolation rather than intensive MLP computation is performed for each query point during rendering . To alleviate the significant memory requirements and processing time to pre-compute the original NeRF MLP ’ s 5D input grid ( 3D position + 2D direction ) , these works leverage spherical harmonics ( Ramamoorthi & Hanrahan , 2001 ) to generate a more efficient 3D representation . Concurrent work ACORN ( Martel et al. , 2021 ) uses a multi-scale block-coordinate decomposition to efficiently represent large-scale images or 3D shapes in a hierarchical way . Unlike our coordinate decomposition , ACORN uses a quadtree or octree to decompose the input coordinate grid into small patches . ACORN first gets patch-level features via a large coordinate-based MLP , and then uses these features to get predicted signals for each coordinate via feature interpolation and a lightweight MLP . Thus the number of queries to the large coordMLP in ACORN is significantly reduced . Though both ACORN and CoordX reduce the number of queries to some parts of the coordinate-based model , they use very different approaches and these two approaches are orthogonal . CoordX can be combined with ACORN to generate additional speedups by replacing ACORN ’ s coordinate encoder with the CoordX model .

Summary: The paper proposes a novel coordinate-based network architecture which proposes to process each of the input coordinates independently in the first layer instead of together in a fully connected layer. This input style results in a speed-up in terms of evaluation of the network, and thus faster training and inference in tasks where coordinate-based MLPs are used, without incurring a significant degradation in terms of the quality of the signal fit. These benefits are demonstrated for the tasks of image representation, video representation, and 3D shape representation.