Amortized Eigendecomposition for Neural Networks

Authors: tianbo  li
zekun  shi
jiaxi  zhao
min  lin
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 08:03:32

Abstract: Performing eigendecomposition during neural network training is essential for tasks such as dimensionality reduction, network compression, image denoising, and graph learning. However, eigendecomposition is computationally expensive as it is orders of magnitude slower than other neural network operations. To address this challenge, we propose a novel approach called "amortized eigendecomposition" that relaxes the exact eigendecomposition by introducing an additional loss term called eigen loss. Our approach offers significant speed improvements by replacing the computationally expensive eigendecomposition with a more affordable QR decomposition at each iteration. Theoretical analysis guarantees that the desired eigenpair is attained as optima of the eigen loss. Empirical studies on nuclear norm regularization, latent-space principal component analysis, and graphs adversarial learning demonstrate significant improvements in training efficiency while producing nearly identical outcomes to conventional approaches. This novel methodology promises to integrate eigendecomposition efficiently into neural network training, overcoming existing computational challenges and unlocking new potential for advanced deep learning applications.

Main paper: 
 introduction: 
 Eigendecomposition is a fundamental technique in linear algebra that finds applications across numerous scientific domains ranging from quantum manybody problems to multivariate statistical analysis.In the context of deep learning, eigendecomposition also plays a crucial role in tasks such as weights normalization [8, 16, 41] , dimensionality reduction [6, 44, 27, 38] , network compression [17, 30] , image denoising [13, 12, 14] , graph adversarial learning [10, 18, 47] . By uncovering the structure of networks, eigendecomposition allows us to enforce low-rankness, ensuring generalization, robustness, and computational efficiency. Eigendecomposition is also instrumental in the spectral analysis of graphs, where it can detect community structure, which is essential in spectral graph neural networks. The ability of eigendecomposition to detect the intrinsic matrix structures and properties makes it a valuable tool in various machine learning tasks with neural networks.Despite its straightforward definition, the computation of eigenvalues is quite challenging. Eigendecomposition algorithms are inherently iterative, involving a sequence of expensive operations such as Arnoldi iteration, QR iteration, and Rayleigh quotient Iteration [39] . Additionally, it usually takes thousands of iterations to reach a desirable tolerance level. For instance, the locally optimal block preconditioned conjugate gradient (LOBPCG) method [22] -a widely-used eigenvalue solver-often requires dozens to hundreds of iterations to achieve convergence [22, 9] . Figure 1 presents a comparative analysis of the execution times for eigendecomposition and other computational operations using PyTorch [33] and JAX [2] . The figure shows that the execution speeds for eigh and svd are remarkably slower-by 4 to 5 orders of magnitude-relative to matrix multiplication. This substantial disparity in execution speed indicates operations such as eigh and svd, once used, will be the bottleneck of computation cost. Conversely, the QR decomposition, often employed for orthogonality, is considerably less computationally expensive. This observation has motivated us to explore the possibility of reducing the iterative computation of eigendecomposition with lower-cost operations.When eigendecomposition is incorporated into the training of a neural network, a nested loop scenario arises, where eigendecomposition acts as the inner loop and the neural network's loss minimization serves as the outer loop. Notably, it is not always easy to be aware of this inner loop of eigendecomposition, as it is encapsulated by the high-level functions provided within deep learning frameworks. However, in this context, this inner loop does not require full convergence during each iteration, given that the remaining parameters have not reached optima. The inner loop can be relaxed and optimized jointly with the training loss, allowing for a more flexible and efficient training process.The key idea can be summarized as follows:Eigendecomposition within a neural network does not have to reach full convergence during each training step; it simply needs to contribute to the desired outcome by the end of the training process.In this paper, we present a novel approach named "amortized eigendecomposition" for training neural networks that require eigenvalues or eigenvectors. Instead of using computationally expensive eigendecomposition decoupled with the training of neural networks, we proposed to relax it into an unconstrained optimization problem on the Stiefel manifold by adding an eigen loss. This relaxation only requires a QR decomposition at each iteration, thus is more efficient. Moreover, through empirical observations, we have found that although the relaxed optimization problem with eigen loss does not involve eigendecomposition in every iteration, the amortized optimization approach consistently achieves the desired results. It achieves nearly identical performance to traditional methods but with significantly improved speed.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The main claims consist of theoretical analysis of the loss functions which is stated in Thm. 1 and Thm. 2 and verifications of our method in numerical experiments are done in Sec. 5.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
