Gradients of Functions of Large Matrices

Thank you so much for your positive evaluation!

To answer your questions:

1. We agree that these directions are interesting for future research!
2. We use 10 Rademacher samples to match GPyTorch's default settings. Appendix H explains all Gaussian-process-related parameters.
3. The errors are relative mean-squared errors (using the square root of the machine epsilon in the current floating-point accuracy as a nugget to avoid dividing by zero). We accidentally omitted this information in the submission and will add it to the next version.

Thanks again for your review! We hope that you continue to fight for this paper's acceptance.

We greatly appreciate your positive assessment! We are happy that you acknowledge how our work "provide(s) immediate quality-of-life improvements to linear algebra users in high-value problems" because this was precisely our goal! We hope that you continue to fight for this paper's acceptance.

Thanks for the positive evaluation!

Before we answer your questions, we would like to reply to your points listed under "Weaknesses" briefly:

Matvecs, sparse/dense complexity: The complexity of our adjoint mirrors that of the forward pass of Lanczos/Arnoldi and depends almost entirely on the efficiency of the matrix-vector product. Therefore, Lanczos and Arnoldi are mainly used on sparse or highly structured matrices like PDE discretisations or Jacobians of neural networks, which is why we use a sparse matrix in Figure 3. With dense matrices, an efficient matrix-vector product can quickly become the limiting factor already for the forward pass (which is what happens in the Gaussian process case study; see the discussion in Appendix G). Finally, unless we misunderstand your sentence

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

we would like to reinforce that the Gaussian process case study does not compare to 'backpropagation through the solver' but an alternative custom gradient operation (alternative to our adjoint), as implemented in GPyTorch; see Equation 5 and/or the introduction of Section 5.

Reverse-mode AD: When you refer to reverse-mode AD in 'why this (the adjoint method) is preferable over reverse-mode AD', do you mean what we call 'backpropagation through the solver'? The adjoint method is one of many ways of implementing reverse-mode AD. If you are looking for a comparison of how backpropagation through the solver relates to our gradient implementation: we show in Figure 3 that our code is far more efficient. See also Appendix A for additional information.

Answers to questions:

1. Thank you for bringing that up. The definition is that $\mathrm{d}x$ is an infinitesimal perturbation. The gradient identity is additional information (and critical for the following derivations). We will update the sentence accordingly.
2. We agree that this is crucial; see Figure 3.
3. Thank you for bringing that up. If we back-propagate ''through'' Lanczos/Arnoldi, the resulting (automatic) gradients would use reverse-mode derivatives of matrix-vector products. Since our algorithm replaces this step, the sentence in line 198 emphasises that our code is matrix-free, just like the forward pass; see also Corollary 4.3. We will revise the sentence to make this distinction more clear.
4. Equation 54 in Appendix F formally defines the loss of accuracy. We will link this information more clearly in the main text.
5. We agree! We dedicate Appendix G to discussing JAX vs KeOps.

Thank you again for your review!

Thanks for the review and the positive assessment!

We would like to briefly reply to the points you list as weaknesses:

Motivation: You are correct that the Lanczos and Arnoldi iterations yield approximate matrix-function-vector products. When we write "exact gradient of the forward pass", we mean that our adjoints yield the exact gradient of the Lanczos/Arnoldi approximation (the forward pass). We do not claim that it is the exact gradient of the matrix function.

Table 3: We agree that the RMSE is crucial. However, since neither method beats the other in terms of RMSE, Table 3 does not highlight any RMSE-winners. For context, kegg_dir also has no winner in the training loss. As for the interpretation of this experiment, we consider it a success that our new black-box algorithm matches the performance of a technique that (i) specialises in Gaussian-process log-determinants and (ii) has been honed for many years. We discuss this in the "Analysis" block on page 7. The 20x runtime difference is not due to our algorithm. Instead, it is because the reference code (GPyTorch) relies on hand-crafted CUDA kernels through the KeOps library. The 20x performance difference is due to the inability to use this library in JAX; see the discussion in Appendix G.

Table 6: Here, too, is the main message that our black-box Arnoldi code can compete with state-of-the-art differential equation adjoints without specialising in ODEs/PDEs. We use the same algorithm for all three case studies, whereas all reference codes specialise in each respective task. For example, the differentiate-then-discretise adjoint for Dopri5 may be fast (as fast as the Arnoldi adjoint for the linear PDE) but only works for ODEs (respectively space-discretised, time-dependent PDEs). Please note how, in Figure A2, our adjoints are more efficient than Dopri5 with a differentiate-then-discretise adjoint in terms of the number of matrix-vector products.

BNN: We respectfully disagree with the "usually" in your statement that "Bayesian NN is usually trained with Variational Inference". There are many strategies for setting up Bayesian neural networks, including the Laplace approximation and variational inference; see Papamarkou et al. (2024). However, unlike other methods, the Laplace approximation is particularly relevant for our work because marginal-likelihood calibration requires (the gradient of) a function of a large matrix: the Gauss--Newton matrix. Our paper does not demonstrate how Laplace approximations compare to other techniques. Instead, it outlines how differentiable linear algebra makes Laplace approximations more efficient. Our results show how this has been successful: Our unspecialised implementation of matrix functions beats specialised Laplace-approximation codes for the VAN model. We believe that this reinforces the importance of differentiable linear algebra for advances in probabilistic machine learning.

Thank you again for the review; we hope that we were able to clarify a few potential misconceptions. In either case, we are looking forward to your reply!

Reference: Papamarkou, Theodore, et al. "Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI." Forty-first International Conference on Machine Learning. 2024.