Paper: Efficient Score Matching with Deep Equilibrium Layers (ICLR-2024)

Reviewer comment (Theory):
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 neighborhood (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*.

Author response:
**Reply:** We have updated our revision to clarify that $f_\theta$ is twice continuously differentiable, which aligns with the practice and guarantees that $z^\ast$ is twice differentiable in a neighborhood of $z^\ast$; see e.g. Theorem 2.1 in Section 2 of  [Lang 1983]. 

[Lang 1983] Serge Lang, Real Analysis. Addison-Wesley, 1983.

Reviewer comment (Theory):
There is a small inconsistency: It is written “Given most activation functions, such as ReLU and Softplus used in equation 5 are contractive,” but this is not correct, as the ReLU is not contractive (its Lipschitz constant is 1).

Author response:
**Reply:** Indeed, ReLU is merely non-expansive, but since the linear function within the activation is already a contraction, then $f_\theta$ will be a contraction as well. In our revision, we removed the mention of ReLU since it is not used in our experiments and in general score matching-based models; see e.g. [Hyärinen 2005, Song 2020].

[Hyärinen 2005] Aapo Hyärinen. Estimation of non-normalized statistical models by score matching. JMLR, 6(4), 2005.

[Song 2020] Song et al. Sliced score matching: A scalable approach to density and score estimation. UAI 2020.

Reviewer comment (Novelty):
DEQs and the phantom gradient method in section 3 are already well-known; what is the contribution here besides plugging them into score-matching?

Author response:
**Reply**: Our approach stands out by integrating DEQs into score matching, allowing for the effective handling of higher-order derivatives and the memory bottleneck of score matching. As a result, we have observed notable enhancements in empirical performance. To the best of our knowledge, our work is the first one that studies DEQ in score matching that requires the computation of higher derivatives.

Reviewer comment (Novelty):
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].

Author response:
**Reply:** Our primary contribution lies in identifying and addressing a specific class of problems in score matching-based density estimation or generative models that involve derivatives in the loss function, leading to higher-order derivatives during backpropagation. To tackle this challenge, we have integrated DEQ layers into score matching and examined the limits of second-order derivatives using the phantom gradient method. This approach has empirically shown significant enhancements in model performance, particularly in handling complex derivative computations efficiently.

Regarding the reference to DEQ models in diffusion models (as per [2]), it's important to note that while [2] does involve DEQ and diffusion models, its focus is distinctly different from ours. The work in [2] reformulates the entire diffusion process of the DDIM model as a DEQ, and their loss function does not involve derivatives.

Reviewer comment (Novelty):
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.

Author response:
**Reply:** As far as we know, the second-order results are new in the literature.

Reviewer comment (Presentation):
In Table 2, it was unclear what exactly DSM, CP, and SSM were referring to. 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.

Author response:
**Reply:** DKEF is a score-based density estimation model. The DSM(denoising score matching), CP(curvature propagation), and SSM(sliced score matching) variants have the same model architecture but with different ways of computing the loss. The DSM and SSM compute DSM loss and SSM loss, respectively, and CP computes an approximated, complex-valued diagonal of Hessian. Thank you for pointing these out, and we have updated the revision to include relevant discussion.

Reviewer comment (Presentation):
The acronym "CP" is used in Table 2 before it is introduced later in the paper.

Author response:
**Reply:** Thank you for pointing this out. CP stands for curvature propagation, and we have clarified this in the revision.

Reviewer comment (Evaluation):
I would have expected Table 2 to include memory usage (similar to Table 3) as this is one of the key motivations for the proposed method.

Author response:
**Reply:** Thank you for the suggestion. We have reported memory usage in Table 2 of the revised paper.

Reviewer comment (Writing):
The description of DKEF itself was unclear — what is \( q_0 \)? Where are the learnable parameters \( \theta \) appearing?

Author response:
**Reply:** $q_0$ is the log density of a mean zero Gaussian distribution with learnable variance. The main learnable parameters in DKEF occur in the modeling of the kernel function $K$.

Reviewer comment (Experiments):
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.