eigendecomposition in neural networks: 
 In this paper, we consider a general class of neural networks that incorporate eigendecomposition. We formulate this family of problems as a constrained optimization problem:EQUATIONHere, the encoder h θ maps the data X ∈ R n×p into a latent space. In addition to the latent representation, the loss function f also incorporates the eigenpair Λ and V of a symmetric matrix A.The matrix A can be constructed from h θ (X), such as a covariance matrix or a similarity matrix, or it can depend solely on the parameters. Notably, A is subject to changes during network training due to its dependency on θ. The computational graph for each iteration of such model structure can be written as,EQUATION)The corresponding algorithm of the above computational graph is shown in Algorithm 1. Solving such problems, however, is computationally expensive, since it requires preserving the eigendecomposition constraint of A after each update of θ. This structure encompasses a wide range of learning problems. We present several representative examples and investigate them in our numerical experiments next.The nuclear norm of a matrix is defined as the sum its the singular values. Due to its convexity, regularization via the nuclear norm is employed to encourage low-rank structures within the learned parameters. This approach proves beneficial in a variety of applications, such as matrix completion [5, 4] , image denoising [13, 12, 14, 45] . Furthermore, eigendecomposition and singular value decomposition are also used for pruning or compressing neural networks, by decomposing the weight matrices of the network and approximating the original network with fewer parameters [17, 30] . The objective function of this type of problem can be written as,EQUATIONHere, η is a regularization coefficient that controls the rank of the parameter matrix. In classical methods, the nuclear norm ∥θ∥ * is usually calculated via singular value decomposition.Whitening is a transformation technique extensively utilized in neural networks to standardize features by ensuring they have zero mean and unit variance, and are decorrelated from each other. In neural network applications, whitening can be categorized into parameter-space whitening and feature-space whitening. Parameter-space whitening, often achieved through PCA/ZCA, is a prevalent method applied during neural network training to improve stability and accelerate convergence [37, 41, 16, 8] . Feature-space whitening, in contrast, applies PCA to the intermediate representations within the network. This process aligns features with the axes of greatest variance, thereby facilitating dimensionality reduction [15, 42, 34, 25] . A concrete example of feature-space whitening is its incorporation into an auto-encoder architecture, which can be expressed mathematically as:EQUATION)where V is spanned by the first several largest eigenvectors of the covariance of Enc θ (X) and θ, ω represent the parameters of the decoder and encoder, respectively.Graph Structure Learning Graph structure learning is an area of machine learning that aims to deduce the latent structure of a graph or network from observed data [23] . This domain has significant applications in graph adversarial learning, which seeks to bolster the robustness of graph neural networks against adversarial attacks [10, 18, 47] . In such contexts, the adjacency matrix is usually often compromised by adversarial modifications while feature matrix X remains unaffected. The primary objective is to learn both a clean graph structure and perform accurate node classification.From this perspective, we follow the approach proposed in [18] and the objective function can be formulated as:min θ, L GNN θ (X, L) -Y 2 + α L * + β L -L 2 F (5)Here L represents a low-rank Laplacian matrix that approximates the original graph Laplacian L.The three parts of the loss function correspond to the node classification loss, low-rank constraint, and unnoticeable adversarial attacks, respectively. However, this approach relies on singular value decomposition at every iteration, which is computationally prohibitive for large-scale networks.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We show in the discussion section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
