$(\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: This work suggests that combining implicit-depth models (e.g. DEQs) with implicit representation approaches (e.g. SIREN) can bring performance improvements. The method is evaluated on image, audio and video representation.