Author response:
**Reply:** We want to clarify that the approach taken in [2] focuses on reformulating the diffusion process of the DDIM diffusion model as a DEQ model, where the loss function does not involve derivatives. This contrasts with our paper, where we utilize a score-based diffusion model (NCSN) and employ sliced score matching (SSM) loss for training, which inherently involves derivatives. Recognizing the significance of this distinction, we have added a clarification in Section 4 of our paper.

Reviewer comment (Writing):
In Section 4.1, "encoder" should be \( q_\phi(z | x) \) and "decoder" should be \( p_\theta(x | z) \).

Author response:
**Reply:** Thank you for pointing them out. In the revised paper, we have corrected them and further improved the presentation.

Reviewer comment (Writing):
There is a typo in the caption of Figure 3: "momory" should be "memory".

Author response:
**Reply:** Thank you for pointing them out. In the revised paper, we have corrected them and further improved the presentation.

Reviewer comment (Reproducibility):
The paper is very sparse on various implementation details of DEQ. Some of the important details that are missing are: 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.

Author response:
**Reply:** We use a simple fixed-point solver, Picard iteration, to compute the fixed point in the DEQ model. Our focus in this paper is to show that using DEQ can significantly improve memory efficiency in score matching. Exploring the benefits of faster fixed-point solvers can be a future work. We aim to attribute differences directly to DEQ integration rather than solver efficiency. This rationale and the number of iterations used for each task have been clarified in Appendix F of the revised paper.

Reviewer comment (Reproducibility):
The paper is very sparse on various implementation details of DEQ. Some of the important details that are missing are: 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.

Author response:
**Reply:** In our implementation of phantom gradients, we chose $\lambda =1$ and $K=2$ for all tasks and datasets, not necessarily because they are optimal, but for simplicity and consistency. This approach allowed us to make fair comparisons with previous models without introducing additional hyperparameters that could complicate the analysis. We recognize the importance of these details and have included this explanation in our paper to clarify our choices and their rationale; see Appendix C for details.

Reviewer comment (Evaluation):
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.

Author response:
**Reply:** Thank you for highlighting the potential enhancements in DEQ-based models by increasing iterations at test time. In our experiments, we opted for a simple tolerance threshold with a maximum number of iterations, consistent for both training and testing phases for simplicity. In the revision, we have included convergence-related hyperparameters in our experiments, now detailed in Appendix F. The tolerance is set to $5e-6$ or $1e-5$, and the absolute difference between consecutive iterate $||z^{t}-f(z^{t},x)||_2$ is employed. We observe that the DEQ maintains good convergence in the experiments, possibly due to the constraint on the Frobenius norm of the linear part of the DEQ layer and the small contractive factor near zero of the Softplus activation function.

Reviewer comment (Theory):
In the discussion in Section 3.2, it is important to point that most of the practical DEQs are not globally convergent — neither uniqueness nor 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.

Author response:
**Reply:** We appreciate your insights regarding the convergence properties of DEQs. In our implementation, we have ensured the existence and uniqueness of the fixed point by constraining the Frobenius norm of the weight matrix to be less than 1 and selecting non-expansive activation functions, making the DEQ layer contractive. While this approach is more straightforward compared to MonDEQs, it effectively guarantees both existence and uniqueness. Incorporating MonDEQ, with its explicit guarantees on these aspects, certainly presents an interesting direction for future work.

Reviewer comment (Experiments):
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.

Author response:
**Reply:** In the revision, specifically in Appendix C, we present a detailed empirical analysis aimed at evaluating the effectiveness of using $K=2$ for phantom gradients. This analysis includes a direct comparison between the phantom gradient approach (with $K=2$) and the exact gradient method for the DEQ block.

We found that the score function remains largely consistent when employing phantom gradients with $K=2$, indicating that the approximation does not significantly deviate from the exact gradient in terms of both the values and trends. Moreover, we observe a strong correlation in the parameter update direction between the phantom gradient method and the exact gradient approach. This observation suggests that even with $K=2$, our method maintains a high degree of fidelity to the exact gradient.

Reviewer comment (Experiments):
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].

Author response:
**Reply:** Our work concentrates on model depth scalability in score matching-based models rather than data dimensionality. It's crucial to note that the challenge of quadratic dependency on data dimensionality stems from the requirement to store second-order derivatives, which is inherent to the score matching (SM) loss function itself. This particular dimensionality bottleneck is effectively addressed by variants like sliced score matching (SSM), which reduces the computational burden of high-dimensional data.