differentiable optimization for eigendecomposition: 
 Before introducing the proposed method, let's begin by examining a more straightforward case: determining the eigenvalues of a symmetric matrix A ∈ R n×n through constrained optimization. 1The eigenvalues of A are denoted asλ 1 ⩾ λ 2 ⩾ • • • ⩾ λ n .The largest eigenvalue of A can be obtained by directly maximizing the Rayleigh quotient:EQUATION)where the maximum value corresponds to the largest eigenvalue and the associated eigenvector is given by the normalized u at the optimum. This method can be generalized to identifying the optimal subspace, which is a common problem in dimensionality reduction and feature selection contexts.The optimal subspace for a symmetric matrix is defined as the subspace spanned by its k largest eigenvectors. This can be found by solving the constrained optimization problem: -y x , which is subject to the constraint x 2 + y 2 = 1. The plot displays the contours of the landscapes of three different loss functions as they vary with x and y: (a) trace loss tr(U AU T ), (b) Brockett's cost function tr(M U AU T ) where M = diag(1, 0.2) and (c) convex loss function tr(f (U AU T )) where f is an element-wise exponential function. The feasible area of the constraint is depicted with a black circle. The red stars signify the optima of the loss in the feasible area. The dashed grey lines represent the true eigenvector direction of A. We see that, the trace loss results in infinitely many optimal solutions due to its rotational symmetry. In contrast, both Brockett's cost function and the convex loss function reshape the optimization landscape, breaking this symmetry and leading to the identification of the correct eigenvectors.EQUATIONwhere the Stiefel manifold U n×k = {U ∈ R n×k | U T U = I} is the set of n × k matrices with orthonormal columns. Maximizing this trace function yields the optimal subspace spanned by the column vectors of U .Rotational Symmetry It's important to recognize, however, that the optimal U of the above optimization problem does NOT represent the eigenvectors of A. This is due to the rotational symmetry of the trace function that for any orthonormal matrix C ∈ U k×k , the equation tr(U T AU ) = tr(C T U T AU C) always holds. As a result, there are an infinite number of solutions to the optimization problem Eq. ( 7), all spanning the same subspace as the desired sets of eigenvectors. A visual illustration is provided in Figure 2a .To accurately obtain the eigenvalues and eigenvectors, it is necessary to refine the traditional trace loss function. This paper introduces two approaches for achieving the correct eigendecomposition. The first method utilizes Brockett's cost function, which applies distinct weights to the diagonal elements of the matrix product U T AU , effectively differentiating the importance of each eigenvalue. The second method involves applying an element-wise convex function directly to the diagonal elements, resulting in an exact eigendecomposition. We will now elaborate on these methods.The first modification made to the trace loss is Brockett's cost function [3] . It enables the extraction of eigenvectors and eigenvalues by solving the following optimization problem:EQUATIONHereM = diag(m 1 , m 2 , • • • , m k ) ∈ R k×kis a diagonal weight matrix with distinct diagonal elements. The matrix M is structured such that all its elements are distinct numbers. Specifically, we can denote these elements as 0 < m 1 < m 2 < • • • < m k . This ordering allows M to assign distinct weights to the diagonal elements of the product U T AU , effectively disrupting the rotational invariance inherent in the trace loss and thus enabling the determination of eigenvalues and eigenvectors through optimization. In fact, M can be any diagonal matrix with distinct diagonal elements. A practical choice could be m i = i/k or m i = i as suggested by [3] . The next theorem illustrates why this approach can yield the exact eigendecomposition. Theorem 1 (Trace inequality with weight matrix). Consider a n-dimensional symmetric matrix A.Let M be a k-dimensional diagonal matrix with elements 0 < m 1 < m 2 • • • < m k . For an arbitrary matrix U ∈ U n×k with orthogonal columns, the following inequality always holds:EQUATIONThe equalities are achieved if and only if U contains the eigenvectors corresponding to the k largest (smallest) eigenvalues of A.This approach can be viewed as an extension of the trace inequalities originally put forward by Von Neumann [40, 31] and further developed by Ruhe [35] . A detailed demonstration is available in the seminal work by Brockett [3] . Additionally, we offer an alternative proof leveraging the Cauchy-Schwarz inequality, which is presented in the Appendix for reference. A more generalized version of this trace inequality is discussed in Liang et al. [26] .The Convex Trace Loss The second method employs a strictly monotonic convex function f , applied element-wise to the diagonal components of U T AU . This perturbation also disrupts the rotational symmetry inherent in the trace loss. The convex nature of f alters the curvature of the loss landscape, thereby ensuring a unique optimal solution corresponding to the eigenvectors. The convex trace loss function, aimed at extracting the k largest eigenvalues, is expressed as:EQUATIONThe optimal U * that achieves the maximum in the above objective are the eigenvectors corresponding to the k largest eigenvalues, and the eigenvalues can be obtained by the diagonal elements of U T * AU * . The next theorem provides a formal validation of this assertion. Theorem 2 (Trace inequality with convex function). Let A be a given n-dimensional symmetric matrix and let U be a matrix of size n × k that resides on the Stiefel manifold. Suppose f : R → R be a monotonically increasing, convex function applied element-wise. The following inequalities hold:EQUATIONThe rightmost inequality becomes an equality if and only if U comprises the eigenvectors of A that correspond to the k largest eigenvalues, The leftmost inequality is met with equality when U is such that all diagonal elements of the matrix U T AU are equal.This theorem is established through the application of Jensen's inequality. The detailed proof is provided in the Appendix B.3 for reference. In the above objective, suitable choices of f include:f (x) = exp(x) on R; f (x) = x α on R + where α > 1; and f (x) = tan(x) on [0, π/2].For finding the k smallest eigenvalues, a simple modification can be made by replacing the function f with an element-wise monotonically increasing concave function and then minimizing the trace loss. A visual representation of the optimization process of the three trace losses is presented in Figure 2 .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: Our code will be released upon acceptance.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
