Gradients of Functions of Large Matrices

Summary

This paper introduces a new matrix free method for automatically differentiating functions of matrices. The computational problem discussed in this paper is of interest because the matrix dimension scales with respect to the size of the dataset (e.g. Gaussian process regression, etc.). The authors' algorithm yields the exact gradients of the forward pass, all gradients are obtained with the same code, and said code runs in linear time- and memory-complexity. The proposed method is also matrix-free and does not form the matrix explicitly which is the key to scalability. The authors' method evaluate functions of large matrices by: 1) first decomposing the matrix into a product of small matrices (e.g. Lanczos, Arnoldi iterations); 2) then evaluating functions on the small matrix produced during the Lanczos/Arnoldi iteration. The authors note that functions of small matrices can be differentiated efficiently. The authors provide the method to compute gradients of the Lanczos/Arnoldi iteration so that backpropagation can be implemented efficiently. The authors use differentiation via adjoints of the Lanczos/Anroldi iterations (this is described in Theorem 4.1 and Theorem 4.2). The parameter gradient expression that can be used for any problem is provided in Corollary 4.3. The authors provide experimental results which provide impressive speedup over GPyTorch over problems of interest including: Gaussian process, physics-informed PDEs, Bayesian neural network tuning.

Strengths

1. This paper tackles an important computational bottleneck in tuning hyperparameters of potentially complex models that are used in scientific and machine learning applications. 2. The proposed method is general and can accomodate any matrix functionals and do not need explicit derivations of the gradients. 3. The authors provide experimental results on several problems of interest to machine learning practitioners.

Weaknesses

There are no particular weaknesses I found on reading the paper. The paper is well written.

Questions

1. There are parallelization techniques that may be useful for even scaling up the proposed method. 2. How many samples do you take for stochastic estimation of log-determinant for the exact Gaussian process result? 3. What norm is the gradient error measured in?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors describe several limitations - but as the authors mention, these extensions describe a more expanded write-up.

Summary

There are some useful iterative methods for calculating important matrix products, to wit, Lanczos and Arnoldi iterations, which apparently did not have known derivatives, until now. The paper provides a generic framework for calculating the derivatives of such iterative linear operator approximations. The core of the paper is interpreting the Lanczos resp Arnoldi iteration as an iterative system and using the implicit derivative trick

Strengths

While these results are not *world* shaking, they seem to provide immediate quality-of-life improvements to linear algebra users in high-value problems. The results unify several known "matrix free" tricks, which is satisfying, and point the way to more. The paper is well-written, compact, and easy to understand, which is a great relief this deep into the review process.

Weaknesses

As far as I can see, few. The paper seems to present what I need to know. It would not be impossible for these results to be known elsewhere in the literature, but if they are known, I have not seen them.

Questions

I have no questions. This paper was well explained, and clear about its goals, and how it approached them.

Rating

8

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

There are many limitations to the methods divulged here, but they seem to be adequately articulated in the paper.

Summary

The paper proposes an adjoint method for functions of matrices that utilize Arnoldi/Lanczos iterations to compute gradients with respect to large dimensional variables and demonstrates their approach's utility on a variety of common compute/memory intensive tasks encountered in machine learning. They derive the adjoints by implicitly differentiating the solution of Arnoldi/Lanczos iterations with respect to the input matrix and vector and then show how this can be used to compute gradients with respect to parameters that the input matrix depends on. They then examine their method on three case studies: Gaussian process hyperparameter optimization, solving for parameters of a discretized two-dimensional (in space) PDE and post-training calibration of a Bayesian neural network. Their results demonstrate an improvement over the standard approaches on a range of metrics.

Strengths

- An adjoint method for the Arnoldi iteration is certainly a useful addition to the automatic differentiation toolbox given the techniques wide use and would thus be of significant interest to the Neurips community. - The paper was mostly clear to read and understand. The conciseness of the experiment section is also a strength as the key aspects of the setups and the results are succinctly presented. - The paper's results indicate the methods utility in obtaining better runtimes, more accurate gradients and/or lower train and test losses on a range of important tasks in machine learning. In particular, the BNN calibration experiment clearly demonstrates how their method can handle matrices that may be too large to fit on a GPU and so it doesn't have to resort to subsequent approximation which they show leads to better performance.

Weaknesses

