Regularizing the Infinite: Improved Generalization Performance with Deep Equilibrium Models

Venue: NeurIPS_cc_2024_Workshop_MLNCP
Published: 2024-09-12 10:20:42

Abstract: Implicit networks, such as Deep Equilibrium (DEQ) models, present unique opportunities for emerging computing paradigms. Unlike traditional feedforward (FFW) networks, DEQs adaptively adjust their compute resources which has been shown to improve out-of-distribution generalization, especially in algorithmic tasks. We demonstrate that this generalization includes robustness to noise making them well-suited for new hardware, such as analog or optical architectures, with higher yet noisy compute capabilities. But do DEQ models consistently outperform FFW networks in generalization? Surprisingly, our findings indicate that this advantage depends heavily on the specific task and network architecture. For equivalent network capacity, DEQ models prove more beneficial as the depth of the network increases—a trend that aligns with hardware systems optimized for deeper networks. We further demonstrate that regularizing the DEQ’s entire dynamic process, instead of random initialization or threshold prescribed in previous work, significantly enhances performance across various tasks, including image classification, function regression, adversarial training, and algorithmic extrapolation, making DEQs a compelling choice for next-generation hardware systems.

Main paper: 
 introduction: 
 Deep Equilibrium Models (DEQs) [3] are machine learning models whose predictions result from fixed-points of a learned transformation. By relying on fixed-point iterations these models effectively implement infinite-depth neural networks using only a finite-depth transformation. The DEQ formulation could be beneficial for bi-level optimization, for example in inverse problems [18] , and latent space optimization in generative models [14] . Additionally, DEQs are deemed to provide improved performance in out-of-distribution (OOD) generalization compared to traditional feedforward (FFW) thanks to the adaptive inference computation implied by the fixed-point formulation, in contrast to the fixed computation budget of FFWs [21, 2, 16] . To optimize memory usage during training, DEQs apply the implicit function theorem [15] , to reduce the memory footprint to O(1) with regards to the number of iterations. Despite memory efficiency, training and inference of DEQs consume more FLOPS compared to FFW networks. A number of methods have been proposed to alleviate the heavy computation of DEQs during training by regularizing the Jacobian of the weights and reducing the number of iterations to convergence [12, 13] , approximating the fixed-points [22] , or learning fixed-point solvers [5] . Despite all these, DEQs still require convergence of the fixed-point iteration for each input batch of data. This raises the question as to whether the extra computation compared to traditional FFW architectures provides sufficient benefits.We demonstrate that generalization capabilities of DEQs extends to the robustness of DEQ models in the presence of noise, which makes them ideal for emerging hardware technologies like analog [27] or optical [29] systems that offer greater compute but tend to be noisier. Counterintuitively, our findings reveal that the advantages of DEQ models over traditional FFWs are not universal; in fact, DEQ models are more robust compared to sufficiently deep FFWs. Additionally, we show that introducing noise into the dynamics of the DEQ acts as a regularizer and enhances OOD generalization. We propose modifications to the fixed-point solvers and stopping criteria to enable more efficient fixedpoint solving in noisy environments. Our study includes a range of tasks across various domains and problems, from image classification and adversarial robustness to algorithmic reasoning. 
related work: 
 Noisy Neural Networks Our method is related to prior work that studies the noising of neural network activations, either through dropout or the addition of additive/multiplicative Gaussian noise [23, 10, 28, 8, 19, 17] . These networks can be categorized as either FFWs [8] , recurrent neural networks (RNNs) [10, 17] , or neural Ordinary Differential Equations (ODEs), which with noise are transformed into stochastic differential equations (SDEs) [28, 19] . It has been demonstrated that noising activations contributes to the improved stability of recurrent networks through explicit [8] or implicit [17] regularization during training. This regularization can promote wider minima in the loss landscape [17, 8] that is advantageous, for example, in adversarial settings [28, 19] . Our work differs from previous work in that it focuses on DEQs with fixed-points, in contrast to SDEs/RNNs that do not necessarily reach a fixed-point. Second, we analytically compare the dynamics of the regularized DEQ and regularized FFW, and study the conditions under which a DEQ is more robust than a FFW. We look at the applications in multiple domains, including regression, classification, algorithmic tasks, and adversarial robustness. In the adversarial setting, we propose to regularize the entire dynamics of the DEQ instead of using costly optimization on the input data [30, 31] .Out-of-Distribution generalization in DEQs Our work is related to the adaptive inference computation of recurrent networks [21] and DEQs [2] to achieve OOD generalization [16] . In the literature, the tasks assigned to DEQ networks are predominantly algorithmic reasoning, such as computing prefix sums and solving 2D mazes [2, 21, 25] . Our work differs in that these studies employ vanilla DEQ models without any form of regularization. We demonstrate that DEQs are not intrinsically superior to FFWs, however, in OOD settings, they tend to perform better under noise when the strength of the distribution shift increases. 