the amortized eigendecomposition approach: 
 The proposed amortized eigendecomposition approach aims to modify the eigendecomposition operation within the neural network's computational graph, as illustrated in Eq. ( 2). This replacement involves two steps: First, the set of eigenvectors form a matrix on the Stiefel manifold reparameterized through computationally efficient operations such as QR decomposition. Then, the loss function is adjusted to ensure that its optimal solutions precisely correspond to the eigendecomposition. The computational graph for this amortized eigendecomposition method is formulated as follows:EQUATIONIn the computation graph, the eigendecomposition operation is circumvented and substituted with a more efficient QR operation. The QR operation is employed to reparameterize the orthogonal matrix Algorithm 1 The conventional eigendecomposition in a neural network outlined in Eq. ( 2)Input: Dataset X, encoder h θ , task fω; 1: Initialize model parameter θ and ω; 2: while not converged do 3: compute A from h θ (X); 4:V , Λ = eigh(A); 5:compute fω(h θ (X), V , Λ); 6: update θ, ω by gradient descent; 7: end while Algorithm 2 The amortized eigendecomposition technique outlined in Eq. (12) Input: Dataset X, encoder h θ , task fω, and η;1: Initialize model parameter θ, ω and W ; 2: while not converged do 3: compute A from h θ (X); 4:U = QR(W ), Σ = diag(U T AU ); 5: compute fω(h θ (X), U , Σ) + ηg(A, U ); 6:update θ, ω, W by gradient descent; 7: end while on the Stiefel Manifold, leading to a substantial acceleration at each iteration. Moreover, instead of forcing the eigendecomposition constraint at each iteration, we relax it to an eigen loss which is jointly optimized with the training loss of the neural network as a nested optimization loop. This training process is outlined in Algorithm 2, where the key difference of our amortized optimization for the eigendecomposition approach is highlighted in red background color.Reparameterize the Stiefel Manifold There are three prevalent methods for reparameterizing an orthogonal matrix: through the matrix exponential, the Cayley transform, and QR decomposition. Due to the QR decomposition's better numerical stability and efficiency for non-square matrices U , we employ it for reparameterizing a matrix with orthonormal columns:EQUATIONIn this formulation, W is dimensionally consistent with U . QR decomposition is more computationally efficient than eigendecomposition and singular value decomposition as shown in Fig. 1 .Additionally, the backward computation of the QR decomposition is well-defined and has been efficiently optimized in modern deep learning frameworks, such as PyTorch and JAX. For details on the back-propagation process of the QR decomposition, see [36] .Relaxation with Eigen Loss Previously, we observed that optimizing the Brockett-type or convex trace loss directly enables us to obtain precise eigenvalues and eigenvectors. For any loss function that depends on the eigenvectors or eigenvalues, as specified in Eq.( 1), we can transform this loss into a regularized version incorporating the eigen loss. This relaxation allows us to forego the need for explicit eigendecomposition at every iteration, while ultimately achieving equivalent outcomes. We now examine several general scenarios that encompass the majority of cases and explore how to implement this relaxation technique.In the general case, the model loss in Eq. ( 1) depends on both the eigenvectors and eigenvalues. This constrained optimization problem can be relaxed by introducing an eigen loss as a regularizer, which is formulated as:EQUATIONIn this reformulation, U ∈ U n×k , is reparameterized via a QR operation as shown in Eq. ( 13). This relaxation circumvents the need for eigendecomposition at each iteration by using the computationally cheaper QR decomposition, while still ensuring that the optimal solution corresponds to the precise eigendecomposition of the matrix A θ .The second type of optimization problem involves scenarios where the model loss is independent of the eigenvalues and depends solely on the eigenvectors, such as in the latent-space PCA network expressed in Eq. (4). To enhance the efficiency of the solution process, the problem can be reformulated to include a trace penalty term, such as Brockett's cost, which is given byEQUATIONIn this formulation, a stop gradient operation is applied to A, since the eigen loss involves A which relies on the parameter θ. By introducing the stop gradient operation, we prevent this regularization term from propagating gradients back to θ. Besides the above relaxations, further simplification is possible if the loss function conforms to the structure of Brockett's cost function or a convex trace function. That is, when the model loss f ω is a non-uniform linear combination of the eigenvalues, or it is monotonic and convex (or concave) to the eigenvalues, the trace penalty term is unnecessary. The loss function will inherently converge to the correct eigendecomposition. This principle is exemplified in the context of the nuclear norm regularization problem, which will be illustrated later. This can be formulated as:EQUATION 