- Although the paper empirically shows how backprop through a matrix function is slower than their adjoint method for Arnoldi iterations in Fig 3, the fact that a sparse matrix (as opposed to a dense one) was utilized motivates some skepticism over whether we would expect to see this same runtime relationship for a dense matrix. In fact, the exact GP results suggest that implementing matrix-matrix (/-vector) operations more efficiently can be enough to make reverse-mode AD work quicker for dense matrices. More discussion on why the adjoint Arnoldi iteration is more appropriate would strengthen their argument for the improved efficiency of their method. - The description of the adjoint systems in section 4.1 is lacking some explanation to make clear the key aspects of the adjoint methods and why this is preferable over reverse-mode AD. Reverse-mode AD also only utilizes vector-Jacobian products and so is a kind of "matrix-free" method. Explanation of how their method differs from this would make more apparent their contribution.

Questions

- The last part of the explanation of implicit differentiation in section 4 was confusing. On the one hand $\mathrm{d}\rho$ is defined in eqn (9) but the form $\mathrm{d}\rho = \langle \nabla_{\theta} \rho, \mathrm{d}\theta \rangle$ is given on the following line. It was not clear as to why these two should be equal. Being more concrete on the definition of $\rho$ would help here - As a solution to the second weakness listed above an algorithm comparing adjoint Arnoldi/Lanczos with reverse-mode AD may help to make clear the differences between the methods. - On the sentence starting with "Gradients ..." on line 198, doesn't reverse-mode AD also only use matrix-vector products as well? - A definition of "loss of accuracy" in Table 2 should probably be given in the main text. - A bit of discussion as to why the JAX low-memory matrix-vector product implementations cannot compete with KeOps would be a worthwhile addition to the paper.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations on the conceptual aspect of the approach are discussed at the end of section 3 but some discussion in light of the experimental results would also be of value

Summary

This paper proposes a new approach to perform automatic differentiation for function of large matrices. Specifically, the paper outlines the backward computation of the matrix-vector product f(A(\theta)) * v where A(\theta) is the jacobian of a large NN that will not fit into memory. The proposed approach uses the Lanczos and Arnoldi iteration method to factorize the matrix-vector product, which permits inexpensive backpropagation, and then derive the backward computation for the Lanczos and Arnoldi iteration via the adjoint method. This technique is tested in three different scenarios that require the evaluation of such matrix-vector product: Gaussian processes (GP), Physics Informed ML with PDE, and Bayesian NN.

Strengths

- Theoretically interesting approach to auto diff of large Jacobian vector product - Potentially useful in scenarios where functions of large matrices are involved in the objective function - Interesting case studies detailing several such scenarios and how to apply the proposed method in those scenarios - Improved performance in majority of experiments - Generally clear and self-contained presentation, easy to read and follow

Weaknesses

1. Motivation is somewhat unclear. The paper claims that the proposed method provides exact gradient of the forward pass, but it doesn't seem true when the Lanczos & Arnoldi iterations themselves do not yield an exact factorization (Eq. 7). This method, although theoretically quite interesting, should belong to the same class with other approximation methods. As such, I would like the authors to explicitly discuss the benefit of this approximation compared to previous work (which are currently positioned as being inferior to an exact gradient method) 2. Empirical results are not convincing: - Table 3 seems to highlight the wrong improvement. It is moot to compare training loss when the loss landscapes are different. The real performance measure should be RMSE, in terms of which both methods converge to the same values (although the proposed Arnoldi method is 20x slower -- so I'm not sure what's the benefit here) - Case study 6 again shows the same performance between the Arnoldi method & Dopri5. - Bayesian NN is usually trained with Variational Inference. Diagonal method is a rather crude approximation, and it is not surprising that the proposed method outperforms it.

Questions

I have no further question since the method is sound & clearly presented. My only problem is that the empirical results are really limited and do not show any significant improvement over existing methods.

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

No potential negative societal impacts

Dear all

I have read the authors' response and the other reviewers' comments on the submission. I am overall happy with the authors' rebuttal response and would vote for acceptance of this paper.

Thanks for the response

First of all, my apology for the late chiming in. My work email changed recently so I didn't receive any notification...  Please understand that my overall sentiment is still positive although my score fell on the skeptical side. I really do like the technique and I'm happy to raise my score to 6 for the theoretical merit. Nonetheless, there are still parts of the empirical results that are not convincing to me, Regarding table 3: I'm not fully agreeing with what is considered a success here. I feel like some signal showing the practicality of your method is sorely needed. Maybe the right way to move forward is to set up an implementation of GP without the KeOps backend, so things can be compared on equal footing. Regarding BNN: I agree that there are many strategies for setting up BNN. Perhaps I missed the point on the Laplacian approximation being relevant to your method. However, I disagree that it is irrelevant to compare to other techniques because it would be interesting to show that how your improvement fares against other sota methods.