preliminaries: deq's (perturbed) dynamics: 
 The forward pass iteration of a DEQ can be written asEQUATIONwhere x ∈ R l is the input vector, z [m] ∈ R d represents the intermediate state after the m-th iteration (with m = 0, . . . , M -1), f θ : R d × R l → R d defines the transformation at each layer, θ denotes the set of parameters across the layers, ∥ • ∥ 2 denotes l 2 norm, and ϵ is a small constant. To find the fixed-point z * of Eq. 1, solvers such as Broyden's method [7] and Anderson acceleration [1] are employed. These solvers find the solution z * when either the stop condition in Eq. 1 is met or the threshold M is reached. For distribution shifts in data, which can be formulated as x + ∆x, the dynamics of the DEQs change via:EQUATIONHere, z [m] represents the state of the unperturbed system after the m-th iteration, while ẑ[m] denotes the state of the perturbed system at the same iteration. As Eq. ( 2) suggests, in order to bound the difference between the fixed-point of the model incurred by the original input x and the fixed-point of the distribution-shifted x + ∆x, previous work [31] has utilized an optimization process that corrects perturbed inputs x via a second iterative process at each step of the dynamics to decrease second term in Eq. ( 2) by reducing the entropy of the perturbed inputs in classification tasks. In contrast, we propose regularizing the entire dynamics (both terms in Eq. ( 2)) by introducing random perturbations throughout the evolution of the system. This approach aims to ensure that, for perturbed data, the perturbed and unperturbed dynamics remain closely aligned. This is achieved while being agnostic to the learning task and without requiring a second optimization per step of the DEQ dynamics.To regularize the dynamics of the DEQ, we introduce perturbations to the intermediate states.Depending on the architecture, this can be achieved by perturbing either the states z or implicitly through perturbing the input injection x. Since the magnitude of these states is not known a priori, we choose a multiplicative (input dependent) noise form to ensure we have control over the strength of the signal-to-noise ratio (SNR) of the states. Formally, the states and inputs of f θ in Eq. 1 could be perturbed as follows:EQUATIONwhereEQUATION) Note that SNR x and SNR z are hyperparameters chosen to fix the SNR with respect to each variable during the dynamics. Additionally, we consider a stronger case where noise is introduced before the last nonlinearity of the unit-cell. We refer to this as the 'signal' noise with SNR s . For example for a simple MLP unit-cell z m+1 = σ(W z m + x) + ϵ s , noise ϵ s scales with the signal W z + x. The fixed-point solution of Eq. 3 is used to minimize the loss function L(θ) for learning the task:EQUATIONwhere z * is the fixed-point solution and ϵ * z is the corresponding noise. To avoid inefficient application of solvers, in the presence of noise, when the fixed-points are already reached, we modify the solvers and stopping criterion to accommodate the variability introduced by noise at each step. We aim to find solutions of the formEQUATIONwhere z * is the fixed-point of f θ . Inspired by the stochastic approximation framework, we adopt the Robbins-Monro algorithm [24] , which samples the stochastic function f θ :EQUATIONHere, z * represents the solution that is the average of the intermediate values up to the M -th iteration, while α serves as a hyperparameter. This method is referred to as 'stochastic unrolling.' 1 Rather than averaging all intermediate values, we implement averaging over a window of width w during the solver step of DEQ, which reduces the memory footprint of our model to the size of w-often a small number. Additionally, we employ more sophisticated solvers such as Anderson and Broyden. These solvers feature an internal averaging mechanism (e.g., the memory parameter in Anderson), which has the potential to obviate the need for explicit averaging of solutions. As for the stopping criterion, we utilize the relative norm applied to the averaged solution within the window-see appendix B.In this section, we intend to compare regularized DEQ and FFW to determine if a fixed-point network is necessarily associated with stronger robustness. For this we use a simple one-layer fully connected network as the unit-cell for both networks. We assume an additive noise perturbation and refer to the multiplicative perturbation in the Appendix A.1.2 and A.2.2. For DEQ, we haveEQUATIONwhere W ∈ R d×d is the weight matrix, x ∈ R d is the input injection vector, η is a scalar representing the noise amplitude, ξ ∈ R d is a noise vector with elements drawn from a Gaussian distribution N (0, I), and σ is the nonlinearity with σ ′ being its derivative. The sensitivity of the DEQ's hidden state with respect to the noise vector ξ can be analyzed through the derivative dz m dξi , where ξ i is the i-th element of ξ. This derivative, at iteration m + 1, is given by: where s m is the function argument defined as s m := W z m + x + ηξ, diag(•) is the diagonal matrix formed from the vector argument, and δ i is the Kronecker delta vector, which has a one at the i-th position and zeros elsewhere. The FFW network at layer m is described by the following equation:EQUATIONEQUATIONHere, W m ∈ R d×d is the weight matrix for layer m, and the rest of the terms are as previously defined for the DEQ. The sensitivity of the FFW's hidden state to the noise vector is similarly assessed by dz m dξi , resulting in:EQUATIONTo compare the sensitivity of DEQ and FFW to noise, we examine the influence of their respective update equations on the propagation of perturbations to state z m . We note that for the DEQ model, the update Eq. 9 iteratively applies the same weight matrix W . The norm of the derivative of the hidden state with respect to the noise, i.e. dz m dξi , is affected by the spectral radius ρ(W ) m . Conversely, the FFW model's update Eq. 11 involves a product of distinct weight matrices W m . The sensitivity derivative now depends on the spectral radius of the matrix product ρ(W 1 W 2 . . . W m ).In both models, the nonlinearity σ plays a crucial role in modulating sensitivity. Assuming a Lipschitz continuous σ with Lipschitz constant L, the derivative norms could be (loosely) upper bounded byEQUATIONDEQs can be constrained using Jacobian regularization [6] to achieve smaller spectral radii. Similarly, FFWs could also be regularized to have smaller Jacobians. Although this process is more costly for FFWs since the regularization needs to be computed at each layer, as opposed to only at the fixed-point for DEQs, we empirically found that DEQs still maintain smaller Jacobians than their FFW counterparts with Jacobian regularization-see Fig. 13 in Appendix D.2. 