experiments: 
 In this section, we present an evaluation of our approach, focusing on four specific tasks. Firstly, we demonstrate the convergence properties of our method empirically. Next, we measure the efficacy and efficiency of our amortized eigendecomposition technique over nuclear norm regularized autoencoder and latent-space PCA using the MNIST dataset. Lastly, we assess the effectiveness of our approach in the context of graph adversarial learning tasks. We implement our approach with the deep learning framework JAX [2] . All the experiments of our approach are conducted on a single NVIDIA A100 GPU with 40GB memory. The two fundamental questions we investigate are as follows:• Does our approach accurately identify the eigendecomposition and singular value decomposition?• How does the efficiency of our method compare to that of traditional techniques?In this experiment, we evaluate the numerical error and convergence speed of our algorithm applied to solving eigendecomposition. We randomly generate ten symmetric matrices of size 1000 × 1000.The first 50 eigenvalues of these matrices range from 1 to 50, while the remaining eigenvalues lie between 0 and 1. Our objective is to compute the 50 largest eigenvalues by minimizing Brockett's cost function and convex trace loss (we adopt f (x) = x 1.5 as the convex function). To achieve this, we employ several optimization algorithms, including Adam [20] , Adamax [20] , Yogi [46] , SGD, and L-BFGS [28] . These algorithms are provided by the Optax and JAX-opt libraries [7] . We measure the mean square error (MSE) of the eigenvalues to the number of training iterations. The results are illustrated in Figure 3 . This result demonstrates that both loss functions are capable of identifying the correct eigenvalues with a small numerical error of 10 -9 . However, there is a noticeable difference in convergence speed.For both trace loss functions, the Adam and Adamax optimizers outperform the others, achieving faster convergence rates. Brockett's cost function, which introduces a linear combination of the trace elements, is more numerically stable compared to the convex trace loss, resulting in faster convergence. This experiment validates the efficiency of the differentiable optimization framework for computing the k largest eigenvalues.In this experiment, we apply the amortized eigendecomposition approach to the nuclear norm regularization problem, as outlined in Eq. ( 3). The experimental framework entails training an auto-encoder on the MNIST dataset by minimizing the reconstruction loss with a nuclear norm regularizer applied to the weight matrix θ ∈ R n×k of the encoder's last layer. We employ a relaxation technique to the original problem defined in Eq. ( 16), which can be expressed as:min ω,θ Dec ω Enc θ (X) -X F +η∥θ∥ * --→ min ω,θ,W Dec ω Enc θ (X) -X F +η k i=1 θu i 2 (17)where u i 's are the orthogonal column vectors of U , which are parameterized by W through Eq. ( 13).The architectures of the encoder and the decoder are constructed as a 2-layer MLP with hidden layer dimensions of D = 128, 256, and 512. For comparison, we also implement the approach based on singular value decomposition (using the svd function). It should be noted that in the current versions of JAX, both the eigh and svd functions are limited to operations on full matrices.The average execution time per iteration for the baseline backbone with only reconstruction loss, the backbone with svd, i.e. LHS of Eq. ( 17), and the backbone utilizing amortized eigendecomposition, i.e. RHS of Eq. ( 17) are reported in Table 1 . We denote these execution times as t 0 , t 1 , and t 2 respectively, and define the speed-up ratio for our approach relative to the svd as:EQUATION)This ratio represents the improvement in execution speed of our eigendecomposition method compared to the standard svd, relative to the baseline backbone performance.We investigate the effectiveness of our approach for the latent-space PCA method, as described in Eq. ( 4), using the MNIST dataset. Computing the eigenvectors of the large-scale covariance matrix in each iteration significantly increases the computation overhead, while our method amortizes this cost by jointly minimizing an additional loss. Moreover, we also aim to ensure that the first two eigenvalues are significantly larger than the subsequent ones. In the following objective function, the additional term resembles Brockett's cost function while the trace of the covariance matrix ensures homogeneity:EQUATIONwhere cov represents the covariance function. The architecture for both the encoder and decoder mirrors that of the nuclear norm regularization model, with the exception that there is a linear projection aligning with the direction of the principal components. The average execution times are reported in Table 1 . More results and analysis are provided in Appendix A.2. The experimental results demonstrate that our approach achieves an average training speed improvement of 40% compared to the conventional eigendecomposition approach.Additionally, we conducted a scalability study, with the results presented in Figure 4 . This study examined the scaling behavior of latent PCA on the Celeb-A-HQ (256x256) dataset [29] by varying both the depth and width of the backbone autoencoder, with average execution time per iteration reported. The largest model in our tests, a 64-layer autoencoder with a 4096-dimensional latent space, contains over 1 billion parameters. From these results, we draw two main conclusions. First, our amortized eigen loss substantially reduces the eigendecomposition training time without significantly increasing the computational load of the backbone, as evidenced by the close alignment of the red (backbone) and green (backbone + our approach) lines. In contrast, the traditional eigendecomposition approach (blue line) scales steeply with increasing dimensionality, whereas our approach exhibits a much slower growth rate. Second, eigendecomposition emerges as the primary computational bottleneck within these neural network architectures, while the fully-connected layer computation remains minor, particularly at large widths (>2000). This is reflected in the widening gap between the backbone (red line) and backbone + eigh (blue line) as dimensionality increases. Notably, increasing the depth of the backbone while keeping the hidden dimension constant results in minimal change in execution time, indicating that the cost of fully-connected layers is small relative to eigendecomposition-further underscoring the results shown in Figure 1 .In this study, we explore the robustness of graph convolutional networks (GCNs) [21] by implementing adversarial attacks on graph structures. The objective of this problem is described in Eq. ( 5). Our approach simplifies the attainment of a low-rank structure by optimizing:EQUATIONwhere L = U T diag(U T LU )U , with U being parameterized by Eq. ( 13). The motivation for this formulation is that at optimum, the columns of U correspond to the top-k eigenvectors of L. Then L becomes the best rank k approximation of L under the Frobenius norm which corresponds to the terms L * , L -L 2 F in Eq. ( 5). This formulation allows the GNN to operate on a low-rank graph L, which has been shown to enhance robustness against adversarial attacks on the graph structure.Our architecture consists of a three-layer GCN, which is utilized for semi-supervised node classification tasks on several citation networks, namely Cora, Citeseer, and Pubmed. Each layer has a hidden dimension of 32. The dropout rates are set to 0.4 for Cora and Citeseer, and to 0.1 for Pubmed, to prevent overfitting. For optimization, we employ the Adam algorithm with a learning rate of 10 -3 .The adversarial attacks are executed by perturbing the graph structure through the random addition and deletion of edges in the adjacency matrix. We quantify the extent of these perturbations using a contamination rate, which is defined as the ratio of the altered edge count to the total node pairs. Additionally, we propose a graph modification based on the original objective, as detailed in Eq. ( 5). The graph Laplacian is represented by a symmetric matrix, on which the eigh function is applied to compute the eigenvalues. The corresponding execution times are documented in Table 1 . For a more comprehensive experiment results and analysis of our findings, please refer to the Appendix A.3. 
