Efficient Score Matching with Deep Equilibrium Layers
Authors: Yuhao Huang
Qingsong Wang
Akwum Onwunta
Bao Wang
Venue: ICLR 2024 poster (2024)
Keywords: score matching
deep equilibrium model
density estimation
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We improve the memory efficiency of score matching by leveraging deep equilibrium models

Abstract: Score matching methods -- estimate probability densities without computing the normalization constant -- are particularly useful in deep learning. However, computational and memory costs of score matching methods can be prohibitive for high-dimensional data or complex models, particularly due to the derivatives or Hessians of the log density function appearing in the objective function. Some existing approaches modify the objective function to reduce the quadratic computational complexity for Hessian computation. However, the memory bottleneck of score matching methods remains for deep learning. This study improves the memory efficiency of score matching by leveraging deep equilibrium models. We provide a theoretical analysis of deep equilibrium models for scoring matching and applying implicit differentiation to higher-order derivatives. Empirical evaluations demonstrate that our approach enables the development of deep and expressive models with improved performance and comparable computational and memory costs over shallow architectures.

Published as a conference paper at ICLR 2024
EFFICIENT SCORE MATCHING WITH DEEP EQUILIB- RIUM LAYERS
Yuhao Huang', Qingsong Wang', Akwum Onwunta", & Bao Wang" 'Department of Mathematics and Scientific Computing and Imaging (SCI) Institute University of Utah, Salt Lake City, UT 84102, USA *Department of Industrial and Systems Engineering Lehigh University, Bethlehem, PA 18015, USA
ABSTRACT
Score matching methods, which estimate probability densities without computing the normalization constant, are particularly useful in deep learning. However, the computational and memory costs of score matching methods can be prohibitive for sigh-dimensional data complex models, particularly due to the derivatives or Hessians of the log density function appearing in the objective function. Some existing approaches modify the objective function to reduce the quadratic compu- tational complexity for the Hessian computation. However, the memory bottle- neck of score matching methods remains for deep learning. This study improves the memory efficiency of score matching by leveraging deep equilibrium models. We provide theoretical analysis of deep equilibrium models for scoring match- ing and applying implicit differentiation to higher-order derivatives. Empirical evaluations demonstrate that our approach enables the development of deep and expressive models with improved performance and comparable computational and memory costs over shallow architectures.
INTRODUCTION
40 is a class of density estimation methods, which avoid computing the normalization constant of the density function; see Section 2.1 for a brief re- view of score matching methods. Score matching is particularly useful when the density function parameterized by deep neural network (DNN). At its core, score matching utilizes the fact that the (Stein) score function the gradient of log density function does not require the computation of the normalization constant; one can estimate the density function by minimizing the discrep- ancy between the score functions of the true data distribution and the model distribution. Given its flexibility, score matching methods have been applied to various tasks requiring density estima- tion [19
Score matching [16 and its variants 33
34
43
especially diffusion models; e.g., B7 38
Despite the remarkable advantages of score matching, its applications to DNNs face computa- tional and memory challenges, arising from optimizing the objective function of the score matching method, which necessitates the computation of the derivative and Hessian of the log density func- tion [16] [39] This step, coupled with the subsequent hackpropagation, which additionally computes the gradient with respect to the weights, significantly inflates the computational graph for training DNNs 19 25 39 40).
Several strategies have been proposed to address the computational and memory challenges of score matching methods, focusing on reducing the quadratic computational cost of the Hessian. For in- stance. methods such as approximate backpropagation [19 and curvature propagation (CP) (25 provide estimations for the Hessian. Additionally, there are efforts to modify the score matching objective functions, including denoising score matching (DSM) 43 which avoids the Hessian computation by considering an objective function with a perturbed data distribution; see equation ? Likewise, the sliced score matching (SSM) [40] is proposed to replace the Hessian matrix compu- tation with the Hessian vector product; see equation 3 These methods alleviate the computational cost of the Hessian term in the score matching methods, but the memory bottleneck that arises from optimizing the objective function that involves derivatives of the log density function remains. This
Correspond to wangbaon wangbaonj@gmai com

---

Published as a conference paper at ICLR 2024
memory bottleneck restricts the depth and complexity of models, limiting their potential in deep learning applications. For example, in Section 4.1 a VAE variant is trained with SSM objective function; while increased model depth offers noticeable performance gains, it comes at a substan- tial cost of memory and computational overhead: see Table for details. Deep equilibrium model presents attractive solution to address these challenges. In DEQ, the hidden layer is represented as the equilibrium point of a nonlinear fixed-point equation that is equivalent to a weight-tied infinite-depth neural network (45). Importantly, by applying the implicit function the- orem [ g] at this equilibrium, DEQ facilitates the computation of the derivatives of the hidden representation without referencing intermediate activation values. As a result, DEQ enjoys constant memory efficiency while being a deep and expressive model 15
(DEQ) (3 45
Inspired by the memory efficiency and remarkable performance of DEQ we propose to address the computational and memory bottlenecks of score matching methods by leveraging DEQ. Using the implicit function theorem, the score function can be computed in terms of the fixed point of the DEQ without storing the intermediate activation values. However, this approach requires computing the inverse of the Jacobian matrix of the DEQ, which can be expensive. We then use efficient implementation of the backpropagation that utilizes the so-called phantom gradient [12] in the presence of higher-order derivatives. As a result, the proposed DEQ-based score matching methods allow for the design of deep and expressive models with computational and memory costs comparable to those of shallow architectures. We demonstrate the effectiveness of the proposed method through empirical evaluations of both density estimation and generative modeling.
1.1 OUR CONTRIBUTIONS
We present a systematic study on integrating DEQ into density estimation tasks or when the loss function involves derivatives in general. Our work makes the following key contributions:
We integrate DEQs into score matching models address their memory challenges.
We provide an efficient implementation of the backpropagation and analyze its convergence in the presence of higher-order derivatives.
Through empirical studies, we confirm that the integration of DEQs allows for the design of deeper and more expressive models without elevating computational and memory demands
1.2 ADDITIONAL RELATED WORKS
Beyond the score function, some works consider higher-order derivatives of the log density func- tion 129.. For example. [26] terms these derivatives as higher-order scores and presents general- ization of DSM. FDSSM, introduced in (28), adopts finite difference to approximate the directional derivatives of the log density function. It uses random projection approximate the gradient terms. As observed in their experiments [28), Section 6. the variance introduced by the random projection can outweigh the computational benefits of the finite difference approximation.
Another additional line of related work is using DEQ to improve the memory efficiency of deep learning models. [15 combines DEQ with the implicit representations to achieve improved perfor- mance with less memory. In optical flow estimations, DEQ model (7] demonstrates faster computa- tion of the flow, significantly reduced memory needs, and improves over state-of-the-an models.
1.3 ORGANIZATION
This paper is organized as follows: Section 2 provides an overview of the score matching method and its challenges, followed by an introduction to DEQ. In Section 3 we present our proposed method for integrating DEQs into score matching models. Specifically, Section 3. describes DEQ for score matching models, Section 3.2 discusses the well-posedness of the fixed-point equation, and Sec- tion 3.3 presents the implicit differentiation for oigher-order derivatives. Section 3.4 describes the efficient implementation of backpropagation in the presence of higher-order derivatives. Section presents the empirical evaluations of the proposed method on density estimation and generative modeling tasks.
2
BACKGROUND AND PRELIMINARIES introduction to DEQ.
This section provides overview of the score matching method and its challenges, followed by an
2.1
SCORE MATCHING FOR UNNORMALIZED DENSITY ESTIMATION
In machine learning, statistical models often rely on unnormalized density function (-).). which is proportional to the normalized (probability) density function po(+) = 2, (-),,, where Zo the nor-
2

---

Published as a conference paper at ICLR 2024
malization constant. The normalization constant depends on the parameter and is often intractable to compute. Thus the maximum likelihood estimation, which minimizes the Kullback- (KL) divergence between the true data density Pdata(-) and the model density Po(-) is infeasible.
Score matching (SM). Instead of using KL divergence to quantify the discrepancy between the true density Pdata and the normalized density po(-), one can alternatively use the Fisher divergence Dr that compares the score functions V log Pdata (x) and V log (og))( of the two densities:
1
D3 (po, Pdata)
Eu Ewrppaara
| lo V log Pdata (z: (z)
2
Ex Ba-pann
V (o(g)
V log (xataa))))
2
Dr (No, Pdata)
showing that Fisher divergence avoids the computation of the normalization constant Ze. However, it requires computing the score function of the true data density Pdata (-), which is inaccessible.
Based on some regularity assumption of the data density pdata()) Hyvärinen [66 finds that the
Fisher divergence can be rewritten as Dp = JsM + C, where C is a constant independent of 0 and I is defined follows:
1
(0) == Emcopaata [[| to (x) + Ba [tr (V Po (
(1)
where tr(-) denotes the trace operator. Then, the proposed score matching method in 16 finds the optimal parameters 0 by minimizing the objective functiin JsM-
Despite its theoretical advantages, the score matching method faces computational and memory challenges when using DNNs to model the score function. The challenge comes from the computa- tion of the derivative and the trace of Hessian terms in the objective function JSM- For example, the common implementation of automatic differentiation like PyTorch [30 requires computing the full Hessian matrix V2 log (2x) even when only need its trace. The quadratic computational com- plexity of the Hessian matrix computation is prohibitive for high-dimensional data or deep models. Furthermore, when optimizing the objective function JsM with backpropagation. it requires tak- ing the derivative of the score function V (og and the Hessian matrix V log pe (æ), which necessitates storing their computational graph which leads to prohibitive memory costs. Several methods have been proposed to reduce the computational and memory costs of score matching.
Denoising Score Matching (DSM). 43 proposes to perturb the data æ according to noise distri- bution Pa (m(zz) and then estimate the score function of the perturbed data ă. Specifically, when the noise distribution is Gaussian with variance o², DSM minimizes the following objective function:
JDSM (0) = Pam-Pin
Puata By (#))z)
V log Po
This method avoids the computation of the Hessian matrix V log (a but at the cost of recovering the perturbed data Po (w) instead of the true data distribution Pdata(z). Mcanwhile, the model is sensitive to the parameter a and requires heuristics and careful selection of the value of a [50
Sliced Score Matching (SSM). [40] proposes a random projection-based approach for efficient score matching by replacing computing the Hessian matrix V2 log (o(x) with Hessian vector prod- ucts. SSM minimizes the following objective function:
(v V. (og)(()))) + + V2 loggo(2)0)
JSSM (0) E,, Eeuleerrraat
(3)
where v is random vector sampled from the normal distribution Rademacher distribution. The term 3 (v V z log (o))) can be computed analytically, and equals to V log po lo)
This leads to SSM with the variance reduction (SSM-VR) objective function [40]
JSSM-VR (0) = Ba
2 log ( + デ V (o
(4)
Even though the above methods reduce the computational cost of the Hessian matrix, they still include derivative in terms with respect to the data æ in the objective function, which requires storing the computational graph for backpropagation and leads computational and memory bottlenecks.
3

---

Published as a conference paper at ICLR 2024
2.2 DEEP EQUILIBRIUM NETWORK
DEQ, first presented in [3] represents the hidden representation, z², as the equilibrium of a specific fxxe-ppoint iteration equation:
fo(=", x))
where fe represents neural network parameterized by 0, and denotes the input data. When using an iterative solver for this fixed-point equation, the equilibrium of the DEQ corresponds to the final hidden representation of an infinite-depth neural network. By leveraging implicit differentiation during training. DEQ eliminates the need to store intermediate hidden representations; results in constant memory usage despite being a deep model significant computational advantage. The theoretical convergence of DEQ has been extensively studied in [5). Further applications prove the model's s versatility and competitive performance across various tasks. These include language modeling [3], semantic segmentation
4 optical flow [7]. diffusion models [1]. and maximum
a a-posteriori (MAP) estimates [42)]
3 DEEP EQUILIBRIUM LAYERS FOR SCORE MATCHING
In existing score-based density estimation using DNNs, the output of the DNN is typically the log density functionn and the Fisher divergence-related objective functions are optimized. These objec- tive functions often involve the derivatives of the log density function and necessitate the computa- tion of higher-order derivatives in backpropagation. We propose use DEQ for the core interme- diate layer of the DNN architecture that accounts for the main depth and complexity of the model. Then, we can utilize the memory efficiency of DEQ to reduce the overall memory consumption.
3.1 GENERAL STRUCTURE OF DEEP EQUILIBRIUM LAYERS
The general structure of our DFQ-assisted score matching model is outlined as follows: An input data, F, first being transformed through simple multi-layer perceptron (MLP). Subsequently, it is processed by a deep equilibrium layer, which finds the fixed point x" of the following equation:
z (n+1) = [=( (m)),,,;)
,x),
where So is a neural network parameterized by 9. The computed fixed point, z", is then converted through simple network into the log density, log (= ) or score function, V log Po (2*), depending on the task. Therefore, the main computational demands revolve around DEQ.
Specifically, the iterative functions So consist of one few sequential blocks of the following form:
(5) with a being an element-wise activation function, ( is the transformed input for this block, and W is the weight matrix. This structure can represent either fully connected layer when W is dense matrix or a convolutional layer when W is a convolutional kernel. Therefore, it provides a flexible framework for adapting the existing density score estimate models into DEQ.
(n+1) o (W x (x) + g(y)),
3.2 THE WELL-POSEDNESS OF THE FIXED-POINT EQUATION
In order to correctly apply DEQ, the fixed-point equation must be well-posed; that is, the fixed point 2 must exist and be unique. The soll-posedness of the fixed-point equation can be guaranteed by constraining the Frobenius norm of the weight matrix W to be less than one. In particular, we apply the following weight normalization to W: W - W W/X([] + +), where > is hyperparameter, e is small positive number, and [ denotes the Frobenius norm of W. The computational overhead of this normalization is minimal. This constraint ensures the spectral norm of the weight matrix W to be less than one, and hence the map = - Wz + g(y) is a contraction. Given most activation functions, such as Softplus used in equation are non-expansive, the composite map + (W + g())) is also a contraction mapping in and is fo x). Thus, Banach's 's fixed-point theorem [46] guarantees the fixed point's well well-posedness. Furthermore, the direct Picard iteration will converge linearly to the fixed point z".
3.3 IMPLICIT DIFFERENTIATION FOR HIGHER- ORDER DERIVATIVES
In this subsection, we describe the computation of the score or its derivatives with respect to the input data æ in the presence of the DEQ component. Let æ (I), I2, , T Idd be the input
data. We use x E R. to denote the fixed point of DEQ. Then the i-th component of the score function V, log (=(=)) for 1
< i < d can be computed as follows:
& log P=(=") 0x,
a log (z")) 0z"
(6)
dz
oz
4

---

Published as a conference paper at ICLR 2024
8 log Fo( for S S d can Du;Oa;
I. Likewise, the second-order derivative of the log density function be computed as follows:
8º (o(")) 8² log po(=") ) Oz Oz Oz. day
a log p=
(7)
dz,
0z-0z"
Oz dr,
Therefore, the main computational challenge is to compute the term
Oz de:
and
0= dw,day
For this, we
apply the implicit function theorem to obtain the following result: Proposition 1. Given Jo(z, x),, a continuously d
that is contraction mapping.
Let z² is the fixed point of the equation fo(z", ..) = =", then the matrix I
0fa
invertible.
Furthermore, the derivative of x" with respect to x can be computed as follows:
0z
0
05e
(8)
dz,
dz
da,
and if fe is additionally twice continuously differentiable, the second-order derivative of x" with respect to x can be computed as follows:
882= 0w:02;
(1
8 to 0w:0us
50 0zOz,
8 fo
82"
8 fe
8z" 8z" dz, 0my
(9)
8z
D=Oz:
0z,
3z3z
To the best of our knowledge, equation is well-known in the DEQ literature; see c.g. B1. while equation 9 is new in DEQ context. In practice, the Softplus activation function is used in DNNs to approximate the log density function, ensuring fe is twice continuously differentiable [40]. When using the reverse-mode automatic differentiation, the score function and its derivatives can be com- puted by only using the states of the DEQ at the fixed point
3.4 EFFICIENT IMPLEMENTATION WITH DEQ
There have been efforts devoted to improving training DEQs, including the use of Anderson ac- celeration 3 and separate neural block 55) to accelerate the fixed-point finding. Effective tech- niques have been employed for the forward pass, such as reusing the fixed point from previous itera- tions [5] Jacobian regularization (6. Meanwhile, the inexact gradient [12] can render the com- putational cost of backpropagation negligible without sacrificing the model's performance 7 2 31 In our DEQs, we adopt the strategy of reusing the fixed point from prior iterations to accelerate the forward pass. Additionally, we use inexact gradients, which we will describe in detail in this sub- section. Particularly, we will analyze its behavior for higher-order derivatives.
Proposition highlights DEQ's capability in computing the score function and its derivatives without storing the intermediate activation values. Yet, the computation of the matrix inversion 0= ---) can be expensive. Inspired by the success of inexact gradient in training DEQs, we for computing the first and second-order derivatives of the score function. The implementation of phantom gradient takes the following two steps:
(I 8fe
adopt the phantom gradient [12]
Step Compute the fixed point =* of the DEQ by performing forward iterations until convergence with automatic differentiation turned off.
Step 2 Perform K additional forward iterations from the fixed point x" with automatic differenti- ation turned on.
We provide a PyTocch-style pscudocode for this implementation in Algorithm B in the appendix. Since we need to perform multiple forward iterations to compute the fixed point z², the additional forward iterations in Step only incur small computational overhead. Proposition[]shoow that the derivative with respect to the input data computed by this method converges to the true derivatives as the number of forward iterations K increases. In our experiments, we set K = 2 and find that the models already achieve good performance, and this aligns with findings in training DEQs BI 7.
S K} be generated by performing K. forward iterations from z", i.e., set 20) z" and z 2(++1) fe (k), x) for Then the derivative (z(K)
Proposition 2. Consider function fo(z. x) with fixed point Let (z(x))
matches approximation of derivative of the implicit function x" x; by using K-th order Neumann series of (I JJ- -1, i.e.
1 k K. Let J
0fe
8z ="
dz,
(k)
K-1
05e
Σ J
(10)
k=
5

---

Published as a conference paper at ICLR 2024
0g (R) da,
Hence as K - oo, the derivative
converges to the true derivative, i.e.,
(z(K)
0z(*)
lim K0000
dr;
Similarly, if fo is twice continuously differentiable, the second derivative 0²x(K) dz, at satisfies
02((K) 022(1)
lim K-000 dr,dr, dr,dr,
Again, the regularity assumption in the above propositions aligns with the assumption in score matching based models [6 40). Moreover, as far as we are aware, the second-order results in the above proposition are new in the DEQ context.
4 EXPERIMENTS
In this section, evaluate the performance of the proposed DEQ-assisted score matching models, including score matching variational autoencoder (SMVAH) 40 in Section [4. and noise condi- tional score network (NCSN) [40] in Section 4.4 for generative modeling. We also consider deep kernel exponential families (DKEF) 44 in Section 4.2 and nonlinear independent components es- timation (NICE) [T] in Section for density estimation. In our experiments, we utilize the open
source code from 40 37 and adopt the same training setup for fair comparisons.
4.1
SCORE ESTIMATION IN VAE WITH IMPLICIT DISTRIBUTION
VAE 18] learns latent variable z from the observed data which contains decoder pe( xx) that models the conditional distribution of x given the latent variable 2 and an encoder go (z)x) that approximates the posterior distribution of the latent variable A VAE model is trained by maximizing the following evidence lower bound (ELBO):
L( Φ) Em--parra E [=~ag(=(o) [log log (x]=)) (=)m)( [log (o(x)))]]]]
(11)
Therefore, typical training procedure assumes that qo (=(x) is simple distribution to make the computation tractable. Following [40] Section 6.2.1], the score estimation techniques can utilized to compute the term Vg [log q (2(2(²)) [=)] directly thereby allowing 9 (z|z) (r) to be implicit distribution a distribution without tractable density. In [40]. a variant of VAE proposed that employs score estimator network to facilitate the ELBO computation, and the score estimator
network is trained by minimizing the SSM objective function JSSM in equation 3 we refer to this model SSM VAE; see Table 6 in the appendix for the detailed model architecture.
We consider two image generation tasks: CelebA 24 and C (0) CelebA is a dataset that contains 64 64 x 3 color images that identify celebrity face attributes, and Cifarl0 contains 10 classes of 32 x 32 x 3 color images. Each dataset is split into train, validation, and test sets with 70%, 20%, and 10%, respectively. In our experiments, we examine the performance of SSM VAE regarding its score estimator network depth. We augment the score estimator network with addi- tional convolutional layers and term the model as Aug(k)-SSM VAE, k or 16. We also turn the augmented CNN layer into a (convolutional) DEQ block and term the model as DEQ-SSM VAE. The detailed model architecture is shown in Table 7 in Appendix D33 We train the models using Adam [17, for 10 iterations, with learning rate 1e-4, weight decay 1e-12, and batch size 128.
Figure and Table I compare the performance of SSM VAE, Aug(k)-SSM VAE, and DEQ-SSM VAE in terms of SSM loss, ELBO, FID score [13] and memory usage. The comparisons between SSM VAE and Aug(k)-SSM VAE show that increasing the number of encoding layers improves the model's performance, but it also significantly increases the memory usage and training time. In contrast, DEQ-SSM VAE outperforms the other models terms of SSM loss, ELBO. and FID score while costing the smallest amount of memory overhead compared to the baseline SSM VAE model.
4.2 DEEP KERNEL EXPONENTIAL FAMILIES FOR DENSITY ESTIMATION
DKEF is an unnormalized density estimation model proposed in 44 that parameterizes the unnor- malized log density as log (z) = f(x) + log (xx),, where (o(x)) is a base Gaussian distribution and f is mixture of kernels. Specifically, f (xx EL_ azk(xx, 21), where z; are inducing points, a are the mixture weights, and k(x, ) is the kernel function. The model can be trained with score- matching techniques with loss functions SM, DSM or SSM. Alternatively, the model can the

---

Published as a conference paper at ICLR 2024
DEQ VAE
SSM
Aug(8) -SSM
Aug (16) SSM
S3H Loxs
-303
-3540
FID Scare
8
Memoryyaa
15-
ELBO
-3560
80
-0e3
-3580
40%
Terations
Iterations
Iterations
-2e3
2540
120
5SM Loos
-2550
$100
3+
2500
80
40K
100k
-2570
50k
100k
10k
30c 50k
7ck
100c
Iteations
Ieeraiios
Itarations
Figure 1: Comparing the performance SSM VAF, Aug(k)-SSM VAF, and DEQ-SSM VAE in terms of SSM
loss, ELBO, FII) score, and memory usage. First row: CelebA. Second row: Cifar10.
Method
Dataset
FID. (1UK Her))
FID. (S0k ite))
FID (100K iter)
Parrm
Memory
Trainin time/iter(s)
SSM VAE
96. 23.23-00.50
66.59
61.80++0.56
8.8M
4.0
8.22
SUS(H)SSM VAE
CelebA
79.8710.81
61 20.61
56.02.2.57
15M
37.57
SU VAE
76.43310.2
59. 59.20
54. 64.27:0.48
19M
16GB
73.27
DEO-SSM VAE
74.21=20.2
55.93 55.93 sa
64.17=0.. 48
9.7M
5.4 GB
21.95
SSM VAE
131.74+=0.40
83.93+0.42
77.35-+0.98
7.1M
38
3.97
Aug(8)-SSM VAE
Cifarto
135. (15.28820..0
81.41
70.05
13M
6,3 GB
17.45
A6g VAL
135.18g0.42
27.85
20.77.0.224
17M
GB
DEQ SSM VAE
134.2420.39
77.72 38
68.13_0.8
7.5M
GB
9.98
Table 1: Performance comparison of models for image generation on CelebA and Cifarl0: DEQ-SSM VAE
leads in FID during training, except at 10k iterations where SSM VAE excels. DEQ-SSM VAE shows superior efficiency in parameters, memory, and time over increased depth.
SM loss but with curvature propagation (CP) to approximate the diagonal of the Hessian. We refer to these models as SM, DSM, SSM, and CP, respectively.
We consider the density estimation two datasets: high-dimensional Gaussian. The Gaussian datasets we estimate are 100 and 200 dimensional, re- spectively; each of the datasets is split into train, validation, and test sets with 4860, 600, and 540 samples, respectively. Each of the UCI datasets is split into train, validation, and test sets with 70%, 20% and 10%, respectively. In the experiment, we examine the performance of SM DKEF by creasing the number of layers of the kernel network such as k-Layer-SM DKEF (k- 3, 8, 16). We can also turn the sequence of MLP layers into a DEQ block termed as DEQ-SM DKEF.
|1 and
We set the number of kernels to be 3. For the general score matching function, the feature extractor uses MLP with different number of layers with hidden dimension 30 and sets Softplus as the activation function. For the DEQ variant, we turn the feature extractor into one DEQ block in the form of equation with fully connected weights with dimension 30 and activation function Softplus. The models trained on the UCI and high- sigh--imensional Gaussian datasets are the same except for the latent dimension which is 64 on UCI and 200 on Gaussian. The model trained by Adam [17] for 200/300 epochs with batch size 32/ /188 on the high-dimensional Gaussian/UCI datasets, with learning rate 1e-3, weight decay 1e-12. We follow the procedure in 440, which trains the model using the score matching method with the objective function JSM introduced in Section 2.1
4.2.1 UCI DATASETS:
Table 2 shows the results the three UCI datasets. The results show that DEQ-SM DKEF outper- forms the k-Layer-SM DKEF in terms of SSM loss and likelihood. Moreover, DEQ-SM DKEF is more efficient in terms of memory usage and training time compared k-Layer-SM DKEF. The sults also show that DEQ-SM DKEF outperforms DSM and CP in terms of SSM loss and likelihood.
Parkinson Test LL
WhiteWine Test LL (?)
Metries
SM Loss
PPrmm Mem
SM Loes
Te (f)
PPrrmm Mem
SM Loss
+Param Mem
(.)
(t)
(4)
(1)
(4)
(3)
(
DSM
71.2=1.2
-16-740.9
I0k
-3.8
-14.141.2
17.8
-14.841.0
0.4
CP
23.7
16..4
10c
0.7
12.2
16.3341.5
-0..3+..1
16.1
0.6
SSM
111.94102
-15.140.9
10%
0.5
27.24...9
-13.8+0.7
%
0.5
-33.7+6.2
-14.840.5
Sk.
0.5
SSM-YR
121.147.9
-15.040.8
10k
0.5
27.245.9
-13.810.7
9
0.5
-33.745.2
-14.840.5
%%
0.5
DEQ-SMM DKEFY
143.248.9
-1..920.9
10k
42.945.9
-12.940.9
0.6
48.7 =6.5
-14.2+0.9
06
123.740.6
-15.140.7
10k
1.2
27.2466.9
-13.840.7
1.0
-33.746.2
-14.8+0.8
1.0
8-Laye-5MM DKFF
124.843.5
-16.340.9
268
2.1
29.5
-15.5-20.8
25k
1.6
-34.7
-14.4-20.5
25k
1.6
16- SM DREF
127.644.2
-13.440.8
51k
3.6
-3166+5.2
-12.740.8
50k
2.7
-35.5
-14.5-008
50k
2.7
DEQ-SM DKEF*
152.11147
-12.4_o.7
10k
-44.511.1
-12.210.7
45.011447
-13.91005
1.1
Table 2: Test SM loss and test log-likelihood of different models on the Parkinson, RedWine, and WhiteWine datasets. Mem stands for memory with unit GB.
7

---

Published as a conference paper at ICLR 2024
DEQ SSM DKEF
3-Layer- SN DKEF
8 8-Layer LSN DKEF
16 - Loo-Laye--a SM DKEF
-75
Log-Ilkellhoos
-12
-80
-12
SM Lo55
-100
-13
-155
14
-120
-14
-15
-150
50
100
200
50 100 Epochs
200
-160
Test Log-
Epochs
Test SSM
Likelihood
Memory(GB)
-15.0
Log-likellhood
-12
-20.0
SM Loes
-30.0
-37.0
-13
-30.0
-40.0
-14
-4500
-15
Log-
Epachs
Epochs
Test SSM
Likellhood
Memory(GB)
Log-like (epliikelihood
-13
20..0
-14
-20.0
SM Loss
-30.0
30.0
44.0
-0..0
-15
50
100
200
-14
50
100
200
Test Log- Likelihood
Epochs
Epochs
Test SSM
Memory(GB)
Figure 2: iSsilog-ikclhhhhos/mmmmryy usage of different DKEF models Parkinsom/RedWine/WWiie Wine
(Parkinson: rowl; RedWine: row2: WhiteWine: row3) datasets.
4.2.2 HIGH-DIMENSIONAL GAUSSIAN
For the nigh-dimenniionl Gaussian datasets, the comparison of DEQ-SM DKEF, k-Layer-SM DKHF shown Table 3 and Figure 3 in terms of test loss, test log- lieeliie,oood, Fisher divergence, and memory usage. The results show that the DEQ-based sliced score matching model can save memory while having the best performance the test dataset compared the k-Layer-SM DKEF model.
3-Layer-
L-yayer-
16-Layer-
Divergence
SM Locs
-60
300
Memory/(GB
.65
-70
30
320
-75
100
200
300
100 Epochs
200
300
25
100
200
300
Epochs
Epochs
Figure 3: SM loss/lo--ikellhoodeee likelihood/Fiher divergence on the test dataset, of the 200-dimensional Gaussian, from
diffcrent DKEF methods. The memory used shows the memory for the training process.
Data Dim
Test SM Loss |
Test 1.
Test Fisher-Div
#Param
Memory
Layyerr
.43.00
3.8 GB
8.I.aye-SMM SM DKEF
100- dimension
-43.66
-140.85
6.56
41k
GB
16-Layer-SM DKEF
-44.2
-141.25
13.2 GB
DEQ-SM DKEF
-46.03
-138.86
6.72
22k
3.9
Ground Truch
N/A
-137. 14
0.0
N/A
N/A
3.aayer: DKEF
8-Layer-SM DKEF
200-
-6699
-313.35
33.09
50k
8.7 GB
16-Layer-SM DKEF
dimension
-67.58
-311.84
32.27
14.8 GB
DEQ-SM DKEF
-75.03
-292.56
25.42
33k
5.6 GB
Ground Truth
N/A
-283.79
0.0
N/A
N/A
Table SM loss/log
loss/oo--ikellhood//Feooooooeeee likelihood/Fisher Divergence of different DKEF models high dimensional Gaussian.
4.3 NICE MODEL ON MNIST
Metries MLE
Test SM Loss
Test
#Param
Mem Used
NICE flow model [10 is another den- sity estimate model. Unlike DKEF, its
-579
791
-1694
-1517
DSM
.1)
3035
-4363
log-likelihood are tractable and can be
DSM 1.74)
-97
-8082
directly trained with MLE. Following 40). we adopt NICE [0] as san- ity check for DEQ-assisted score matching. We train it as well as a DEQ variant with SSM loss and compare it with the models such as curvature propagation(CP) (25 DSM 43 ap-
Approx BP
48
-2288
SSM
2450+62
2058+00
15M
6,3 GB
DEQ-SSM
2548
-1766 as
Table 4: SM loss and log likelihood for NICE Models on MNIST: the results of DSM. MLE, CP, and Approx BP results are from 440; SSM and SSM DEQ results are averaged over three random seeds for memory usage.
8

---

Published as a conference paper at ICLR 2024
proximate back- proppppppgttiin (Approx BP) [19] trained with SM loss, also include maximum likelihood estimation (MLE) method as the likelihoods are tractable.
In this experiment, we consider the MNIST dataset [9], which contains 60K 28 28 grayscale images. We set the train, validation, and test datasets with 70%, 20%, and 10% of the whole dataset, respectively. We follow the setup in [40] and use it as baseline model termed SSM NICE. At the core, SSM NICE contains 4 blocks with each containing 5 fully connected layers with 1000 hidden dimensions. Our DEQ-SSM NICE variant turns each of the blocks into fully connected DEQ block of the form in equation 3 with 1000 hidden dimensions. Data are dequantized by adding uniform noise in the range 1/1512 1/512], and transformed using logit transformation, log(z) log(1 .r.. We train the model using Adam for 1000 iterations with a learning rate 1e-3, weight decay le- and batch size 128.
Table 4 compares the performance of DEQ-SSM against SSM and a few other baseline models, showing that DEQ-SSM outperforms SSM both test SM loss and test oikeliikli Meanwhile, DEQ-SSM takes less memory compared to SSM.
4.4 NCSN
We examine our DEQ integration in generative tasks using score-based generative model called NCSN [37] with the training objective of SSM-VR. It is worth noting that our DEQ integration orthogonal to the recent DEQ-based generative models [1, where DEQ is used form the diffusion process in DDIM model [36,, and its training objective does not involve derivatives. We follow the model setting in (37) whose core is 4-cascaded RefineNet [23]. We use the SSM-VR as the objective function in this model and term the baseline model as SSM-VR NCSN. Our DEQ variant replaces the 4-cascaded RefineNet with two residue blocks with one of them turned into
DEQ, we term this model as DEQ-SSM-VE VR NCSN. We also consider a variant of the model with two simple residue blocks, which we term as TwoRes-SSM--V NCSN.
We train the model on the Cifar10 dataset (2D with 128 batch size for 200K iterations using Adam optimizer. For Adam, we set the learning rate as le-4, with weight decay 1e-12 and batch size 128. The results shown in Figure and Table (able[5cooppae the SSM-VR loss, FID score, and memory cost for retaining the computational graph of V,, log (x) for backpropogation, as well as the overall memory cost. The results show that DEQ-SSM-VR NCSN achieves significantly reduced memory footprint and faster computation while maintaining strong performance. In contrast, when the DEQ component is turned off in DEQ-SSM VR NCSN, the resulting model TwoRes PARRes-SS--VE VR NCSN is only marginally faster than its DEQ counterpart, but this came at the expense of notably diminished performance. In Figure in the appendix, we provide samples generated by DEQ-SSM-VR NCSN.
DEQ DSQ-SSM- VR NCSN
SSM VR NCSN
TwoRes SSM VR NCSN
-3370
3oci8
Vallidation
-1380
FiD Score
Memory
15GB
Timet
1.0
13300
-1400
160
32h
40
Joh
32h
22
125
256
512
Training Time
Training
CPUS
(a)
(b)
(c)
(d)
Figure 4: Performance comparison of NCSN models for Cifar10 generation: (a) Test loss Training time,
(b) FID score vs. Time, (c) Memory usage for computational graph retention in backpropagation, (d) Average iteration time vs. GPU count ('OMM' indicates 'out of memory'). ).
CONCLUSION
Model
FID
Time(ms)lter #Param M. of V. log Fie (a: Tot Mem
SSM-VRRCCS.
10.3111.4
871210
28M
7536 MiB
40.4
TwoRes SSM VR NCSN 50.77
391
9M
1701 MiB
18.5
In this work, we integrate DEQs into score matching models to address their mem- ory constraints. We provide a convergence analysis of ap- plying implicit differentiation to higher-order derivatives in the setting of DEQ, By strategically incorporating DEQs into core parts of existing models, we enhance depth and reduce memory re- quirements without compromising their performance. Our empirical experiments across various models, datasets, and tasks demonstrate that including DEQs in score matching not only signifi- cantly reduces memory usage but also improves computational efficiency and performance.
DEQ-SSM-VR NCSN
40.99+1.0
481
9M
2104 MiB
20.1 GB
Table 5:
Comparing NCSN-based image generation models on Cifarl0.
FID score, iteration time, and memory usage for computational graph and total training memory.

---

Published as a conference paper at ICLR 2024
ACCNOWLEDGGEMENT
This material is based on research sponsored by NSF grants DMS-2152762, DMS-2219956, and DMS-2208361 and DOE grant DE-SC0023490.
CONTRIBUTIONS STATEMENT
This paper started from discussion between BW and AO in workshop using DEQ for com- puting Hessian. BW introduced the idea to QW and YH. QW explained DEQ and the phantom gradient technique YH. YH implemented the whole idea, made the system work, and conducted all experiments. QW and YH wrote the first version of the paper focusing on non-experiments and experiments sections, respectively. All authors involved in writing the paper.
ETHICS STATEMENT
This paper introduces memory memorr-efficient approach for score matching using the deep equilibrium model. The proposed methods are validated in the classical benchmark problems, including density estimation and generative modeling. We do not see any potential ethical issues our research.
REPRODUCIBILITY STATEMENT
To ensure reproducible research, we have made the following two major efforts: First, we include sufficient background materials and provide detailed mathematical derivation. Second, we submitted the code in the supplementary materials ensure the experimental results can be easily reproduced.
REFERENCES
[1] CB Allendoerfer. Theorems about differentiable functions. Calculus of Several Variables and Differentiable Manifolds. New York: Macmillan, 1974.
[2] Cem Anil, Ashwini Pokle, Kaiqu Liang, Johannes Treutlein, Yuhuai Wu, Shaojie Bai, J. Zico Kolter, and Roger B Grosse. Path independent equilibrium models can better exploit test-tine time computation. Advances in Neural Information Processing Systems, pp. 7796-7809, 2022.
[3] Shaojie Bai, Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances Neural Information Processing Systems, 32, 2019.
[4] Shaojie Bai, Vladlen Koltun, and J Zico Kolter, Multiscale deep equilibrium models. Advances Neural Information Processing Systems, 33:5238-522 2020.
[5] Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Neural deep equilibrium solvers. Interna- tional Conference on Learning Representations, 2021.
[6] Shaojie Bai, Vladlen Koltun, and Zico Kolter. Stabilizing equilibrium models by jacobian regularization. International Conference on Machine Learning, pp. 554-565. PMLR, 2021.
[7] Shaojie Bai, Zhengyang Geng, Yash Savani, and Zico Kolter. Deep equilibrium optical flow estimation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 620-630, 2022.
[8] Jérôme Bolte, Tam Le, Edouard Pauwels, and Tony Silveti-Falls. Nonsmooth implicit differ- entiation for machine-learinng and optimization. Advances in neural information processing systems, 34:1 111111333315544 2021.
[9] Li Deng. The mnist database of handwritten digit images for machine learning research [best of the web]. IEEE signal processing magazine, 29(6): 141-142, 2012.
[10] Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear independent components estimation. arXiv preprint arXiv: XXiiv:411..8 1410.85 16, 2014.
10

---

Published as a conference paper at ICLR 2024
[11] Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. //archive. ics. uci edu/ml
URL
http:
[12] Zhengyang Geng, Xin- Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin. On training implicit models. Advances in Neural Information Processing Systems, pp. 24247-24260, 2021.
[13] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[14] Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Ad- vances in neural information processing systems, 33:6840-6851, 2020.
[15] Zhichun Huang, Shaojie Bai, and Zico Kolter. Implicit²: Implicit layers for implicit repre- sentations. Advances in Neural Information Processing Systems, 34:9639-9650, 2021.
[16] Aapo Hyvärinen. Estimation of non- non-oomallizd statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[17] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv; 1412.6980, 2014.
[18] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.611 1312.6/ 14, 2013.
[19] Durk P Kingma and Yann LeCun. Regularized estimation of image statistics by score match- ing. Advances in neural information processing systems, 23, 2010.
[20] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[21] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[22] Serge Lang. Real Analysis. Addison Weelly, 1983.
[23] Guosheng Lin, Anton Milan, Chunhua Shen, and Ian Reid. Refinenet: Multi Multi-path refinement networks for high-resolution semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1925-1934, 2017.
[24] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pp. 3730-3738, 2015.
[25] James Martens, Ilya Sutskever, and Kevin Swersky. Estimating the Hessian by back- propagating curvature. arXiv preprint /206.
2012.
[26) Chenlin Meng, Yang Song, Wenzhe Li, and Stefano Ermon. Estimating high order gradients of the data distribution by denoising. Advances in Neural Information Processing Systems, 34: 2021.
[27] Avik Pal, Alan Edelman, and Christopher Rackauckas. Continuous deep equilibrium models: Training neural odes faster by integrating them to infinity. arXiv preprint arXiv:2201.122200 2022.
[28] Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu. Efficient learning of generative models via finiteddifference score matching. Advances in Neural Infor- mation Processing Systems, 33:19175-19188. 2020.
[29] Matthew Parry, Philip Dawid, and Steffen Lauritzen. Proper local scoring rules. The Annals of Statistics, 40(1):561-592, 2012.
[30] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin. Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performrfrre deep learning library. Advances in neural information processing systems, 32, 2019.
11

---

Published as a conference paper at ICLR 2024
[31] Ashwini Pokle, Zhengyang Geng, and Zico Kolter. Deep equilibrium approaches to diffusion models. Advances in Neural Information Processing Systems, 35:37975-37999 2022.
[32] Zaccharie Ramzi, Florian Mannel, Shaojie Bai, Jean-Luc Starck, Philippe Ciuciu, and Thomas Moreau. Shine: Sharing the inverse estimate from the forward pass for bi-level optimization and implicit models. In ICLR 102-International Conference on Learning Representations, 2022.
[33] Martin Raphan and Eero Simoncelli. Learning to be bayesian without supervision. Advances in neural information processing systems, 19, 2006.
[34] Martin Raphan and Eero Simoncelli. Least squares estimation without priors supervision. Neural computation, 23(2):374-200) 2011.
[35] Saced Saremi, Arash Mehrjou, Bernhard Schölkopf, and Aapo Hyvärinen. Deep energy esti- mator networks. arXiv preprint 18Xiv 08306, 2018.
[36] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. lernational Conference Learning Representations, 2021. BtLLhhtps://oeeeee
net/ wtt///orrmm//
[37] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data dis- tribution. Advances in neural information processing systems, 32, 2019.
[38] Yang Song and Stefano Ermon. Improved techniques for training score-based generative mod- els. Advances in neural information processing systems, 33: 13:12-28-1 2020.
[39] Yang Song and Diederik P Kingma. How to train your baaeyy models. arXiv preprint arXiv:2101.03288. 2021.
[40] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach density and score estimation. In Uncertainty in Artificial Intelligence, pp. 574 584. PMLR, 2020.
[41] Yang Song. Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34: 1415-1428, 2021.
[42] Russell Tsuchida and Cheng Soon Ong. Deep equilibrium models as estimators for continuous latent variables. In International Conference on Artificial Intelligence and Statistics, pp. 1646 1671. PMLR, 2023.
[43] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-165 2011.
[44] Li Wenliang, Danica J Sutherland, Heiko Strathmann, and Arthur Gretton. Learning deep kernels for exponential family densities. In International Conference on Machine Learning, pp. 6737-6746. PMLR, 2019.
[45] Ezra Winston and Zico Kolter. Monotone operator equilibrium networks, Advances in neural information processing systems, 33:10718-10728, 2020.
[46] Kösaku Yosida. Functional analysis. Springer Science & Business Media, 2012.
12

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 MISSING PROOFS IN SECTION
In this section, we provide the missing proofs in Section
Proof of Proposition L Since the function (z, 2) is contraction mapping, the Jacobian matrix 8 is
ffe has its magitudes of the eigenvalues less than one. Therefore, the matrix 0s
8=
invertible. Then the implicit function theorem (e.g. Theorem 2.1 in Section 2 of [22) implies that in a neighborhood of x=, x² is a twice differentiable function of æ. Then we can differentiate through fixed point equation z= fo(=", x)) to obtain
Oz 0 f
Ofe
Oz
0z, Dr,
Oz
dz:
Oz
0
-1
0.e
0m:
Oz
0x:
2"
Similarly, we can differentiate through the fixed point equation 2" = fo(2",
twice to obtain
822²
0fe
f2
fe
0z
dr,dr, +
0z,0z;
Oz
dz,Ōr, 02 fo Ozdr,
Ozdr,
dz,
Oz 0 fe
Oz Oz
Ōz, dzdz
Ox; dx,
a5e
80 0
8z
I
Oz
On,,Or, 1- Ox,
Ozdr,
x Oz:
8 fo Oz
82 f0
0z 0z"
+
Ozdr,
Ozōz
Oz, dx,
This concludes the proof.
Proof of Proposition
We prove this result by induction on K. For K 02(1)
have
0.e
0
Oz (0)
0
dz;
dx;
dz
Ōr, 0z: =
2(0)
20)
This proves the base case. Now assume that the statement holds for - 1, we have
Oz(K)
0
02(K-1)
Ōx:
0x,
Oz
Oz,
(K-1)
11-11
0fe
0fe
K
i f
k 0
Ox.
0z
0z
Dx.
K-1
k a
Oz
dx,
k=0
In particular, when K - 00, we have
Oz*
k 0
0.e
-1
a
Dz,
Oz
Ox:
=
Oz
dx,
k==
This concludes the proof of the first part.
Now move to the second part, using the contraction of tensors, will show that there is
(k)
fe
82
fe
AK dr,Ox,
+ B
dx,0s;
Ozdr,
Ox:
8S0
010
8
0
0
+ BR Ozdr,
Cx
dr,
Oz0z
dr.
0z,
13

---

Published as a conference paper at ICLR 2024
where the matrices Ax, Bx, Cx satisfy the following recursive equations:
0fo
Aк
Ak-1 +I,
0z
050
K-2
0fe k
Bк
Bk BK-14 +
0z
Oz 1
km0
050
K-2
050
K-2
050
)))
Cк
CK-1+ +
dz
Oz
Oz
z''
k=
A==
with initial conditions Ao I, Bg
0, Co
0. We prove this by induction on K. For K = 1, we
have
(22
0fe
8z
f0
80
Oz
Or,,Or,
Oz
2(0) dr, Jr
dr,On, 210)
Ozdr,
100 Ōr,
8 dzdz, 2(0) dr, Ozdz
θz(0) se
Oz 0z
+
dr.
dz,
200)
S0
Or,0rr
This verifies the base case.
Now assume that the statement holds for K 1, we have
(K)
θfe
(21(-1))
fe
f0
1x(
+
0x;0x;
Oz
x(K-1) 0x,dx,
dr,ōs 1(x-11
0zdr,
Ōx:
z)K-1)
S0
0z(K-1 1)
8² fo
0z(K (z(K-1)
+ Ozdr:
OzOz
Oz.
1(K 1)
2(K 1)
θfe
8 fo dx.0xr
fe
fe
Ak
+ +B Ozdr;
0z z(K-1)
Ox:
x1(
*(K 11
02 fo
Ofe
8
afe Ofo Dx, Oz,
+BK Ozdx;
+ Ck
Ox,
Ozdz (1-
(K--1)
82 fe + 0r,0u,
8 fo
R-2 Σ
0.e
fe
+ dzdxz
Oz z(K-1)/
Øx: (2)
k=0
f0
Ofo
Ofo
+ Ozdx,
K=O
Oz
(((-1))
(x(-1)
f0
K-2
0
R-2
0
0
Ozoz
Oz
Ox:
az
Dx, z(k-1>
A==
2(K
1)
2(K
By collecting terms and using the fact that the value of 2(-1)) equals z", verify that the matrices
Ak, BK, CK satisfy their respective recursive equations. Since the matrices 0f. is contraction
8z
mapping, the recursive equations imply that the matrices Ak, Bx, CR converge to the following limits:
Ofo
-1
lim Ak K0000
0z
""
050
lim
BR
I
K
Oz
0
lim
Ck
I
K-00
Oz
14

---

Published as a conference paper at ICLR 2024
Hence, we have the following limit:
02 fe Oz 12.)
(I
fo fo
Oz
lim K0000 0z,dz
0x, x
dzdz, x" Ōz;
8f fo Oz 00i0
Oz 0z
Ozor,
0z.
Ozōz
0r. dx,
dr
This concludes the proof.
B IMPLEMENTATION OF DEQ
Algorithm Implementation the fixed point iteration in PyTorch-style pseudocode.
1: procedure FIXIDPIINTITRATIONG K)
2: with torch.no torch.no_grad():
z' = RootSolver(fo.22 4: (0) ="
3:
5: for = to 10 K do
6: 2k++1 (x)
7: end for
8: return (K))
9:
end procedure
C EMPRICAL ANALYSIS OF PHANTOM GRADIENT ERRORS
In DEQ integration, without any additional damping. This is because we find that this simple implementation is sufficient for the DEQ integration to work well practice. In this section, we provide additional empirical analysis of the phantom gradient error in the DEQ model. We follow the setting of the experiments in the DKEF model Section 4.2 and use the Pakinson dataset for the experiments. We train the model with SM and SSM loss respectively and at each training step, we compute the score function using the exact gradient for the DEQ block and the phantom gradient with K 2. We also compute the loss gradient with respect to the parameters used in the two cases.
simple implementation of the phantom gradient with K = 2 and
C.I ANALYSIS OF ERRORS IN THE SCORE NORM
We first consider the score norm error, which defined as the relative error of the scores obtained from outputs of the DEQ with K 2 using phantom gradient and DEQ with true gradient computed through implicit differentiation. We use 8 to denote the score from the exact gradient for the DEQ block and to denote the score from the phantom gradient with K 2. The relative error is defined as
$
err
100%.
||s|
Figure 5 we show the distribution of the score norm error from training the same model on the Pakinson dataset with SM and SSM losses, respectively. In both cases, the score norm error small, which indicates that the phantom gradient is a good approximation of the exact gradient. Additionally, in 6 show the corresponding distribution of the cosine similarity. The results demonstrate that the scores obtained from the exact gradient and the phantom gradient with K are highly correlated.
C.2 ANALYSIS ON THE PARAMETER GRADIENT COSINE SIMILARITY
In this section, we consider how the enror stems from the phantom gradient affects the optimization process. To this end, consider the cosine similarity of the gradient of the loss function with
15

---

Published as a conference paper at ICLR 2024
4
4
Density
Density
2
2%
4%
6%
2%
4%
6%
Score Norm Err (SM)
Score Norm Err (SSM)
(a)
(b)
Figure 5: (a) and (b) show the distribution of the score norm error from SM and SSM, respectively. The rela- tive error concentrates around zero, indicating that phantom gradient with = good approximation of the exact gradient.
6
Distribution
Density
4
4.
2
0.98
0.99
0.99 Score Cos Similarity (SSM)
0.98
Score Cos Similarity (SM)
(a)
(b)
Figure 6: (a) and (b) show distribution of the score cosine similarity from and SSM respectively. The scores obtained from the exact gradient and the phantom gradient with K highly correlated.
respect to the parameters from the exact gradient in the DEQ block and the phantom gradient with K = 2. The result is shown in Figure 7 The results show that the cosine similarity is close to one, which further indicates that the phantom gradient is good approximation of the exact gradient.
10
8
Density
Density
2
4.
2
0.94
0.96
0.98
1.00
0.6
0.8 Grad Cos Similarity (SSM)
1.0
Grad Cos Similarity (SM)
(a)
(b)
Figure 7: (a) and show the distribution the cosine similarity of the gradient of the loss function with
respect the parameters from the exact gradient in the DEQ block and the phantom gradient with K 2 from SM and SSM, respectively. cosine similarity shows that loss gradient from the phantom gradient highly correlated with the loss gradient from the exact gradient.
16

---

Published as a conference paper at ICLR 2024
D ADDITIONAL DETAILS ON EXPERIMENTS
D.1 DKEF MODEL
The deep kernel exponential families (DKEF) model parameterizes the unnormalized density as log po(x) = f()) + log (0(x), where (o(x) is a base distribution and f(x) is deep kernel expo- nential family. In our implementation, the deep kernel exponential family is defined as
f(x)
E ayk(z, 21)
Jema
where induced points, aj mixture weights. The kernel function k(æ,, x1) is defined
kax, 21) = Σ Pr exp (- (-20² (x) 0r(2))11²)
Here, (-(-)) is the r-th feature extractor, is the mixture coefficient, and σ, is the length scale. We the DEQ layer model the feature extractor (.).) of the form equation 5 with fully connected weights and Softplus activation function. Additionally, (3) in this case is single layer MLP also with Softplus nonlinearity.
D.2 NICE MODEL
Nonlinear independent components estimation (NICE) is a flow-based density estimate model that has four coupling layers with each layer containing five dense hidden linear layers with Softplus activation function. The DEQ model is used to replace the five dense hidden linear layers in each coupling layer to be of the form equation with Softplus activation function. The function gyy) in this case is single layer MLP with Softplus nonlinearity.
D.3 VAE MODEL.
We include the architecture used in the VAE model in Table 6 The model is based on the model in [40] with the addition of the [Augmented Pre-processing Layer] in the score estimator. The specific configuration of the [Augmented Pre-processing Layer] is shown in Table
D.4 NCSN MODEL
Noise conditional score network (NCSN) 37,38] 38 is a score- based generative model that learns the score function from data and then uses the score function to generate samples. The score function is parameterized as a neural network. In our implementation, the DEQ residual block contains two components that form map z -> h(z) }(y) where h is (normalized) residue block with each convolution weight has its Frobuinus norm constrained to ensure contractiveness and g(y) is ResBlock that transforms the input data, in the original NCSN implementation (37], See Table and Table for details.
E SAMPLED IMAGES GENERATED BY DEQ-ASSISTED GENERATIVE MODELS
In this section, provide the sampled images generated by DBQ-assisted models in Section 4 This includes the images generated by DEQ-SSM VAE model on the CelebA dataset (Figure and the images generated by DEQ-SSM based NCSN model the Cifar 10 dataset (Figure
F CONVERGENCE OF THE FIXED-POINT SOLVER
In experiments, utilize the simple fixed point iteration (Picard iteration) to solve the fixed point of the DEQ layer. This is because we aimed to attribute differences directly to DEQ integration, not solver efficiency. In this section, provide the convergence results of DEQ-based models in
17

---

Published as a conference paper at ICLR 2024
Name
Configuration
5 conv; channels; stride 2 x 2; padding 2; ReLU
5 X conv; 2m channels; stride 2 x 2; padding 2; ReLU
Implicit Encoder
5 conv; 4m channels; stride 2 2; padding ReLU
5 x conv; Sm channels; stride 2 x 2; padding 2; ReLU
512 Dense: ReLU
D2 Dense
Dense; ReLU
5 X 5 conv T; 4m channels; stride x 2: padding 2: out padding 1; ReLU
Decoder
5 X 5 conv ; 2m channels; stride x 2; padding 2: out padding 1; ReLU
5 x 5 conv
channels; stride 2 x 2: padding 2: out padding 1; ReLU
x 5 conv T: channels; stride x 2: padding 2: out padding 1; Tanh
Concat|æz, Softplus(Dense())))
[Augmented Pre-processing Layer]
5 conv; m channels; stride 2; padding 2; Softplus
5 5 conv; 2m channels; stride 2 x 2; padding 2: Softplus
Score Estimator
5 conv; 4m channels; stride 2: padding 2: Softplus
5 5 conv; 8m channels; stride 2 2; padding 2: Softplus
512 Dense; Softplus
D₂ Dense
is hyper-parametee be set define the number of channels of cach convolution hidden layer. e denotes the number of channels of the original image data. D₂ is the dimension of the latent space. The [Augmented Pre pre-processing Layer] contains layers in addition to the base model in 40 that augments the input data. The specific configuration of these layers is shown in Table
Table 6: The basic model of cach part of the VAE model used in Scction 4.1
Model
Specification
Sequential CNN
concat Ø. Softplus(Dense(e)) I: 5 conv; channels; stride x 1; padding 2: Softplus
(n-layers)
2: 5 x conv; m channels; stride x 1; padding 2: Softplus
:5 x conv; m channels; stride x 1; padding 2: Softplus concat[æ, Sofiplus(Dense())) f(z) Softpouuuuu W g(y)
DEQ
Table 7: Specifications of the [Augmented Pre-processing Layer] in the score estimator of the VAE model 6 The input transformation g(y) in the DEQ shares the same architecture as one layer Sequential CNN. Meanwhile, the weight W is the linearized version of such convolutional layer.
Figure 8: Images generated by DEQ-SSI SSM VAE model on the CelebA dataset.
Figure Images generated by DEQ-SSM based NCSN model on the Cifar10 datasct.
Section In Table 10 we present the numerical tolerance using absolute error in I-znnorm between two consecutive iterations, the maximum number of iterations, and the averaged number of iterations observed in the experiments.
18

---

Published as a conference paper at ICLR 2024
CAM-VRNCS
DEQ-SSM-VR NCSN
3 x Conv2D,
No: 3 > 128
3 x 3 Conv2D,
Nc: 3 - 128
ResBlock, Ne: 128 -> 128
ResBlock, No: 128 - 256
ResBlock, Ne: 128 - 256
DEQ-ResBlock, Nc: 256 - 256
ResBlock down, Nc: 256 - 256
ResBlock, Ne: 256 - 256
ResBlock, No: 256 - 256
ResBlock down,
Nc: 256 + 256
ResBlock down, Ne: 256 - 256
ResBlock,
No: 256 - 256
ResBlock, Nc: 256 - 256
RefieBlockk, No: 256 - 256
ResBlock down, Nc: 256 - 256
RefineBlock, No: 256 128
ResBlock, No: 256 - 256
3x3 Conv2D,
No: 128 -
RefineBlock, Ne: 256 256
RefineBlock, Ne: 256 + 128
RefineBlock, Ne: 128 + 128
RefineBlock, No: 128 128
3x3 Conv2D,
Nc: 128 + 3
(a)
(b)
Table The architectures of SSM-VR NCSN (a) vs. DEQ-SSM-VR NCSN (b). Ng denotes the number
channels.
DEQ-ResBlock
g(y)
f(z) == h(z) + g where (1)2) is (esBlocknnormaized)) with Nc: 256 - 256
3 3 Conv2D, No: 3 + 128
ResBlock, No: 256 + 256
Table 9; The architecture of first ResBlock in (b) Table that transforms the input data to the hidden variable. No denotes the number channels.
in the DEQ-SSM-VE NCSN. g the first CNN layer and the
Model
Absolute Error
Max Iterations
Average
Error
Iterations
DEQ-SM DKEF DEQ-SSM DKEF DEQ-SSM VAE (CelebA)
5e-6
128
12.8
5e-6
128
14.5
le-5
128
18.9
DEQ-SSM VAE (Cifar10) DEQ-SSM NICE DEQ-S BEQ-SSM NCSN
le-5
128
15.2
5e-6
128
11.8
1e-5
128
16.6
Table 10: List of corvergence-related parameters, including absolute error, the maximum number of iterations,
and the average number of iterations observed during experiments. The results show that by constraining the Frobenius norm Softplus activation function and also reusing the learned fixed point from the previous forward pass, we often obtain good convergence results.
linear map in the DEQ layer as well as the small contractive factor near zero
19


Reviews:
[{"review_id": "2kwKefwvMD", "reviewer": "Reviewer_H3Py", "summary": "Score matching methods attempt to learn probability distributions by computing a loss on the derivative of the log likelihoods, thereby circumventing the need to compute the normalising constant. Unfortunately, they are time and particularly memory intensive, on account of requiring to compute the Hessian during backpropagation. This paper uses deep equilibrium models (DEQs) to improve time and memory consumption. Experiments show that with the same computational and memory budget, DEQ models can outperform non-DEQ models.", "questions": "- As far as I understand, one introduces a few DEQ layers (equation 5), which compute some candidate fixed point of the score function, which is then passed to the actual score function. Is that correct? How then does the DEQ (5) become close to the true score function? Are you just relying on the Fisher divergence loss? \n- Is the following paper related at all? They compute fixed points of the score function using a DEQ, but for a rather specific model some kind of generalisation of PCA. Deep Equilibrium Models as Estimators for Continuous Latent Variables, AISTATS 2023\n- It is mentioned that the Hessian is memory intensive. However, can't the HVP be computed without ever having to form the Hessian? This would then just be linear rather than quadratic in the number of parameters, without even having to approximate. I can understand time intensive, but not necessarily memory intensive. This is discussed somewhere in the third paragraph of the introduction and section 3, but could the authors please further clarify?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This is overall a good quality paper.\n- The paper cites an appropriate amount of related literature and places itself well within the literature.\n- The use of phantom differentiation is nice (mentioned in section 3.4). Compute the forward iterations of the FP solver with auto-diff off, then turn it on and run K more steps. \n- Experiments are convincing and broad, and show that DEQ models consume less memory and achieve better predictive performance than some existing methods. (Contribution +)\n- The paper is moistly clearly written. The illustrations are clear and informative. Tables are presented well. (However I do have questions about section 3.1, as below, hence the low \"Presentation\" score). (Presentation ++)\n- The mathematics and formulation are mostly easy to understand, and they appear to be okay. (Soundness ++)\n- This paper solves an important problem of speeding up training in score-based methods (Contribution ++)", "weaknesses": "I am willing to raise my score if the authors can respond to my queries below.\n\n- My biggest concern is related to my first Question in the box below. Why is this method actually principled? As far as I understand, a DEQ model is used as a kind of surrogate for the true score function, and then the fixed point of this surrogate is plugged into the score function and the Fisher Divergence is minimised. Does this result in the DEQ model being somehow close to the true score function, and in particular, the fixed point of the DEQ becomes close to the fixed point of the score function? I would appreciate it if the authors could spend some more time describing the method itself (i.e., I feel section 3.1 is too short). (Presentation --)\n- Proposition 2. It seems to require $K \\to \\infty$, however the authors state they use $K=2$. There is no quantification of the error in finite $K$ settings. Furthermore, it is not clear how error in the derivatives translates to errors in the score, or errors in the Fisher divergence. (Soundness -)\n\nMinor:\n- The first sentence in the abstract doesn't seem to make sense, however I could be wrong. Double check with someone else. I would write \"Score matching methods, *which* estimate probability densities without computing the normalization constant, are particularly useful in deep learning\"\n- I believe there is an incorrect statement in section 3.2. \"The well-posedness of the fixed-point equation 5 can be guaranteed by constraining the Frobenius norm of the weight matrix W to be less than one\". This would be true if the activation were contractive (a property which is later mentioned in two sentences), however it is not necessarily true in general. One needs the weight matrix to have a norm less than 1 AND to have a contractive activation.\n- I am not sure what the value is in providing Figure 2 in the main paper. First, it is entirely qualitative. Second, it is only for the newly introduced DEQ model and not for the other models. In my opinion this could be left in the appendix, with a comparison with the other methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NSHiMYpu5g", "reviewer": "Reviewer_CgBT", "summary": "Training models with score matching objective can be computationally expensive as it involves computation of higher-order derivatives like Hessian. Therefore, alternate objectives like sliced score matching (SSM) and denoising score matching (DSM) were proposed.\n\nThis work proposes to use deep equilibrium model (DEQ)-based architectures to train networks with different score matching objectives in order to reduce the memory requirements. DEQs use implicit gradients which reduce the memory footprint of backward pass as it does not require creation of an explicit computation graph for all the layers of the network. However, computing implicit gradients for DEQs need  matrix inversion. This work uses phantom gradients to circumvent matrix inversion. \n\nThe benefits of using DEQ-based architecture for SSM and DSM has been shown on different density estimation tasks with NICE and deep kernel exponential family (DKEF) models. Further, benefits of using this architecture is also shown for generative modeling with VAEs and NCSN. Overall, the experiments show that using DEQs leads to lower ELBO and test SSM loss, as well as reduced memory requirements.", "questions": "1. Could the authors clarify which tasks in the experiments sections use a DEQ model that needs computation of Hessian? As far as I can tell, only DKEF trained on UCI datasets uses the original score matching loss. Besides this, all other tasks use sliced score matching (SSM) loss or denoising score matching loss, both of which do require computation of Hessians. In that case, could the authors state the memory requirements of the networks in Table 2. Also, could the authors add these numbers for SSM and SSM with variance reduction (SSM-VR) in Table 2? \n2. There are several missing details in the tables of experimental sections:\n    1. The size of models in terms of number of parameters should be indicated in Tables 2, 4, and 5.\n    2. The implementation details of DEQ-SSM NICE state — “Our DEQ-DDM NICE variant turns each of the blocks into a fully connected DEQ block of the form in equation 5 with 1000 hidden dimensions.” It is unclear what g(y) is in this case. \n    3. Please include the size of CelebA images used in Section 4.1", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "As indicated by results in Figure 1-5 and Tables 1-5, DEQs outperform or match the performance of non-weight tied architectures for optimizing with SSM objectives on density estimation tasks involving NICE and DKEF architectures, as well as on score estimation for VAEs and NCSN. Further, they consistently have lower memory footprint. The plots also seem to indicate that DEQ converge faster and have lower training and test loss, while matching or outperforming perceptual quality of generated images in terms of FID.", "weaknesses": "- The paper is very sparse on various implementation details of DEQ. Some of the important details that are missing are:\n    - The choice of fixed point solver is unclear. From the code shared in the supplementary, it seems like fixed point iterations were used. However, there are other fixed point solver methods (e.g. Newton methods, Anderson acceleration, Jacobi, Gauss-Siedel etc.) that can converge faster, and it is unclear why simple fixed point iterations were preferred over these other faster methods. I presume that memory footprint could be one of the reasons. The reasoning behind various design and implementation should be explicitly stated. Further, for each task, the number of fixed point iterations should be stated as this is an important hyper parameter. \n    - The original work on phantom gradients (Geng et al. 2022) implements these gradients as a convex combination with parameter $\\lambda$ i.e. $z^\\star = \\lambda f(z^\\star, x) + (1 - \\lambda) z^\\star$. However, it seems that this paper doesn’t follow this as per Algorithm 1. Is there an empirical evidence of optimality for the choice of  $\\lambda=1$? Further, the paper also skips on values of hyper parameter $K$ used for phantom gradients for different tasks and datasets.\n    - The paper is very sparse on implementation details of different network architectures. In general, choice of g(y) in Equation 5 is unclear for different model architectures for instance in Table 7. Similarly, DEQ architectural details, especially choice of g(y) are missing for other architectures.\n- There is limited empirical analysis on the nature of convergence for different architectures. We don’t know how convergent the fixed points are and if there are additional benefits of using additional test-time iterations of the DEQ architecture to squeeze improved test-time performance. It will useful to report the final values of the absolute difference between consecutive iterates of DEQ $\\|| z^{t} - f(z^{t}, x)\\||_2$ or relative difference measure as $\\frac{\\|| f(z^t, x) - z^t \\||_2} {\\|| z^t\\||_2}$. Even reporting the final values of these quantities at $z^\\star$ is useful.\n- In the discussion in Section 3.2, it is important to point that most of the practical DEQs are not globally convergent — neither uniqueness not existence of the fixed point is guaranteed--- and thus they are not well-posed. This applies to the DEQs used in this work as well. Further, it is possible to converge to limit cycles as noted in [2][3]. This is the motivation behind deep equilibrium models like MonDEQs [1] which have explicit guarantees on uniqueness and existence of a fixed point.\n- This paper derives closed form expressions for second order implicit gradients used for backprop in DEQ, but empirical analysis of stability and performance of these second order implicit gradients is missing. \n- Scalability analysis for score matching is missing. As DEQs have constant memory requirements, these models should be more scalable compared to non-weight tied models as the dimension of data increases. It would be useful to recreate a figure similar to Figure 2 in Song et al. 2020 [4].\n- I am not sure if a fair comparison is being made between non weight-tied and DEQ models across all the tasks. Ideally, one should equalize either for the number of parameters or FLOPs. However, these numbers are not included for all the tasks/results in the experiments section.\n\n[1] Winston, Ezra, and J. Zico Kolter. \"Monotone operator equilibrium networks.\" Advances in neural information processing systems 33 (2020): 10718-10728.\n\n[2] Bai, Shaojie, J. Zico Kolter, and Vladlen Koltun. \"Deep equilibrium models.\" Advances in Neural Information Processing Systems 32 (2019).\n\n[3] Anil, Cem, et al. \"Path Independent Equilibrium Models Can Better Exploit Test-Time Computation.\" Advances in Neural Information Processing Systems 35 (2022): 7796-7809.\n\n[4] Song, Yang, et al. \"Sliced score matching: A scalable approach to density and score estimation.\" Uncertainty in Artificial Intelligence. PMLR, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "l8PeGmHXCQ", "reviewer": "Reviewer_hmh8", "summary": "**Edit: Score increased after rebuttal.**\n\nThis paper proposes the use of deep equilibrium layers (DEQs) for the task of score-matching. The motivation for doing this is that score-matching techniques are typically expensive (in both time and memory) as they require computing gradients and/or Hessians of the model, due to the need to store the entire computational graph in memory for backprop. In contrast, DEQ methods do not need to store the entire computational graph. In Section 3, Propositions 1 and 2 give results on several derivatives needed to train score-based models with DEQ layers. Section 4 provides a fairly thorough empirical evaluation of the proposed methods across a variety of tasks, including density estimation and generative modeling. The proposed methods generally show comparable or better performance than the baselines, with low memory costs.", "questions": "See weakness section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Score-matching techniques are a fairly hot topic, and this paper aims to solve an important issue with these models, namely their high memory costs.\n- The experimental methodology is generally sound, and the provided evidence generally supports the claims made throughout the paper.\n- The proposed method shows consistently favorable performance in comparison with the considered baselines, while still retaining lower memory costs.\n- To the best of my knowledge, this is the first work to use DEQs in the context of score-matching for density estimation tasks.", "weaknesses": "- The main weakness of the paper is that the technical contribution of the paper is somewhat limited. In particular, the authors apply standard DEQ modeling techniques to the task of score-matching. While this combination is novel, the techniques are a relatively straightforward application of existing DEQ methodology. Moreover, DEQ models in the context of diffusion models has been previously explored [2].\n- It was unclear in some places what was novel in this work and what has been derived previously. For instance, in Proposition 1, Equation 8 has certainly already appeared in the DEQ literature (Theorem 1 of [1]), but it was unclear to me if the second-order results (Equation 9) was novel. Similarly, Proposition 2 (at least in part) appears in [3], and it was not entirely clear to me what aspects of this were novel.\n- The results in Section 4.2.1 were a little confusing to me.\n     - What exactly are DSM, CP, and SSM referring to in Table 2? My understanding is that these rows correspond to some version of the DKEF model fit with these score-matching techniques, but the details were unclear. \n     - Moreover, the acronym \"CP\" is used here before it is introduced later in the paper. \n     - I would have also expected Table 2 to include memory usage (similar to Table 3) as this is one of the key motivations for the proposed method.\n     - The description of DKEF itself was also unclear -- what is $q_0$? Where are the learnable parameters $\\theta$ appearing?\n- Given that there exists previous work on DEQ diffusion models [2], I might have expected a comparison with the methods in [2] in Section 4.4, or at least some rationale for not doing so.\n\n### Minor Comments\n- In Section 4.1, \"encoder\" should be $q_\\phi(z | x)$ and \"decoder\" should be $p_\\theta(x | z)$\n- Typo in the caption of Figure 3: \"momory\" should be \"memory\"\n\nIn general, I found this paper to be well-motivated with thorough and convincing empirical results, but perhaps limited in its novelty.\n\n### References\n\n[1] [Bai et al., Deep Equilibrium Models](https://proceedings.neurips.cc/paper_files/paper/2019/file/01386bd6d8e091c2ab4c7c7de644d37b-Paper.pdf)\n\n[2] [Pokle et al., Deep Equilibrium Approaches to Diffusion Models](https://arxiv.org/pdf/2210.12867.pdf)\n\n[3] [Geng et al., On Training Implicit Models](https://arxiv.org/pdf/2111.05177.pdf)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Dky2fVU4tz", "reviewer": "Reviewer_YwTh", "summary": "The paper proposes a novel score-matching method that uses deep equilibrium networks (DEQs) to overcome the memory bottleneck usually associated to score-matching.\n\nThe main contributions of the paper are the proposed method involving DEQs, a theoretical analysis involving implicit differentiation for higher-order derivatives, and an empirical validation that compares the DEQ model with some other methods for solving score-matching problems, like sliced score-matching.\n\nThe theoretical analysis consists of two propositions. The first proposition is concerned with computing the first and second derivatives of the solution mapping of the fixed point equation used in the DEQ and these formulas involve inverting a matrix (typical for implicit differentiation methods). The second proposition is focused on approximating the derivatives of the solution mapping formulated in Proposition 1 without inverting any matrices, by using a truncated Neumann series approximation.\n\n**Edit: after the rebuttal by the authors, the assumptions on the types of activation functions considered have been strengthened and my doubts about the theoretical contributions have been addressed. I therefore raise my soundness score and overall score to a 6.**\n\nThe empirical validation is split into four parts, each centered around a task that can be solved by score-matching methods. In each of these tasks the frugality with respect to memory is validated without showing loss in accuracy - the plots suggest that the most accurate methods in terms of various metrics (Frechet Inception Distance, slice score-matching loss, ELBO, etc) is often the DEQ method, although it is not often the fastest method nor is it always the most memory frugal.", "questions": "Are there any conditions you can impose on f_theta that will ensure that z* is twice differentiable that don't also rule out using non-smooth functions like ReLU in the DEQ?\n\nDEQs and the phantom gradient method in section 3 are already well-known, what is the contribution here besides plugging them into score-matching?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The novelty and empirical validation of the claims in the paper seem worthwhile. I don't know of any methods that combine DEQs with score matching, although DEQs and the phantom gradient method of section 3 are already well-known and studied in the papers that introduced them. It is also clear from the numerical experiments that this performs well empirically and that the memory overhead is not so large, especially compared to augmented sliced score matching methods.", "weaknesses": "The theoretical parts of the paper are not correct.\n\nFor instance in proposition 1 it is required that the function f_theta is continuously differentiable; this is not the case for the ReLU and will necessitate a non-smooth implicit function theorem (see, e.g., [Bolte 2021]). There even appears the second derivative of f_theta but this is not compatible with the stated assumptions.\n\nThen in Proposition 2 there is no justification that z* is going to be twice differentiable (and in the examples using ReLU, it might not even be differentiable). If f_theta is continuously differentiable then one can use the implicit function theorem to show that z* is continuously differentiable in a nieghborhood (assuming an invertibility condition holds on the partial derivatives, which will as you are assuming a contractive mapping) but this does not give twice differentiability of the solution map z*.\n\nI find these to be major gaps in the theoretical analysis that do not appear to be trivial or easily fixed.\n\nSmall inconsistency: It is written “Given most activation functions, such as ReLU and Softplus used in equation 5 are contractive,” but this is not correct, the ReLU is not contractive (its Lipschitz constant is 1).\n\n\"Nonsmooth Implicit Differentiation for Machine-Learning and Optimization\" - Jérôme Bolte, Tam Le, Edouard Pauwels, Tony Silveti-Falls, NeurIPS 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)