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 an interesting tweak to the network architecture to accelerate CoortMLP. The idea is to split the input coordinates along the dimensions and then share weights before fusion.  The authors analyze the theoretical upper bound (as far as the MAC ops are concerned) and show about 2X speedup on actual machines.