discussion and conclusion: 
 In this study, we address a class of deep learning problems that incorporate eigendecomposition within their constraints or objective functions. Such problems are prevalent in applications like nuclear-norm regularized denoising, network compression, graph structure learning, and whitening normalization. The traditional approach requires performing eigendecomposition or singular value decomposition within the computational graph, which becomes the bottleneck in the training process. To circumvent this computation overhead, we introduce an amortized eigendecomposition framework integrating a relaxation eigen loss into the learning objective, which relies on a set of orthonormal vectors. These vectors are reparameterized efficiently through QR decomposition, thereby substantially reducing the computational cost of eigendecomposition in each iteration. Furthermore, the differentiable nature of QR decomposition allows for its seamless incorporation into the neural network training workflow. Our experimental results demonstrate that, when applied to network tasks, our algorithm not only accelerates the process but also maintains the precision of the conventional method with eigendecomposition. The method proves particularly beneficial as a differentiable top k eigensolver in environments where the backward gradient computation for top k eigendecomposition, such as in JAX or PyTorch, is not well-supported. Consequently, our approach provides a viable alternative for integrating top k eigendecomposition into neural networks.Limitations. While our method excels in scenarios where eigendecomposition operations on large matrices are embedded within neural networks, it is important to note that when employed as a pure numerical eigensolver, it does not offer any speed or precision advantages over conventional methods. Thus, its suitability is specifically aligned with applications where eigendecomposition is a substantial component of the neural network training process. In this experiment, we present the convergence analysis of our amortized eigendecomposition approach concerning the accurate estimation of singular values. We tackle the nuclear norm regularization problem outlined in Eq. (17) . In each iteration, we compare the diagonal elements of U T θ T θU with the exact singular values of θ. We conducted tests on θ of sizes 128 × 128, 256 × 256 and 512 × 512. The mean square error (MSE) is depicted in Figure 5 .Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: This is included in this paper in the Experiment section and Appendix A.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: The statistical significance may not be the primary focus. The evaluation and validation of our method rely on other metrics, such as accuracy and convergence rate.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We have indicated the type of GPU in the Experiment section.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: this paper is conducted with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This paper is a pure theoretic paper instead of an application, thus no societal impacts is involved.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [No] Justification: Not relevant to our research. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We have done this. CC-BY 4.0 Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: The results indicate that initially, the singular values obtained by our approach do not match the correct singular values within a small error. However, within just a few iterations, our approach rapidly converges to the correct values, demonstrating the effectiveness of the loss function in accurately estimating singular values.In this task, we evaluate the latent-space PCA method as outlined in Eq. ( 19) using the MNIST dataset. The architecture for both the encoder and decoder mirrors that of the nuclear norm regularization model, with the exception that there is a linear projection U U T aligning with the direction of the principal components. We aim to extract the first two principal components in the latent space.The experimental outcomes are depicted in Figure 6 . Specifically, Figure 6 (a) displays the convergence trajectories for both the reconstruction loss and the eigenvalue loss as defined in Eq. (19) . It is observed that the reconstruction loss curves for the conventional eigh function and our amortized eigendecomposition strategy are nearly indistinguishable. However, for the eigen-loss, our method initially registers lower values compared to the eigh function but eventually converges to equivalent values. This demonstrates the efficacy of the amortized optimization approach. can vary significantly. However, when regularization is applied (with η = 0), the distribution of the embedding becomes more uniform, indicating a more even spread across the principal components.The regularization applied to the eigenvalues also influences the sparsity of the network's architecture. This effect is depicted in Figure 6 (c), which shows the weight matrix of the last layer in the encoder. When η = 0, the network prioritizes learning weights that optimally reconstruct the images, resulting in a relatively dense weight configuration. Conversely, when η = 0, the eigenvalues corresponding to the principal components are encouraged to concentrate on the first few, leading to a decline in variance for the subsequent components. This, in turn, promotes sparsity in the weight matrix, as the network assigns less importance to the less variable components.This study investigates the robustness of graph convolutional networks (GCNs) [21] by conducting adversarial attacks on graph structures. These attacks involve perturbing the graph structure by randomly adding or removing edges in the adjacency matrix. To measure the magnitude of these perturbations, we introduce a contamination rate, which represents the ratio of the number of modified edges to the total number of node pairs.For each adversarial scenario, we commence by randomly determining a contamination rate within the range of [0, 0.02] for the Cora and Citeseer datasets, and [0, 0.001] for the PubMed dataset. This rate dictates the proportion of edges to be randomly added or removed from the graph. A greater contamination rate signifies a more severe attack. The nodes are partitioned randomly into training, validation, and test sets with respective proportions of 60%, 20%, and 20%. The classification accuracy on the test set is then recorded.The robustness of the GCN when subjected to graph contamination is illustrated in Figure 7 . Each data point on the graph corresponds to a specific attack instance and its resultant test accuracy. Additionally, we have applied a polynomial regression with a 5th-degree basis to model the general trend of the relationship between the contamination rate and the test accuracy. The findings indicate that the graph convolutional network exhibits enhanced robustness when convolutions are applied to graph signals of a lower rank. In this section, we provide the proofs for our theoretical findings.Lemma 1 (Jensen's inequality for convex functions). For a convex function f , given n real numbers x 1 , . . . , x n and non-negative weights w 1 , . . . , w n such that n i=1 w i = 1, Jensen's inequality is expressed as:EQUATIONThe equality holds if and only if one of the following conditions is satisfied:1. f is linear;2. x i 's are equal;3. one of w i 's is one and the rest are zero.Theorem 1 (Trace inequality with weight matrix) Consider a n-dimensional symmetric matrix A.Let M be a k-dimensional diagonal matrix with elements 0 < m 1 < m 2 • • • < m k . For an arbitrary matrix U ∈ U n×k with orthogonal columns, the following inequality always holds:EQUATIONThe equalities are achieved if and only if U contains the eigenvectors corresponding to the k largest (smallest) eigenvalues of A.Proof. We can consider the case where A has all nonnegative eigenvalues ε j ⩾ 0, without loss of generality, as we can always shift the diagonal elements of A to meet this condition:EQUATIONLet u ij be the i, j-th entry of U . The Right-hand-side (RHS) of the leftmost inequality can be expressed as,EQUATIONWe denote the actual eigenvectors of A as V with entries {ν jl }, and {λ s }, s = 1, . . . , n as the eigenvalues, we have:EQUATIONSubstituting this into (24), we obtain:EQUATIONEQUATIONEQUATIONWe defineEQUATIONthen the RHS becomes,EQUATIONand,EQUATIONIt's worth noting that for any s,EQUATIONThus we have,EQUATION)Applying the Cauchy-Schwartz inequality, we get,EQUATIONThe equality holds if and only if for all s and a fixed l,EQUATION)orEQUATIONThis is only true when g si are either 0 or 1, with at most one non-zero value for each s and i. In other words, G = {g si } must be a permutation matrix. This implies that U is equivalent to the eigenvectors of V , albeit with a possible permutation.Thus this inequality reaches a minimum if we rearrange the indices of the eigenvalues:EQUATIONThe last inequality is a result of the rearrangement inequality. This completes the proof of the first inequality.For the upper bound of Eq. ( 22), we apply the Cauchy-Schwartz inequality to Eq. ( 30) and obtain,EQUATIONEQUATIONHere again, the Cauchy-Schwartz inequality achieves the maximum when U is equivalent to the eigenvectors of V , up to a possible permutation. By applying the rearrangement inequality, we further obtain the final bound:EQUATIONThis yields the desired upper bound for the second inequality of the main result.Theorem 2 (Trace inequality with convex function) Let A be a given n-dimensional symmetric matrix and let U be a matrix of size n × k that resides on the Stiefel manifold. Suppose f : R → R be a monotonically increasing, convex function applied element-wise. The following inequalities hold:EQUATIONThe rightmost inequality becomes an equality if and only if U comprises the eigenvectors of A that correspond to the k largest eigenvalues, The leftmost inequality is met with equality when U is such that all diagonal elements of the matrix U T AU are equal.Proof. The rightmost inequality We first focus on the rightmost inequality, which provides an upper bound for the convex trace loss.To establish the rightmost inequality, let us consider the left-hand side (LHS) of the equation. Let's denote the diagonal elements of A as {a jj }. Consider the eigendecomposition of A as in Eq. ( 25), we can rewrite a jj in terms of the eigenvalues and eigenvectors:EQUATIONwhere ν ji is the (j, i) entry of its eigenvectors V . The diagonal elements of the matrix product U T AU can then be written as:EQUATIONSubstituting this result back into the LHS, we haveEQUATIONEQUATIONEQUATIONWe can express the summation of the elements in vector form as follows:EQUATIONFor all indices j, l and s, all the addends satisfy:EQUATIONwhich allows us to treat n l=1 n j=1 ν ji ν li u ls u js as a set of weights. By applying Jensen's inequality to the equation previously labeled as Eq. ( 49), we obtain:EQUATIONEquality holds if and only if all the eigenvalues {λ i } are identical, which is generally not the case for an arbitrary symmetric matrix A, or for each index s, there exists an index i such that,EQUATIONwith all other terms being zero. Therefore, U must be a permutation of V :EQUATIONSubstituting this result back into Eq. ( 53), we obtain:EQUATIONWith this, we conclude the derivation of the first inequality.The leftmost inequality. For the leftmost inequality,EQUATIONthe equality holds if and only if j u * ij u ij a jj equals for all i.A similar result for finding the k smallest eigenvalues can be obtained. The next corollary states the resultsLet A be a given n-dimensional symmetric matrix and let U be a matrix of size n × k that resides on the Stiefel manifold. Suppose f : R → R be a monotonically increasing, concave function applied element-wise. The following inequalities hold:EQUATIONThe rightmost inequality becomes an equality if and only if U comprises the eigenvectors of A that correspond to the k smallest eigenvalues.NeurIPS Paper Checklist38th Conference on Neural Information Processing Systems (NeurIPS 2024).Although the findings in this paper can be readily extended to the complex space, we specifically focus on the real space throughout our study. 


