Amortized Eigendecomposition for Neural Networks

We appreciate reviewer fsNU for the constructive feedback on our paper. Here are our responses.

Q1 It seems amortized eigendecomposition would converge slower w.r.t iterations, compared to traditional eigh/svd (because it's a rather indirect way of optimization). And I wonder how much time one can save for the whole training process?

Our amortized eigendecomposition may indeed converge slower in terms of iterations compared to traditional eigh/SVD when used purely as a numerical eigensolver. We acknowledge this limitation in the conclusion section of the paper. However, the primary advantage of our approach lies in replacing the computationally expensive eigh/SVD with a more efficient QR decomposition. While the convergence to the same results may require more steps, the QR decomposition is significantly faster overall. In our experiments, we observed a substantial speed-up with our approach, ranging from 1.4x to 600x across various tasks.

Additionally, we conducted an experiment on the latent PCA task using a backbone with 64 layers and over 1 billion parameters, and an eigendecomposition on a 4096x4096 matrix. The experimental results are shown in the additional PDF page. Our approach achieved at least a 50% reduction in time. Specifically, the average training time for such a backbone with eigh was 0.330 seconds per iteration, whereas our method reduced this to 0.155 seconds per iteration.

Q2 Can this framework be applied to other aspects such as improving the stability of SVD in network training?

Yes, our method can enhance the stability of SVD during network training. QR decomposition tends to be more stable than SVD, especially when the underlying matrix $A$ is full-ranked [see 1]. This stability can be achieved through effective initialization of $W$ (where $U = QR(W)$ ), and the rank of $W$ is preserved during optimization on the orthogonal manifold.

Furthermore, our approach can serve as a robust alternative to truncated SVD (ie. we just want the top k singular values), which is often not differentiable in frameworks like JAX and TensorFlow, or lacks numerical stability, as seen with methods like torch.lobpcg. Our method provides a stable and differentiable approach to matrix decompositions, making it a viable option in these scenarios.

Q3 For large-scale deep learning tasks with complex loss landscapes, there would be many sub-optimal solutions and the eigen loss may not be helpful.

We want to clarify that the primary contribution of this paper is not to claim that the eigen loss itself improves neural network performance. Instead, our focus is on accelerating neural networks that involve eigendecomposition (or SVD) operations. Our approach aims to make these operations faster without compromising the accuracy of the eigendecomposition results. While it is true that involving eigendecomposition can make the loss landscapes more complex, there is extensive research demonstrating its effectiveness in various tasks, such as nuclear norm regularization, graph convolutional networks, and network compression. Our paper’s goal is to provide a more efficient way to perform eigendecomposition within neural networks, achieving speed improvements in an amortized fashion while preserving performance.

Q4 Limited to eigendecomposition which is only applicable to symmetric matrices.

Although this paper focuses on eigendecomposition, our approach is NOT limited to symmetric matrices. The singular values of an asymmetric matrix $A$ can be obtained by performing eigendecomposition on the matrix $A^\top A$ (or $AA^\top$) and taking the square root of the eigenvalues, which is a well-known result. In Section C2 of the appendix, we demonstrate the extension of our approach to SVD. However, we chose to focus on eigendecomposition in this paper because it is a more fundamental and more commonly-used linear algebra operation compared to SVD. For the next version of this manuscript, we will make this point more clear.

[1]. Demmel, James W. Applied numerical linear algebra. Society for Industrial and Applied Mathematics, 1997.

We would like to thank reviewer WiHy for the constructive comments on our paper. Here are our responses.

Q1 Whether the accelerate method can maintain the performance.

In the manuscript, we try to answer this question through two experiments:

• Latent-space Principal Component Analysis: As shown in Figure 5a, the convergence curves of the traditional eigendecomposition and our approach align well. The reconstruction loss curves for both the conventional eigh function and our amortized eigendecomposition strategy are nearly indistinguishable. Initially, our method registers lower eigen-loss values compared to the eigh function, but it eventually converges to equivalent values. This demonstrates the efficacy of the amortized optimization approach.
• Adversarial Attacks on Graph Convolutional Networks: As shown in Figure 6, we find that incorporating eigendecomposition of the Laplacian matrix makes the graph convolutional network more robust to attacks on graph structures.

Both experiments demonstrate that our acceleration method can maintain performance without any degradation. Due to the page limit of the review round, we were not able to include these results in the main context. We make this point more clear in the main context of the paper if additional page is allowed.

We would like to thank reviewer 5kBi for his acknowledgement and valuable comments of our paper. Here are our responses.

Q1 Regarding the scalability.

For additional experiments on large-scale setups, please refer to the general response provided to all reviewers.

Q2 Comparison with randomized SVD.

While randomized SVD can enhance the efficiency of SVD by randomly projecting the matrix into a smaller one and then performing SVD on the smaller projected matrix, there are several key differences between our approach and randomized SVD:

• Applicability: Randomized SVD is more suitable for large sparse matrices with low rank, as mentioned in the document of torch.svd_lowrank. Our method, on the other hand, is a general eigendecomposition/SVD technique that makes no assumptions about the input matrix. For dense matrices, randomized SVD often suffers from lower performance compared to standard SVD. As shown in Figures 3 and 4, our approach achieves error rates close to those of standard SVD on dense matrices.
• Determinism: Randomized SVD is a stochastic algorithm, meaning the results can vary due to inherent randomness. Although this variability can be mitigated by using a fixed random seed, it still introduces a non-deterministic element to the computation. In contrast, our algorithm is deterministic, providing consistent and reproducible results.
• Hyperparameter: The accuracy of the randomized SVD approximation depends on the choice of the oversampling parameter p. If p is too small, the approximation may be poor. Conversely, if p is too large, it reduces computational efficiency. Our method does not require any specification of hyperparameters.

We acknowledge that our algorithm could be further accelerated using randomized linear algebra techniques such as randomized QR decomposition or randomized trace estimators. We plan to explore these possibilities in future work.

We thank reviewer 6W8h for the constructive comments and valuable suggestions. Here are our responses.

Q1 Regarding the scalability.

For additional experiments on large-scale setups, please refer to the general response provided to all reviewers.

We would like to emphasize that, as shown in Figure 1, eigendecomposition/SVD is significantly more expensive than matrix multiplication—potentially up to 1000 times slower. If a network increases in depth while keeping the width constant, the additional matrix multiplications do not significantly increase computational cost. However, if the width of the network increases substantially, the cost of eigendecomposition operations escalates dramatically. For example, in the experiment discussed in the general response, the training time of an MLP model with 64 layers and over 1 billion parameters is still faster than performing an eigendecomposition operation on a 4096x4096 matrix. Thus, the eigendecomposition typically becomes the bottleneck unless the backbone is a significantly large neural network.

Q2 Regarding Eq. 19.

Thank you for highlighting this confusion. Yes, the loss function in Eq. 19 encourages the covariance matrix to be low rank, with the rank being close to or equal to the number of principal components (PCs). Additionally, If the trace normalization is removed, the eigenvalues of the eigen loss can cause the output of the encoder to become excessively large.

If we set $\mathbf{M}=\mathbf{I}$, $\mathbf{U}$ can find the optimal subspace, but the resulting PCs are not orthogonal, meaning the covariance matrix of the PCs is not diagonal. Therefore, $\mathbf{M}$ must have distinct diagonal elements.

We acknowledge that this loss function may not be the clearest way to demonstrate latent PCA. The purpose of introducing this trace-normalized eigen loss (Eq. 19) is to demonstrate that a simple PCA layer in a network can enforce sparsity in the neural network, as shown in Figure 5c. However, we also find that this loss function may cause some confusion. An alternative loss is to avoid concerns about the eigenvalues by using a stop-gradient operator, as in Eq. 15, applied to the covariance matrix of the hidden output. The eigenloss would then become:

$tr(\mathbf{M}\mathbf{U} \text{StopGrad} (\text{cov} (h_\theta(X)) \mathbf{U} )$

In this case, the optimal $\mathbf{U}$ would be the PCA projection, and the trace loss would no longer cause the encoder/decoder parameters to become excessively large or small. Just, sparsity would not be enforced directly. We will modify this loss in the manuscript.

thanks for adding this info