experiments: 
 We We begin our examination by assessing the robustness of DEQ models trained on the MNIST dataset against perturbations applied at each iteration of the network. We compare these models to a multilayer FFW model. For the multilayer FFW, we increased the number of layers, each with different weights, while maintaining a consistent hidden size of 128 in each layer. In contrast, for the DEQ models, we expanded the width of the hidden state to ensure the number of parameters remained equivalent between both models. Consequently, the DEQ widths of 128, 223, 288, 340, 386, 427, and 467 correspond to FFWs ranging from single-layer up to thirteen-layer configurations. We look into the robustness of the DEQ models in the previous section from the sensitivity perspective.The right hand-side of the Fig. 2 shows the sensitivity ratio E ξi dz * dξi of the DEQ to FFW network, where the sensitivities are averaged over 100 noise samples for all dimensions. Note that z * represents the fixed-point for the DEQ or the last layer activations for the FFW. We observe that as we increase the depth of the FFW, the DEQ with the corresponding number of parameters shows increasingly superior robustness. We also examine the eigenvalue distribution of the DEQ and FFW models in the left hand-side of Fig. 2 . It is noteworthy to note that the DEQ models are less sensitive even when FFW networks are also Jacobian regularized-see Fig. 13 in Appendix D.2. This trend remains consistent as the DEQ models become wider or the FFW models become deeper.We conduct experiments on vision tasks using the Multiscale-DEQ (MDEQ) [4] , where multiple scales of the image are iterated to a fixed-point. We use up to four scales with number of filters of 32, 64, and 128 for the MDEQ architecture. In contrast, the FFW architecture maintains a consistent number of 128 filters across all layers. The perturbation applied in both DEQ and FFW architectures is a signal perturbation that is added before the nonlinearity of the unit layer. We plot the accuracy of the two architectures versus the branch/layer of the DEQ/FFW for inference perturbation in Figure 1e . These results are consistent with those from the analytical and small-scale DEQ models.For adversarial experiments, we utilized adversarial training (AT) and dynamics regularization as the defense mechanism. For DEQ models, we propose injecting noise into either z or x, while for FFWs, only z is perturbed. Noise is added at two levels of SNRs, 0dB (denoted as SNR 1 in linear unit) and 10dB, with respect to the noise-injected variable. We also considered the random thresholding technique from previous work [31] . In all scenarios, models are trained with either projected gradient descent (PGD) [20] or TRADES [32] for AT with ℓ ∞ -norm perturbations, ϵ = 0.3 (the maximum allowable perturbation under the l ∞ norm), step size 0. 3a ,b shows the accuracy of the models for DEQ and FFW under PGD attack and various regularization defenses. As seen, DEQ and SNR= 1, x as well as FFW and SNR= 1 usually have the strongest defense. We then use these models and compare the performances of the models against various attacks in Fig. 3c -f and supplementary Fig. 16 . Interestingly, as the number of steps of the PGD attack increases, the gap between DEQ and FFW increases. This is also consistent with the strongest AA attack, in which the number of iterations is chosen automatically. 