Decision: While the use case is somewhat niche (for neural networks relying on eigendecomposition during training), the reviewers found the paper clear, the idea novel and well supported by the experimental setting. The reviewers raised concerns regarding the scalability of the method that were addressed by the authors in their rebuttal. I recommend for acceptance.

Reviews: review: paper_summary: paper_summary: The paper describes a novel method to circumvent the need for explicit eigendecomposition during neural network training. The central insight is that one can simply learn estimates of the eigenvectors (parameterized by a QR decomposition) of interest alongside the original loss function via gradient descent. 
To do so the author's add an auxiliary (and differentiable) loss whose optima is obtained by the appropriate set of eigenvectors.
A summary of the presented results:
- The author's provide both theoretical (proofs) and empirical evidence that the proposed estimation algorithm converges to the true eigenvectors.
- 3 settings where SVD or eigendecomposition are incorporated in neural network training schemes are investigated empirically, in all settings the proposed scheme yields a significant speed up relative explicit matrix decompositions.
strength_weakness: strength_weakness: - Small scale experiments: Each of the three empirical experiments are performed with relatively small networks and datasets. While this is not necessarily a weakness in and of itself, my intuition is that the speed-up benefits will be more marginal as network size increases. This is because, for example as the depth/width of the encoder/decoder increases in the latent space PCA setting, the eigendecomposition becomes less of a bottleneck as more time is spent propogating activations through the network. If this is the case I think it merits discussion in the paper. On the same note it would be nice to have more assurances that the amortization method finds correct solutions (a la Figure 5a) on larger scale problems.
questions: questions: - I would appreciate if the author's could comment on the scalability of the method (or more precisely, how the computational benefits scale with the architecture size).

