$(\textrm{Implicit})^2$: Implicit Layers for Implicit Representations

1 Introduction . The concept of “ implicitness ” has recently been applied to various contexts of machine learning research . One particular thread is implicit representations [ 19 , 22 , 28 , 27 ] , which aims to learn a continuous representation of high-frequency , often discretely-measured signals such as large images . Formally , given a spatial-temporal coordinate input x ∈ Rd , an implicit representation of it is a function Φ : Rd → C , where C is the space of desired quantity ( e.g . color , volume density , distance , etc . ) and Φ is usually parameterized with a neural network . Such implicitness has multiple benefits over the conventional discrete ( e.g. , grid-based ) representations ; e.g. , as Φ is defined on a continuous domain , an implicitly represented input consumes much less storage than the original input ; as another example , Φ is differentiable , hence allowing for the computation of higher-order derivatives . However , the training of these Φ networks themselves are usually quite memory-consuming , since we usually deal with very high-resolution images and videos , and training is typically done in full-batch mode ( e.g. , when training on a batch size of 512× 512 , a simple 4-layer MLP with 1024 hidden units already requires > 16 GB memory just to store the intermediate activations ) . A recent ( orthogonal ) usage of implicitness in deep learning comes from the notion of implicit layers/models , where the word is used to characterize model architectures ( rather than input representations ) . The nomenclature comes from the concept of implicit vs. explicit functions : instead of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . representing a model as an explicit stacking of layers , an implicit model solves a non-linear dynamical system [ 10 ] ( e.g. , ODEs [ 7 , 9 ] or fixed-points [ 1 , 32 ] ) . Importantly , these implicit models decouple the forward and backward passes and analytically differentiate via the implicit function theorem [ 14 ] . This enables these models to use constant training memory , irrespective of the forward pass trajectory . Moreover , recent works on deep equilibrium models [ 1 ] have been shown to be able to achieve results on par with the state-of-the-art explicit models ( e.g. , Transformers [ 31 ] ) on very large-scale vision or NLP tasks [ 1 , 2 ] . However , despite training memory efficiency , these implicit layers are also slower to compute ( e.g. , by 2-3× ) , since the solvers are usually iterative in nature and ultimately require many function evaluations . Despite the success of these “ implicit ” methods in their respective areas , the two concepts so far have rarely intersected : existing models used for implicit representations are neither implicit in nature ( since they are still stacked multi-layer explicit operators trained end-to-end in the feedforward manner ) , nor do they exploit any properties of the implicit functions during training ( e.g. , implicit differentiation ) . In this paper , we argue that implicit representations and implicit layers are remarkably well-suited to each other , and their use together can compensate for their aforementioned drawbacks . Specifically , we propose to combine the best of both worlds by replacing the explicit models used in conventional implicit representation learning with implicit layers , whose output is defined by a fixed-point condition , thus dubbed ( Implicit ) 2 networks . Formally , let x be the input and F be an ( often shallow ) layer , the ( Implicit ) 2 approach learns the following implicit representation Φ : Φ ( x ) = Wz ? + b , where z ? = F ( z ? ; x ) ( 1 ) Importantly , we show that these two “ implicitness ” complement each other well , especially in two important aspects . First , the unique large-batch training scheme of implicit representations as well as the forward/backward decoupling properties of implicit models permit us to amortize the cost of the iterative solver that would otherwise make implicit models slow to train . We show that by fixed-point reuse ( in the forward pass ) and Jacobian approximation ( in the backward pass ) , training an implicit model on implicit representation tasks requires almost no computation overhead . Second , implicit models allow us to train these high-frequency input tasks with a much lower memory consumption while not losing speed or performance , thus significantly lowering the hardware requirement and computation budget of training implicit representations . We evaluate ( Implicit ) 2 on multiple implicit representations for images , videos , and audios , showing substantial improvement upon existing competitive explicit-model methods like SIREN [ 27 ] or MFN [ 11 ] using equivalent-sized networks . 2 Background and Preliminaries . 2.1 Implicit Neural Representations ( INR ) . While high-frequency data such as images and scene geometry have been traditionally represented discretely ( e.g. , pixel/voxel grids or mesh points ) , recent work has demonstrated the possibility of replacing them with continuous functions parameterized by multi-layer perceptrons ( MLPs ) [ 27 , 28 ] . These deep networks have been used to learn differentiable representations of various forms , such as continuous spatial-temporal image representations [ 5 , 8 ] , high-resolution scene representations [ 26 , 22 ] , and volumetric rendering [ 19 , 29 ] , and have been shown to be significantly more compact than the conventional grid-based approaches . However , training these INR models is not easy . For example , the training of implicit representations is usually conducted in very large batch sizes ( e.g. , 512× 512 ) , which renders these ( albeit simple ) MLPs very memory-consuming during training . As another example , recent works have shown that typical deep networks with ReLU/tanh/sigmoid non-linearities fail to capture the fine details of an input signal , and resorted to periodic non-linearities [ 27 ] , Fourier feature encodings [ 28 ] , or repeated application of non-linear filters ( e.g. , Gabor wavelets ) to the input and then multiply with the features [ 11 ] . We introduce below the two latest , highly competitive models that have been developed on this task , which our work will build on ( see §3.1 ) . Sine-activated Network for Implicit Representations ( SIREN ) . To overcome the detrimental effect of traditional non-linearities like ReLU/tanh on modeling fine details and higher-order derivative of the input signals , Sitzmann et al . [ 27 ] proposes to use sinusoidal activation functions that allow explicit supervision on any derivatives of the input signal . The resulting MLPs , though simple , have been shown to achieve state-of-the-art performance in representing images , videos , complex geometries , and more [ 27 ] . Formally , given an input x , SIREN learns an implicit representation by the following L-layer MLP : Φ ( x ) = W ( gL−1 ◦ gL−2 ◦ · · · ◦ g0 ) ( x ) + b ( 2 ) where gi is the i-th layer of the network that computes gi ( h [ i ] ) = sin ( Wih [ i ] + bi ) with a sine non-linearity , and h [ i ] denotes the hidden state at i-th layer . Multiplicative Filter Networks ( MFN ) . More recently , Fathony et al . [ 11 ] proposed to forego the SIREN-like design ( i.e. , MLPs with periodic activation functions ) in favor of multiplicative operations that are proven to be similarly expressive . Specifically , at each layer i of the network , an MFN applies a learnable non-linear filter kernel ( parameterized by θ ( i ) ) g ( x ; θ ( i ) ) on the original input x ∈ Rn , and elementwise multiply it with a linear transformation of the features ; i.e. , formally , Φ ( x ) = Wh [ L ] + b , where h [ i+1 ] = ( Wih [ i ] + bi ) · g ( x ; θ ( i ) ) and h [ 0 ] = g ( x ; θ ( 0 ) ) ( 3 ) Fathony et al . [ 11 ] in particular proposed two instantiations that admit sinusoids or a Gabor wavelet as the filter g , which are called FourierNet and GaborNet , respectively : gFourier ( x ; θ ( i ) Fourier ) = sin ( ω ( i ) x+ φ ( i ) ) gGabor ( x ; θ ( i ) Gabor ) = sin ( ω ( i ) x+ φ ( i ) ) · exp ( − ( γ ( i ) /2 ) · ( x− µ ( i ) ) 2 ) ( 4 ) where θ ( i ) Gabor = { ω ( i ) ∈ Rn , φ ( i ) ∈ R , γ ( i ) ∈ R , µi ∈ Rn } ( similar for θ ( i ) Fourier ) , and ( µ ( i ) , γ ( i ) ) denotes the mean and scale of the Gabor filter . These models have shown performance on par with or better than the periodic-nonlinearity-based networks like SIREN on a range of tasks like image and video representations . 2.2 Implicit Models . Whereas explicit models hierarchically stack a sequence of L non-linear operators to build a computation graph for their forward/backward passes ( where L is often referred to as the depth of a model ) , implicit model instead define the output as a solution to a non-linear dynamical system which in most cases can ’ t be expressed as a fixed computation graph . For example , the Neural ODE model [ 7 ] solves an initial value problem ODE that intuitively corresponds to infinitesimal residual layers [ 13 ] . Another example is the recent deep equilibrium ( DEQ ) models [ 1 ] , whose forward pass involves computing the fixed point of a shallow layer F . Importantly , unlike explicit models , these implicit models do not need to store any intermediate activations during training : to produce gradients in the backward pass , we can directly differentiate through their final outputs ( e.g. , by adjoint method for Neural ODEs [ 7 ] and implicit differentiation for equilibrium models [ 1 ] ) , regardless of how the forward pass was computed . Therefore , the forward and backward passes of implicit models can rely on independent black-box solvers , and these models consume only a constant training memory . Our work is primarily built on the formulation of deep equilibrium models , which we introduce in detail below . Deep Equilibrium ( DEQ ) Models . Given an input x and a layer F ( z ; x ) ( e.g. , a residual block , or a multi-head self-attention layer [ 31 ] ) , a DEQ model aims to approximate an infinite-level stacking of this layer ( i.e . z [ i+1 ] = F ( z [ i ] , x ) with i = 1 , . . . , L and L→∞ ) . To achieve this purpose , Bai et al . [ 1 ] advocated for directly solving a fixed-point condition defined by F and x : z ? = F ( z ? ; x ) ( 5 ) where z ? is called the equilibrium since , intuitively , stacking another F ( · , x ) layer on top does not change the activation . To solve for the equilibrium point , prior works have had success with various kinds of advanced solvers , such as Broyden ’ s method [ 6 , 1 , 17 ] , and Peaceman-Rachford splitting [ 23 , 32 , 24 ] . Then , in the backward pass , we can produce the gradient by the implicit function theorem [ 14 ] directly applied on this equilibrium point z ? , with a loss ` : ∇ ( · ) ` = ∂ ` ∂z ? ( I − JF ) −1 ∂F ( z ? ; x ) ∂ ( · ) ( 6 ) where JF denotes the Jacobian of F evaluated at z ? . As the inverse of a matrix is expensive to compute , in practice the backward pass is usually formulated as another linear fixed-point condition in y : y > JF + ∂ ` ∂z ? = y > , which again can be solved via any of the advanced black-box solvers . Thus theoretically , the model only needs to store x , z ? and F ( z ? ; x ) in memory at training time , which is equivalent to the memory consumption of merely one F layer . Recent works [ 1 , 2 ] have demonstrated that DEQ models can scale to high-dimensional tasks and be as competitive as explicit structures like ResNet [ 13 ] and Transformers [ 31 ] on large-scale computer vision and NLP tasks . However , these implicit models are also slower ( typically by ∼3× ) to train , because both forward and backward passes involve numerous iterative steps . Our work studies the novel combinations of implicit models and implicit representation learning . Specifically , we show that these two “ implicitness ” are surprisingly well-suited to each other , as they almost perfectly compensate for each others ’ drawbacks . We demonstrate that implicit modeling of a simple layer F substantially improves the training speed , memory , and performance on implicit representation tasks when compared to the aforementioned state-of-the-art deep explicit networks . 3 ( Implicit ) 2 Networks As we discussed in §2.1 , learning implicit representation for complex signals is often more effective when projecting the input coordinates onto high-frequency bases ( e.g . sinusoidal functions ) , making the designs of these models typically different from deep networks used in more common applications , like image classifications ( e.g. , ResNets [ 13 ] ) . In this work , we propose to directly build on the designs of these aforementioned state-of-the-art layer architectures by only making minimal modifications , but modeling them as shallow implicit models . We briefly introduce below the two instantiations of our proposed ( Implicit ) 2 approach , which are based on SIREN [ 27 ] and MFN [ 11 ] , respectively , followed by a discussion of how implicit representation tasks can enable us to amortize the cost of training these implicit networks significantly in practice . 3.1 ( Implicit ) 2 Network Architctures Implicit Sine-activated Networks ( iSIREN ) . Motivated by the success of SIREN , we first propose an input-injected variant of SIREN ( denoted as SIREN ( input inj . ) ) suitable for implicit models , i.e. , FSIREN ( z ; x ) = sin ( W ( z + sin ( V x ) ) + Ux+ b ) ( 7 ) where , following prior works on deep equilibrium models [ 1 , 32 , 24 ] , the transformed input x is added to the hidden representation z ( see Fig . 2 ) . Just like the canonical SIREN layer , any order of derivative of this input-injected SIREN variant is also an input-injected SIREN of the same form . Therefore , the proposed variant inherits the same high-frequency and higher-order differentiability properties as SIREN , and we name the implicit model that solves its fixed-point as iSIREN . Implicit Multiplicative Filter Networks ( iMFN ) . Additionally , we introduce a variant based on the recent Multiplicative Filter Networks ( MFN ) [ 11 ] . As MFNs only perform linear transformations on the hidden features , which are multiplied with a non-linear filter function on the original input x , we slightly modify the layer design as follows to introduce an additional input injection : FMFN ( z ; x ) = ( W ( z + g ( x ; θ2 ) ) + b ) ◦ g ( x ; θ1 ) ( 8 ) where g denotes the filter function of choice , such as Gabor or Fourier ( gGabor and gFourier in Eq . ( 4 ) ) filters [ 11 ] . We call the implicit models defined by this layer iMFN . Moreover , just like the original MFN [ 11 ] , we can show that the output of iMFN ( which is the fixed-point of FMFN ) is also a linear combination of the non-linear filter functions ( see Appendix ? ? for proof ) . We provide visualizations of these two architectural instantiations in Fig . 2 . With these variants , we formally introduce the corresponding ( Implicit ) 2 approaches to implicit representation tasks as follows : Φ ( x ) = W ′z ? + b′ , where z ? = F ( z ? ; x ) , and F ∈ { FMFN , FSIREN } ( 9 ) 3.2 Accelerated Training of ( Implicit ) 2 Networks Compared to explicit networks , training with implicit models significantly lowers the memory budget required ( thanks to implicit differentiation ) . This is especially compelling for training implicit representation tasks , which are often memory-bottlenecked . In this subsection , we describe how one can exploit 1 ) the large-batch nature of implicit representation training ; and 2 ) the nature of fixed-point optimizations ( in both forward and backward ) , to significantly improve the speed of these models as well , making them eventually both more memory-efficient and time-efficient than explicit models ( while achieving performance as competitive , or better ; see §4 ) . Fixed-point Reuse . Suppose we fix a training step t and let the corresponding network parameter ( i.e. , weights of the network ) be given by θ ( t ) . Typically , a deep equilibrium model solves for the fixed-point z ? ( t ) given an input x with a black-box fixed-point solver starting with an initial guess z [ 0 ] ( which is usually 0 or a random vector in N ( 0 , I ) ) : z ? ( t ) = FixedPointSolve ( F , θ ( t ) , x , z [ 0 ] ) ( 10 ) where we note that z ? ( t ) does not depend at all on this initial guess z [ 0 ] of the optimization . Depending on the quality of the initial guess , the optimization process itself can take a different number of steps . For example , in the unlikely circumstance where we guessed z [ 0 ] to be z ? ( t ) , the fixed-point solvers will directly converge at iteration 0 . However , in practice , this property is rarely used , as given a complex high-dimensional input x ( e.g. , a high-resolution image ) , it is hard to know ( or learn ) what makes a good initial guess . Therefore , these implicit models ’ aforementioned memory efficiency comes at a cost : without good initial guesses , we must spend more iterative updates in order to secure convergence to a good enough estimate of the fixed point in practice , rendering a slower model . However , we argue that the training on many implicit representation tasks , which uses the entire batch of data ( e.g. , to fit an RGB-color image implicit representation , each pixel of this image is a “ sample ” of this full batch ) , can compensate for this issue by providing us very reasonable initial guesses . In particular , training in the full-batch mode allows us to have direct access to all of the fixed-points z ? ( t ) for all inputs x in the training set immediately after step t. Assuming the model updates between iterations t and t+ 1 are small ( which is true in practice , such as via SGD with a small learning rate η ) , one can safely assume the fixed-points of the updated layer F ( · ; x , θ ( t+1 ) ) do not deviate much from its current estimate z ? ( t ) . Formally , F ( · ; x , θ ( t+1 ) ) → F ( · ; x , θ ( t ) ) as learning rate η → 0 , which implies z ? ( t+1 ) → z ? ( t ) . We empirically verify this on these full-batch training settings ( see Fig . 3 ) , where a majority of the fixed-point convergences ( to a target residual level ε = 0.01 ) finish in exactly 1 step . However , we note that this does not imply all training iterations will gain a similar level of convergence boost ( e.g. , see the first 100 steps in Fig . 3 ) . In particular , at training iteration t = 0 , we do not have a “ previous estimate ” at all , which means we still have to solve for z ? by running the solver for several steps with a neutral initial guess . Thus , we highlight that such fixed-point reuse amortizes the cost of later training iterations , and indeed , the majority of the training of the implicit models . Empirically , we found that performing one step of the fixed-point iteration is sufficient for training the iSIREN and iMFN networks ; i.e. , at training step t+ 1 , we simply perform z ? ( t+1 ) = F ( z ? ( t ) ; x , θ ( t ) ) . ( 11 ) Truncated Backward Gradient . As mentioned in §2.2 , the backward gradient of implicit models ( Eq . ( 6 ) ) can be alternatively obtained by solving a linear fixed-point system y > = y > JF + ∂ ` ∂z∗ ( 12 ) so as to avoid expliciting inverting the matrix . However , recent works have suggested that assuming sufficient stability in the DEQ model layer F , we can also approximate ∂ ` ∂z∗ ( I − JF |z∗ ) −1 by an expansion of it as a Neumann series [ 15 ] : y > = ∂ ` ∂z ? ∞∑ i=0 ( JF ) i ( 13 ) For example , Fung et al . [ 12 ] propose to use the 0- th order approximation and perform a Jacobian-free backward gradient estimation ( assuming Lipschitz F ) ; i.e. , y > ≈ ∂ ` ∂z ? . Although this approximated gradient was empirically shown to be effective in image classification tasks , we found it to substantially hurt the training-time error convergence in implicit representation learning . As a result , we propose to use a truncated backward gradient such that the Neumann series is unrolled for a T > 0 steps to better approximate the exact implicit differentiation . Formally , the proposed gradient has the form ∇̂ ( · ) ` = ∂ ` ∂z∗ T∑ i=0 ( JF ) i ∂F ( z ∗ ; x ) ∂ ( · ) . ( 14 ) In practice , we found T = 1 to be already sufficient for implicit representation tasks ( see Fig . 4 ) , which requires only one evaluation of the vector-Jacobian product ∂ ` ∂z∗ JF . Spectral Normalization . We additionally apply spectral normalization on the layers F we use for the implicit models , which ensures that F ( for forward passes ) and JF ( for backward passes ) are contractive mappings , and is able to guarantee unique and stable fixed-points due to the Banach Fixed Point Theorem [ 3 ] . Specifically , we adopt the power-iteration in [ 20 , 4 ] to scale the spectral norms of W in Eq . ( 7 ) and ( 8 ) after each training iteration if it become > 1 , which incurs little additional computation cost [ 20 ] . However , we note that contractivity is not always necessary for fixed-point convergence ( especially in the presence of stronger fixed-point solvers , as shown in large-scale DEQ models [ 1 ] ) , and weaker regularizers such as weight normalization [ 25 ] may also suffice in practice .

Summary: Summary ------------------ This manuscript proposes the use of implicit differentiation (i.e. using the implicit function theorem) to train variants of so-called "implicit" representations (also known as "coordinate based networks")  the authors develop. Specifically the manuscript extends the formulations of SIRENs and MFNs in the form of Deep Equilibrium  Networks (DEQ) that can thus be trained through implicit differentiation. To make training of those  architectures efficient, the manuscript suggests two approximations. The first approximation lies in the  forward pass and circumvents the use of an expensive Fixed Point Solver. The second approximation lies in the backward pass and circumvents the expensive inversion of typically "large" Jacobian matrices. The authors evaluate their new architectures and training schemes on image, audio and video fitting and show that smaller Implicit^2 networks can fit signals better than SIREN or MFNs.