broader impact: 
 In this work, we examined the robustness of DEQ networks, which learn fixed-point representations of data. We demonstrate that the fixed points of DEQs do not inherently offer improved performance compared to FFW networks. Previous work [2] has observed that DEQs mainly exhibit advantages during inference, particularly when applied to OOD algorithmic tasks. However, in the presence of noise, this advantage does not readily extend to common tasks such as classification and regression and is only noticeable for deeper FFW networks. We hypothesize that this behavior can be explained by the spectral norm of both the FFW and DEQ, at least for simple MLP architectures, which is empirically confirmed in larger MDEQ models. Additionally, we have developed a perturbation regularization approach that is task-agnostic and has proven effective in enhancing robustness over FFWs, including in shallow configurations. All these properties make DEQs a strong candidate for emerging analog hardware as a new computing paradigm [29, 27] that is not only robust to the system's inherent noise but also presents an opportunity to provide new capabilities.Our aim is to understand the robustness of the DEQ/FFW neural networks against intermediate perturbation applied at each iteration/layer of the network.Note: Throughout this section, vectors are denoted by boldface lowercase x letters, matrices are indicated by uppercase and boldface letters X, and scalar values are represented by non-bold, lowercase letters x.We study two cases of additive and multiplicative perturbations where the perturbation is a normal Gaussian noise.We are assuming a unit-cell of the following form for the DEQ:EQUATIONIn this equation,z m ∈ R d , W ∈ R d×d represents the weight matrix, x ∈ R l where x = σ 2 (U x)is the input projection layer, η is a scalar, and ξ is to denote an d-dimension noise vector whose elements are from a Gaussian distribution N (0, 1), and σ is the nonlinearity of the unit-cell.Lets define the pre-nonlinearity activation, i.e. s m := W z m + x + ηξ. We assume ξ is sampled once and remains fixed throughout the evolution of the equation.To characterize the robustness of the DEQ, following [26] , we look at the vector dz m dξi where ξ i is the i-th element of ξ. From 13, we have:EQUATIONwhere δ i ∈ R d represent the Kronecker delta vector, defined as δ i = [0, . . . , 0, 1 i , 0, . . . , 0] t where 1 i is the unit element positioned at the i-th component, and all other elements are zero, • is the element-wise hadamard product and σ ′ sm is the derivative of σ computed at s m .Using Eq. 15, a number of iterations can be written as:EQUATION)EQUATION)EQUATION...where W k denotes k application of matrix W , i.e.W k := W W • • • W k times .We are assuming a unit-cell of the following form for the DEQ:EQUATIONwhereEQUATIONSimilar to the additive noise, the sensitivity parameter dz m dξi reads as:EQUATIONIn Eq. 24, dGm dξi can be further simplified using 22:EQUATION)where s t m is the normalized s m . Using Eqs. 24 and 25, a number of iterations can be written as:EQUATION)EQUATION)EQUATION... (29) (30)We study two cases of additive and multiplicative perturbations where the perturbation is a normal Gaussian noise.We are assuming hidden layers which are of the following form for the FFW network:EQUATIONwhere z 0 = x = σ 2 (U x), and W i represents the weight matrix at the i-th hidden layer.To characterize the robustness of the FFW model, we again look at the vector dz m dξi . From 31, we have:EQUATIONUsing Eq. 33, a number of iterations can be written as:EQUATION)EQUATIONEQUATION)whereW k denotes k application of matrix W k , i.e. W k := W k W k-1 • • • W 1 k times .We are assuming hidden layers which are of the following form for the FFW network:EQUATIONwhereEQUATIONSimilar to the additive noise, the sensitivity parameter dz m dξi reads as:EQUATIONIn Eq. 43, dGm dξi can be further simplified using 41:dG m dξ i = (W m z m + x) t (W m dz m dξi ) ||W z m + x|| 2 = (W m z m + x) t (W m dz m dξi ) G m = s t m W m dz m dξ i (44)where s t m is the normalized s m . Using Eqs. 43 and 44, a number of iterations can be written as:EQUATION)EQUATION... (48) (49)In the case of additive perturbation, we compare the sensitivity of the DEQ to a feed-forward network with one hidden layer. This specific scenario is the only one in which the two networks have the same number of parameters. In this case, we haveEQUATION)EQUATIONHere, M represents the number of iterations required for the DEQ network to converge. We observe that for a single hidden layer in an FFW, the sensitivity is not directly tied to W 1 , instead, it is related indirectly through the derivative of the nonlinearity σ ′ s0 . Nonlinear functions, such as the ReLU, which have a derivative of zero for certain input ranges, tend to be more robust.We propose two mechanisms that could augment fixed-point solvers used in DEQ models in the presence of noise. One is the averaging of the solutions using stochastic approximation, and the other is modifying the stopping criterion based on the averaged values. In particular, we propose to use a stochastic-unrolling (SU) solver based onEQUATIONand a relative-averaged stopping criterion over a window of length w defined as:EQUATIONwhere z[w] is to denote averaging over the window w. We note that other fixed-point solvers, such as the Anderson accelerator [1] or Broyden's method [7] , could utilize the averaged-relative stopping criterion. To compare the performances of the solvers in the presence of noise we use a DEQ with a unit-cell consisting of one hidden layer of size 100 with a fully connected structure and Tanh nonlinearity, where the weight matrix is initialized in U (-1/ √ 100, 1/ √ 100). We consider multiplicative perturbation and two regimes of low (0dB) and high (20dB) SNR levels.Empirically, we find that averaging assists in finding a solution more effectively than the native unrolling solver (see Fig. 4a-d ). Methods such as Anderson's and Broyden's can find solutions without explicit averaging due to their inherent averaging capabilities. We observe that the averagedrelative criterion (see Fig. 4c, d ) is more effective at detecting the fixed-point compared to the relative criterion (see Fig. 4a, b ), potentially preventing unnecessary continuation of the solving process, thereby speeding up the procedure and reducing computation.To further demonstrate the improvement of the average-relative stopping criterion, we apply it to MNIST classification with signal (pre-activation) noising (see Fig. 5 ). The average-relative criterion provides a lower number of iterations for convergence. We also examine the impact of the stochastic solver on performance. Using the MDEQ architecture with four branches on CIFAR-10 classification with test-time signal noise perturbation (see Fig. 6 ), we note that the stochastic solver achieves the same accuracy as the Broyden method. As the gradient tape for the intermediate states of DEQs are not stored by the solvers, loss gradient cannot backpropagate to the input. We followed [31] and employed fixed-point solutions followed by unrolling for a few steps:EQUATIONwith z * j as the solution of the fixed-point solver and j ranging from 1 to K. The resulting state, z * K , post-unrolling was then utilized to calculate the loss and gradients as:EQUATIONwhere ∆x represents the input perturbed, and y is the ground-truth label. We used M = 8 (unrolling threshold), K = 9 and λ = 0.5.Additional experiments for MNIST classification, regression, and Adversarial training is provided.We investigated the robustness of DEQ models in regression tasks. Specifically, we examined the regression of a 1-D function defined by y = 2(exp(- The SNR sweeps for each activation are plotted in Figures 7, 8 , and 9. We note that the robustness of the FFW models is particularly sensitive to the choice of input projection activation in shallower networks. For example, with ReLU, since the derivative of the activation is zero for negative values, shallower FFW models preserve robustness to noise.In comparison to classification tasks, the advantage of the DEQ becomes apparent even with shallower FFW architectures. We believe this difference arises because, in classification, there is a range of perturbation that can be tolerated before misclassification occurs, whereas in regression, perturbations directly affect the accuracy of the function being learned.The accuracy of the DEQ and FFW models on the MNIST classification tasks, as discussed in Section 4.1, is plotted against test-time SNR ranges for three scenarios: no train-time noise (Fig. 10 ), train-time noise with an SNR of 10dB (Fig. 11 ), and an SNR of 0dB (Fig. 12 ). It is noteworthy that the DEQ model exhibits greater robustness to noise compared to the FFW model, especially as the FFW model becomes deeper and at lower SNR ranges. We test additional noise defense mechanisms against adversarial attacks on the MNIST dataset. Furthermore, we present the effectiveness of the best noise defense for DEQs and FFWs trained with PGD (epsilon=0.3, step-size=0.1, and number of steps=10) in Fig. 16 , and for those trained with TRADES (with a loss coefficient of 6, as detailed in [31] ) in Fig. 17 For MNIST classification, we perturb W z + b + x in Eq. 56, which we refer to as signal perturbation.We employed multiplicative noise and adjusted the signal-to-noise ratio (SNR) level using a coefficient.During training, we used SNR values of either 0dB, 10dB, or no perturbation (an SNR of ∞). At inference, we varied the SNR coefficient to evaluate the model under different SNRs.For CIFAR-10 classification, we employed the multiscale DEQ architecture from [4] , which processes the image at various scales in branches of the network. We used up to four scales with hidden channel sizes of 32, 64, and 128 for the DEQ architecture. The FFW architecture has a consistent hidden channel size of 128 across all layers. Tables 3 and 4 include details of the architectures and training. For CIFAR-10 classification, we perturb either the x, z, or activations before nonlinearity in the unit-cell of DEQ or layer of FFW, which we refer to as signal perturbation.The network architecture of DEQ and FFW for regression tasks is similar to the MNIST classification, except that the input projection layer for regression, is either RelU, SiLU, or identity nonlinearity. complexity with mazes of various sizes, specifically 9x9, 13x13, 17x17, 21x21, 25x25, 27x27, 29x29, 31x31, 33x33, 35x35, 37x37, and 59x59. We follow the training procedure outlined in [2] . For DEQ, we use the residual block structure with two convolutions, each with 128 hidden channels. We do not use any normalization and apply a Tanh nonlinearity. For FFW, we use a single hidden layer residual block with the same nonlinearity as DEQ.In the Robbins-Monro algorithm, α is typically inversely proportional to the iteration number m. However, empirical evidence suggests that, for a sufficiently large M , a fixed α is adequate for finding solutions. 
 
 Appendix:  



Reviews: review: title: title: Interesting regularisation method for Deep Equilibrium (DEQ) Models
main_review: main_review: This paper studies the use of injecting noise into the activations (intermediate network states) of DEQs. This is a method motivated by previous findings that such noise injection made RNNs more robust to input perturbations. The proposed method is a straightforward application of this idea and is evaluated on several tasks and benchmarks. Overall, the work is novel and a good fit to the workshop.

**Strengths:**
- Interesting method that seems to work well at improving robustness to noisy inputs.
- I appreciated the attempt to systematically characterise when exactly DEQs are more robust than feedforward (FFW) networks with the same parameter count.

**Limitations:**
- I did not find the theoretical arguments comparing DEQs and FFW models very compelling. The derivation of the network sensitivity is straightforward enough, but the argued implication that DEQs are more robust only in comparison to deeper FFW models does not necessarily follow from the presented analysis. This seems to be more of an empirical finding, and I would suggest the authors to rephrase/soften their statements in the introduction and abstract accordingly.
- The use of Gaussian noise in the inputs seems to me a very limited method to test robustness. For example, is it really surprising that a model trained with noise is more robust to the same kinds of noise? Some comments from the authors on this point would be appreciated.