- I found myself slightly confused by Eq. 19. The numerator of the regularizer encourages high variance along the first two PCs (and synergistically encourages $U$ to converge toward the first two PCs). Doesn't this, alongside the projection and reconstruction loss, encourage the covariance matrix to be low rank (as any variance in lower variance eigenspaces will not be accessible to the deocder)? Does homogeneity need to be preserved simply because without trace normalization the weights of the encoder would grow large and the weights of the decoder would decay? Sorry for getting caught up in this as it really is a detail, but wouldn't this also be a case where we really only care that $U$ forms a basis for the space of the top 2 PCs (because the first linear layer of the decoder could "undo" any such rotation of the space if it were advantageous)? If this is the case, I assume then it would be sufficient to set $M=I$?
limitations: limitations: Yes.
paper_summary: paper_summary: This paper proposes a method named "amortized eigendecomposition" to replace the computationally costly SVD operation in settings where an eigendecomposition is required during neural network training. The proposed method introduces a loss term ("eigen loss") and replaces the full SVD with the less computationally expensive QR decomposition at each iteration. A theoretical analysis shows that the desired eigenpairs are obtained as optima of the eigen loss. An empirical analysis on several tasks, including nuclear norm regularization, latent-space PCA, and graph adversarial learning, demonstrates that the proposed method attains significant improvements in training efficiency while producing nearly identical outcomes to conventional approaches.
strength_weakness: strength_weakness: A discussion of scalability and comparison with alternative methods for fast SVD are two significant points that are missing in this paper. The experiments report results on small-to-moderate matrix dimensions.
questions: questions: How does amortized eigendecomposition scale to larger dimensions?
How does it compare to other fast-SVD methods? In particular, I have randomized SVD in mind, which is known to scale very well. Is there a trade-off, or is there a reason for why randomized SVD is not applicable in this setting? Or is it just still more computationally expensive because it needs to be done in each iteration?
limitations: limitations: The authors have addressed the limitations of their method reasonable well. As mentioned in the sections above (weaknesses/questions), an additional discussion of scalability and comparison to other fast/approximate SVD methods would be nice to have.
paper_summary: paper_summary: This paper proposes a method that decouples the eigendecomposation calculation from training process, instead, they use a eigen loss to jointly optimize it with the training loss of the neural network as a nested optimization loop. The proposed method can speed up the training process of problems that incorporate eigendecomposition within their constraints or objective functions.
strength_weakness: strength_weakness: The theoretical support of this paper is sufficient. However, in the experiment section, except for the speed up ratio, I think the author should report the loss value and the corresponding task performance for a comprehensive comparison. It is important to verify whether the accelerate method can maintain the performance.
questions: questions: I have no questions.
limitations: limitations: The author clearly state the limitation of this paper.
paper_summary: paper_summary: In this paper, the authors proposed a novel framework to integrate eigendecomposition efficiently into neural network training, which is typically costly in the backward process.

The key insight is to compute the decomposition together with training (since the decomposition also changes during training). Rather than forcing exact eigendecomposition $U\Sigma U^T=A_\theta$ every iteration, the authors treat the "eigenvectors" U as optimizable variables and leverage an eigen loss that attracts U to the true eigenvectors of $A_\theta$. The authors include a thorough analysis of different "eigen loss" inspired by the trace inequality (Eq.9, 11); they also provided an efficient implementation that re-parameterizes the orthogonal basis U into an unconstrained matrix $W$ via QR decomposition $U=\text{QR}(W)$.

For convergence analysis, the authors empirically showed that their eigen loss cooperates well with different optimizers for solving eigendecomposition; furthermore, they tested their method on several machine learning tasks and observed a significant speed-up from 1.4x to 600x.
strength_weakness: strength_weakness: The idea of doing eigendecomposition amortized is interesting. However, the stability and convergence rate of this method are not adequately discussed in the paper. Also, the design of eigen loss is very much "handcrafted", and I look forward to more flexible methods.
questions: questions: 1. 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?
2. Can this framework be applied to other aspects such as improving the stability of SVD in network training?
limitations: limitations: 1. The proposed method seems to be designed for relatively simple machine learning tasks. But 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.

2. Limited to eigendecomposition which is only applicable to symmetric matrices. For example, in section 5.2 the authors have to approximate the nuclear norm $\|\theta\|_\star$ with $\sum_i \|\theta u_i\|$ since $\theta$ is not necessarily symmetric. So perhaps an "amortized SVD" would be more significant.

Comments: title: ok
Response to Rebuttal.
comment: thanks for adding this info
Thanks for your detailed explanation, I will increase my score to 8.
Thank you to the author's for their clarifying comments. I keep my original score.
Thanks for the comprehensive rebuttal, especially **Q3** illustrating that the eigen loss can act as a regularizer and preserve networks' performance. I am increasing my rating to 5.