Regarding the benefit of constant memory in DEQs, the results in Tables 1, 2, and 3 show that DEQ-based models outperform baseline models and their increased depth counterparts while maintaining low memory footprints.

Reviewer comment (Experiments):
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.

Author response:
**Reply:** We have included the parameter count for each experiment. Utilizing the parameter efficiency of DEQ models, our DEQ-based models often have a similar parameter count as the models chosen as baselines but with much-improved performance.

Reviewer comment (Reproducibility):
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?

Author response:
**Reply:**  Indeed, only DKEF on Gaussian and UCI datasets uses SM loss and hence computing Hessian is needed to form the loss. However, we want to emphasize that even though SSM loss and DSM loss only involves first-order derivative in the loss, higher-order derivatives still occur during backpropagation for updating the weights. 

We have updated Table 2 to include memory consumption and performance of SSM and SSM-VR.

Reviewer comment (Presentation):
There are several missing details in the tables of experimental sections: The size of models in terms of number of parameters should be indicated in Tables 2, 4, and 5.

Author response:
**Reply:** We have added the parameter count in these Tables.

Reviewer comment (Reproducibility):
There are several missing details in the tables of experimental sections: 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.

Author response:
**Reply:** $g(y)$, in this case, is a simple MLP transform of the input with $y\rightarrow \sigma(Wy+b)$, where $\sigma$ is Softplus and $W, b$ are learnable weights.

Reviewer comment (Experiments):
There are several missing details in the tables of experimental sections: Please include the size of CelebA images used in Section 4.1

Author response:
**Reply:** The size of each image in CelebA is $64\times 64\times 3$, and we have added the description in our revision.

Reviewer comment (Theory):
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 --)

Author response:
As in traditional score-based density estimation models, the log density function is modeled using neural networks, with loss functions related to Fisher divergence terms like score matching (SM), denoising score matching (DSM), or sliced score matching (SSM). 

Our approach transforms the core layers of these models into DEQ layers. In this case, the fixed points act as an intermediate hidden representation. This modification significantly increases the model's depth, allowing for more complex representations while maintaining a low memory footprint. The score function, being the derivative of the log density function, is thus derived from the neural network's output, computed via autodifferentiation with Neumann approximation for efficiency.

We have extended the discussion about the neural network modeling log density function and DEQ layer before Section 3.1 for clarity.

Reviewer comment (Theory):
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 -)

Author response:
Thank you for your feedback on Proposition 2. In the revised version, specifically in Appendix C, we present a detailed empirical analysis aimed at evaluating the effectiveness of using $K =2$ for phantom gradients. This analysis includes a direct comparison between the phantom gradient approach (with $K=2$) and the exact gradient method for the DEQ block.

We find that the score function remains largely consistent when employing phantom gradients with $K=2$, indicating that the approximation does not significantly deviate from the exact gradient in terms of both the values and trends. Moreover, we observed a strong correlation in the parameter update direction between the phantom gradient method and the exact gradient approach. This observation suggests that even with $K=2$, our method maintains a high degree of fidelity to the exact gradient.

Reviewer comment (Writing):
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".

Author response:
Thank you for your suggestion. We have updated our manuscript to improve the writing further.

Reviewer comment (Theory):
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.

Author response:
That sentence is supposed to convey that the linear map inside the activation function is contractive. We have fixed the typo in presenting the linear map in the revision and clarified that contractive linear and non-expansive activation ensures the well-posedness of the fixed point equation 5.

Reviewer comment (Presentation):
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.

Author response:
Thank you for pointing this out. Indeed, these sampled images are only for visualization purposes. We have moved them from the main text into Appendix E. The FID scores presented in Table 1 present a better quantitative comparison.

Reviewer comment (Theory):
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?

Author response:
As explained in the response to Q1. The original density estimation models output the log density function and use the Fisher divergence-related loss function for parameter optimization. Our DEQ integration replaces the core parts of the previous model with DEQ layers. The fixed point doesn’t necessarily have a meaning that is directly related to the score function but rather is just used as a hidden representation.