Amortized Eigendecomposition for Neural Networks
Authors: Tianbo Li
Zekun Shi
Jiaxi Zhao
Min Lin
Venue: NeurIPS 2024 poster (2024)
Keywords: Eigendecomposition
SVD
Amortized optimization.
Primary area: deep_learning_architectures
TL;DR: Don't use eigendecomposition or SVD in neural networks, there is a better way doing that

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.

Amortized Eigendecomposition for Neural Networks
Tianbo Li Zekun Shi Jiaxi Zhao?, Min Lin'
SEA AI Lab
School of Computing, National University of Singapore 3 Department of Mathematics, National University of Singapore "litbSsea. com
Abstract
Performing eigendecomposition during neural network training is essential for tasks such as dimensionality reduction, network compression, image denoising. and graph learning. However, cigendecomposstion is computationally expensive as it is orders of magnitude slower than other neural network operations. To address this challenge, we propose 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 more affordable QR decomposition at each iteration. Theoretical analysis guarantees that the desired eigenpair attained as optima the eigen loss. Empirical studies on nuclear norm regularization, latent-space principal component analysis, and graphs adversarial learning demonstrate significant improvements training efficiency while produc- ing 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.
Introduction
Figendecompositiinn a fundamental technique in linear algebra that finds applications across numer- scientific domains ranging from quantum many- body 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 reduc- tion [6, 44, 27, 38], network compression [17, 30),
10a
Moan CPU Time Ratio to matmul
qr
2.1E5
1.5E5
eigh
SVD
1.3E4
10a
2.0E3
image denoising [13, 12, 14), graph adversarial learning [10, 18, 47). By uncovering the struc- ture of networks, eigendecomposition allows to enforce low-rankness, ensuring generalization, robustness, and computational efficiency. Eigen- decomposition also instrumental in the spectral analysis of graphs, where it can detect community structure, which is essential in spectral graph neu- ral networks. The ability of cigendecomposition to detect the intrinsic matrix structures and prop- erties makes a valuable tool various machine learning tasks with neural networks.
10²
1.0E3
10a
PyTorch
JAX
Figure 1: A comparison that illustrates the for- ward execution time of three linear algebra op- erations: qr, eigh, and svd, when performed on a 10000 10000 matrix using PyTorch and JAX. The presented values represent the mean ratios of the execution time relative to that of matrix multiplication (matmul) of 100 runs.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Despite its straightforward definition, the computation of eigenvalues is quite challenging. Eigende- composition 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 presents a com- parative 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 by 4 to 5 orders of magnitude- rellatve matrix multiplication. This substan- tial 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 the inner loop and the neural network's loss minimization serves as the outer loop. Notably, it is not always easy to be aware this inner loop of eigendecommosition, 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:
Figendecomposition within a neural network does not have to reach full convergence during each training step; simply needs contribute to the desired outcome by the end of the training process.
In this paper, we present a novel approach named "amortized eigendecompositionn for training neural networks that require eigenvalues or cigenvectors. Instead of using computationally expensive cigendecomposition 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. achieves nearly identical performance to traditional methods but with significantly improved speed.
2 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 constrained optimization problem:
min f (ho (X).
s.t. V AV - A
Here, the encoder hg maps the data X E RX into a latent space. In addition to the latent representation, the loss function ff also incorporates the eigenpair A and V of symmetric matrix A. The matrix A can be constructed from he(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 O. The computational graph for each iteration of such model structure can be written as,
X he XX)
A (V,A)
(ho(
(2)
encoding
eigendecompositionn
downstream task
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 cigendecomposition constraint of after each update of 0. This structure encompasses a wide range of learning problems. We present several representative examples and investigate them in our numerical experiments next.
Nuclear Norm Regularization The nuclear norm of a matrix 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 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 the network and approximating the original network with fewer parameters [17, 30). The objective function of this type of problem can be written as,
min (ho(()) + =]|00|-
0
Here, is a regularization coefficient that controls the rank of the parameter matrix. In classical methods, the nuclear norm ||O|| is usually calculated via singular value decomposition.
Whitening in Neural Networks Whitening is a transformation technique extensively utilized in neural networks to standardize features by ensuring they have zero mean and unit variance, and 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 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:
min w,
Dec (VVTEEnce
Ence (x)
X
(4)
where V is spanned by the first several largest eigenvectors of the covariance of Ence (xce(() (X) and 0, w represent the parameters of the decoder and encoder, respectively.
Graph Structure Learning Graph structure learning is 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, follow the approach proposed in [18] and the objective function can be formulated as:
min GNNN (X. Ż) -Y + + - + B|| l O.L
(5)
Here L represents a low-rank Laplacian matrix that approximates the original graph Laplacian L. The three parts of the loss function correspond 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.
3 Differentiable Optimization for Figendecomposition
Before introducing the proposed method, let's begin by examining a more straightforward case: determining the eigenvalues of symmetric matrix A Raxn through constrained optimization. The eigenvalues of A are denoted as > > X > ... > Xn The largest eigenvalue of A can be obtained by directly maximizing the Rayleigh quotient:
u Au. \ max
uER"
u Tu
where the maximum value corresponds to the largest eigenvalue and the associated eigenvector is given by the normalized at the optimum. This method can be generalized 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 cigenvectors. This can be found by solving the constrained optimization problem:
Tr AU)
max UEUnkk
Although the findings in this paper can be readily extended to the complex space, we specifically focus on the real space throughout study.
3

---

-2
2
2
(a) Trace loss
(b) Brockett's ' cost function
(c) Convex trace loss
Figure 2: An illustration on disrupting the rotational symmetry of the trace loss. We aim solve
0.8 0.2
eigenvectors for a 2-dimensional symmetric matrix A
with three loss functions. We which is subject to the constraint x²² 1² 1.
0.2
0.1)
parameterize orthonormal matrix U (
The plot displays the contours of the landscapes of three different loss functions as they vary with and y: (a) trace loss tr(U tr(UAUT)) (b) Brockett's cost function tr ((MUAUT) where M diag(1, 0.2) and (c) convex loss function (U (UA 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.
where the Stiefel manifold Unxk (U E Rank UT 1} is the set of x 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 the rota- tional symmetry of the trace function that for any orthonormal matrix C E Uk*k the equation tr(UT AU) = tr AUC) always holds. As a result, there are an infinite number of solutions the optimization problem Eq. (7), all spanning the same subspace the desired sets eigenvectors. A visual illustration is provided in Figure 2a.
To accurately obtain the cigenvalues and eigenvectors, it is necessary 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 the diagonal elements of the matrix product U 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 exact eigendecomposiitoon. We will elaborate these methods.
The Brockett's Cost Function The first modification made to the trace loss is Brockett's 's cost function [3]. It enables the extraction eigenvectors and eigenvalues by solving the following optimization problem:
tr(MU AU).
max UEURK
Here M diag(m1, m2,.... M,....,ma .,)) E k is 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 assign distinct weights to the diagonal elements of the product U 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
0 < mi < m2
< my. This ordering allows M to

---

elements. A practical choice could be m, = i/k or m. = i as suggested by [3]. The next theorem illustrates why this approach can yield the exact eigendecomposition.
Theorem (Trace inequality with weight matrix), Consider a n-dimensional symmetric matrix A. Let M be k-dimensional diagonal matrix with elements 0 < mı < 12 < my. For arbitrary matrix U E U with orthogonal columns, the following inequality always holds:
k
Σ mm--A ma-L41de M AU) &
Σ mide.
(9)
The equalities are achieved if and only U contains the eigenvectors corresponding to the l 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 Schwarz inequality, which is presented in the Appendix for reference. A more generalized version of this trace incquality 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º AU. This perturbation also disrupts the rotational symmetry inherent in the trace loss. The convex nature of 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 le largest eigenvalues, is expressed as:
max UetT*N
trf(UT AU)
(10)
The optimal U, that achieves the maximum in the above objective are the eigenvectors corresponding to the k largest cigenvalues, and the eigenvalues can be obtained by the diagonal elements of U. AU.. The next theorem provides formal validation of this assertion.
Theorem 2 (Trace inequality with convex function). Let A be given n-dimensional symmetric matrix and let U be a matrix of size n x k that resides on the Stiefel manifold. Suppose f R. - be a monotonically increasing, convex function applied element- wesee wise. The following inequalities hold:
1s (Mraa) < <I((F( & S(N).
(11)
The rightmost inequality becomes an equality and only 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 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 include: 1; and fxx) tan(z) [0, π /2) For finding
fx) cxp(x) on R; fx) _ on R where the i smallest eigenvalues, a simple modification can be made by replacing the function f with an element- dementtwise 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.
The Amortized Eigendecomposition Approach
The proposed amortized cigendecomposition approach aims to modify the eigendecompositiin operation within the neural network's 's computational graph, as illustrated in Eq. (2). This replacement involves two steps: First, the set of eigenvectors form matrix the Stiefel manifold reparameterized through computationally efficient operations such as QR decomposition. Then, the loss function is adjusted ensure that its optimal solutions precisely correspond to the eigendecomposition. The computational graph for this amortized eigendecomposition method is formulated as follows:
(X)
> (ho( +ng (U, U
(12)
encoder
reporamttrtiaation
In 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
5

---

Algorithm 1
The conventional cigendecomposi- tion neural network outlined Eq. (2)
Algorithm The amortized eigendecomposition technique outlined in Eq. (12)
Input: Dataset X, encoder he. task fw:
Input: Dataset X, encoder hag task fwo and fy;
I: Initialize model parameter and wi
1: Initialize model parameter 0, and W
2: while not converged do
2: while not converged
3: compute A from (xx);
3: compute A from ho X);
4: V. A Tx==ighh eigh( A);
4: U QR(W).
diag AU):
5: compute ( A):
5: compute [ (X)) U. ug(A
update 0, w by gradient descent;
6: update 0, w. W by gradient descent;
7: end while
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 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 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:
U QR(W).
(13)
In this formulation, W is dimensionally consistent with U. QR decomposition is more computa- tionally efficient than cigendecomposition 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-propagatiin 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., we can transform this loss into a regularized version incorporating the eigen loss. This relaxation allows us to forego the need for explicit eigendeconposition 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:
(V, A) is largest eigenpair
min (X),
(X), A)
min [ho (X), - AOU). e.W
(14)
In this reformulation, U NjR is reparameterized via QR operation 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 the precise eigendecomposition of the matrix Ag-
The second type of optimization problem involves scenarios where the model loss is independent of the cigenvalues 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 trace penalty term, such as Brockett's 's cost, which is given by
min (ho V)
V largest eigenvectors
min O.W
StopGrad
(15)
In this formulation, stop gradient operation is applied to A, since the eigen loss involves A which relies on the parameter O. By introducing the stop gradient operation, we prevent this regularization term from propagating gradients back to 0.

---

Brockett's Cost Function
convex trace loss
10a
10
MSE of eigenvalues
10-1
MSE of eigenvalues
10-1
10-3
10- -
10-5
10-5
10
10-7
10
10
5000
10000
15000
20000
0
5000
10000
15000 20000
Step
Step
adam
adamax
yogi
sgd
lbfgs
Figure 3: Convergence analysis on finding 50 largest eigenvalues random 1000 1000-dimensional symmetric matrices. (a): Convergence curves using Brockett's 's cost and convex trace loss (f = x 5..55. ). (b) The fine-tuning convergence series of similar matrices.
Besides the above relaxations, further simplification is possible if the loss function conforms to the structure of Brockett's 's cost function convex trace function. That is, when the model loss fw is non-uniform linear combination of the eigenvalues, or it is monotonic and convex (or concave) to the cigenvalues, the trace penalty term is unnecessary. The loss function will inherently converge to the correct cigendecomposition. This principle exemplified in the context of the nuclear norm regularization problem, which will be illustrated later. This be formulated as:
min (X),
and
min / (X)(X))U,E) 0.W
(16)
non-uniforn linear combination of
5 Experiments
In this section, we present an evaluation of our approach, focusing 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 encoder 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 single NVIDIA A100 GPU with 40GB memory. The two fundamental questions we investigate are as follows:
regularized auto-
Does our approach accurately identify the eigendecomposition and singular value decomposition?
How does the efficiency of our method compare that of traditional techniques?
5.1 Convergence
In this experiment, we evaluate the numerical error and convergence speed of our algorithm applied solving eigendecomposstion. We randomly generate ten symmetric matrices size 1000 1000. The first 50 eigenvalues of these matrices range from 1 to 50, while the remaining eigenvalues lie between 0 and Our objective is compute the 50 largest eigenvalues by minimizing Brockett's ' cost function and convex trace loss (we adopt f(r) 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.
as the convex function). To achieve this,

---

Table 1: Evaluation of execution times per iteration on three tasks.
Backbone time (s/iter)
Backbone+ eigh/svd time (sfiter)
Backbone+ our method time (s/iter)
Speed-up
Task
Dimension
4-
to
4
12
62
128 x 128
5,275E-2
8.323B-2
6.025E-2
4.06x
Nuclear norm
256 x 256
5.600F-2
1.209E-1
6.080F-2
1
regularization
512 x 512
7.186B-2
2.616E-1
7.366E-2
105.4x
256 2
4.178E-3
1.446B-2
1.117E-2
1.47x
Latent-space
512 x 2
6,792B-3
2.918B-2
2.224E-2
1.45x
PCA
1028 x 2
1.434E-2
7.018F-2
5.467F-2
1.39x
2708 x 16
1.021E-3
1.769F--2
1.732F-3
23.4x
Low-rank
3312 16
1.367E-3
2.825B-2
2.498E-3
23.7x
GCN
19717 x 16
1.931E-2
4.941E+0
2.731E-2
615.2x
This result demonstrates that both loss functions capable of identifying the correct eigenvalues with small numerical error of 10-9. However, there is 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, 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.
5.2 Nuclear Norm Regularization
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 auto-encoder on the MNIST dataset by minimizing the reconstruction loss with a nuclear norm regularizer applied to the weight matrix 0 Ra*k of the encoder's 's last layer. We employ relaxation technique to the original problem defined in Eq. (16), which can be expressed as:
min Dec (Ence - we0
?) 0
min w.O.W
Dec (Enca X +1
Σ 1=1
0u, || (17)
where u,'s '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 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.
2-layer MLP with hidden layer
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 backhone utilizing amortized eigendecomposition, i.e. RHS of Eq. (17) are reported in Table 1. We denote these execution times as to. 41, and fa respectively, and define the speed-up ratio for our approach relative to the svd as:
speed-up f to
(18)
1/ to
This ratio represents the improvement in execution speed of cigendecomposition method com- pared to the standard svd, relative to the baseline backbone performance.
5.3 Latent-space Principle Component Analysis
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

---

Depth
Depth 16
backbone
backbone
Well cock time (xitter)
backbone+elgh
We llclkookkk time (silter)
backbone-+eigh
0.25
backbone+ours
0.25
backbone+tors
0.20
0.15
0.15
0.05
0.05
128
256
512
1024 2048 4096
128
256
512
1024
2048 4096
Hidden dimension
Hidden dimension
Depth 32
Depth 64
0.35
backbone
backbone
Wall clock time (x/iter)
0.25
backbone+
Wall-cledd clock time (sitte)
0.30
backbone-teigh
backbone+ours
backbone+ours
0.25
0.20
0.20
0.15
0.10
0.10
0.05
0.05
128
256 512 1024 2048 4096
128
256
512
1024 2048 4096
Hidden dimension
Hidden dimension
Figure 4: A comparison of scaling in the latent PCA task using the Celeb-A-HQ (256x256) dataset. The backbone autoencoders used in this study consist entirely of fully- fully-connected layers with ReLU activation, all maintaining the same dimensions. Between the encoder and decoder, applied both an eigen solver from the JAX eigh function and our amortized eigendecomposition method. We varied the depth of the autoencoders across 8, 16, 32, and 64 layers, and explored dimensionalities of 128, 256, 512, 1024, 2048, and 4096. The results present the average execution time per iteration over 100 runs. Notably, the largest model tested, featuring an autoencoder with 64 layers and dimension of 4096, comprises to 1.0 billion parameters.
additional term resembles Brockett's cost function while the trace of the covariance matrix ensures homogeneity:
Dec. (UCa x (MU Cov cho (X))U)
min w W.O.W
(19)
tr(cov(ha(XX)
where 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 More results and analysis are provided in Appendix A.2. The experimental results demonstrate that approach achieves an average training speed improvement of 40% compared to conventional eigendecomposition approach.
Additionally, we conducted a scalability study, with the results presented in Figure 4. This study examined the scaling behavior 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 our tests, a 64-layer autoencoder with 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. 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 culle-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-iiiee further underscoring the results shown in Figure 1.
5.4 Adversarial Attacks on Graph Convolutional Networks
In this study, we explore the robustness of graph convolutional networks (GCNs) by implementing adversarial attacks on graph structures. The objective of this problem is described in Eq. (5). Our approach simplifies the attainment of low-rank structure by optimizing:
min
|[CNN(XX.1 L) YIl, mtt((MM LU),
(20)
where = diag(U 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 approximation of under the Frobenius norm which corresponds to the terms [l.. Σ = in Eq. (5). This formulation allows the GNN to operate on 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 classifica- tion 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 the original objective, 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 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 the training process. To circumvent this computation overhead, we introduce an amortized eigendecomposition framework integrating 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 differentiabbee 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 & eigensolver in environments where the backward gradient computation for top k eigendecomposition, such in JAX or Py Torch, is not well-supported. Consequently. our approach provides a viable alternative for integrating top & 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 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 substantial component of the neural network training process.
10

---

References
[1] P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, 2008
[2] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Mandeerman-Milng and Qiao Zhang. JAX: composable transformations Python+NumPP programs, 2018.
[3] Roger W Brockett. Dynamical systems that sort lists, diagonalize matrices, and solve linear programming problems. Linear Algebra and its applications, 146:79-91, 1991.
[4] HanQin Cai, Jialin Liu, and Wotao Yin. Learned robust pea: scalable deep unfolding approach for high-dimensional outlier detection. Advances in Neural Information Processing Systems, 1977-77996,,,, 2021.
[5] Emmanuel J Candès and Terence Tao. The power of convex relaxation: Near-optimal matrix completion. IEEE transactions on information theory, 56(5):2053-2080, 2010.
[6] Tsung-Hat Han Chan, Kui Jia, Shenghua Gao, Jiwen Lu, Zinan Zeng, and Yi Ma. Pcanet: A simple deep learning baseline for image classification? IEEE transactions on image processing. 24(12):5017-5022 2015.
[7] DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojeiech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020.
[8] Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, et al. Natural neural networks. Advances in neural information processing systems, 28, 2015.
[9] Jed A Duersch, Mciyue Shao, Chao Yang, and Ming Gu. A robust and efficient implementation of lobpeg. SIAM Journal on Scientific Computing, 40(5):C655-C677) 2018.
[10] Negin Entezari, Saba A Al-Sayouri, Amirali Darvishzadeh, and Evangelos E Papalexakis. All you need is low (rank) defending against adversarial attacks on graphs. In Proceedings of the 13th international conference on web search and data mining, pages 169-177, 2020.
[11] Farzan Farnia, Jesse M. Zhang, and David Tse. Generalizable adversarial training via spectral normalization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[12] Shuhang Gu, Qi Xie, Deyu Meng, Wangmeng Zuo, Xiangchu Feng, and Lei Zhang, Weighted nuclear norm minimization and its applications to low level vision. International journal of computer vision, 121:183-208 2017.
[13) Shuhang Gu, Lei Zhang, Wangmeng Zuo, and Xiangchu Feng. Weighted nuclear norm min- imization with application to image denoising. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2862-2869, 2014.
[14] Cho-Jui Hsich and Peder Olsen. Nuclear norm minimization via active subspace selection. In International Conference on Machine Learning, pages 575-583. PMLR, 2014.
[15] William W Hsich. Nonlinear principal component analysis by neural networks. Tellus A. 53(5):599-66 615, 2001.
[16] Lei Huang, Dawei Yang, Bo Lang, and Jia Deng. Decorrelated batch normalization. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition, pages 791-800, 2018.

---

[17] M Jaderberg, A Vedaldi, and A Zisserman. Speeding up convolutional neural networks with low rank expansions. BMVC 2014-Proceedingg of the British Machine Vision Conference 2014. British Machine Vision Association, 2014.
[18] Wei Jin, Yao Ma, Xiaorui Liu, Xianfeng Tang, Suhang Wang, and Jiliang Tang. Graph structure learning for robust graph neural networks. In Proceedings of the 26th ACM SIGKDD international conference knowledge discovery & data mining, pages 66-74, 2020.
[19] Anton Johansson, Niklas Engsner, Claes Strannegård, and Petter Mostad, Improved spectral norm regularization for neural networks. In Modeling Decisions for Artificial Intelligence: 20th International Conference, MDAI 2023, Umeå, Sweden, June 19-22, 2023, Proceedings, page 181-201, Berlin, Heidelberg, 2023. Peringer-Verlage
[20] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[21] Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference Learning Representations, 2017.
[22] Andrew V Knyazev. Toward the optimal preconditioned eigensolver: Locally optimal block preconditioned conjugate gradient method. SIAM journal on scientific computing, 23(2):117- 541, 2001.
[23] Sandeep Kumar, Jiaxi Ying, José Vinícius de M Cardoso, and Daniel P Palomar. A unified framework for structured graph learning via spectral constraints. Journal of Machine Learning Research, 21(22):1-60, 2020.
[24] Rongjie Lai and Stanley Osher. A splitting method for orthogonality constrained problems. Journal of Scientific Computing, 58:431-444 2014.
[25] Shen Li and Bryan Hooi. Neural pea for flow-based representation learning. In Proceedings of the Thirty-Firs International Joint Conference on Artificial Intelligence, IJCAI-22, pages 2235.. International Joint Conferences on Artificial Intelligence Organization, 7 2022.
[26] Xin Liang, Li Wang, Lei-Hong Zhang, and Ren-Cang Li. On generalizing trace minimization principles. Linear Algebra and its Applications, 656:483-509, 2023.
[27] Zhouchen Lin, Minming Chen, and Yi Ma. The augmented lagrange multiplier method for exact recovery of corrupted low-rank matrices. arXiv preprint UrXiv:/009.555 55055, 2010.
[28] Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, (5(1):503-528, 1989.
[29] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[30] Marc Masana, Joost van de Weijer, Luis Herranz, Andrew D. Bagdanov, and Jose M. Al- varez. adaptivee deep network compression. The IEEE International Conference on Computer Vision (ICCV), Oct 2017.
[31] Leon Mirsky. A trace inequality of john von neumann. Monatshefte für mathematik, 79(4)3303- 306, 1975.
[32] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral Normaliza- tion for Generative Adversarial Networks. 2 2018.
[33] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[34] Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pca. directions (by accident). In Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition, pages 12406-12415, 2019.
12

---

[35] Axel Ruhe. Perturbation bounds for means of eigenvalues and invariant subspaces. BIT Numerical Mathematics, 10(3):343-354, 1970.
[36] Matthias Seeger, Asmus Hetzel, Zhenwen Dai, Eric Meissner, and Neil D Lawrence. Auto- differentiating linear algebra. arXiv preprint arXiv: 2017.
[37] Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening and coloring batch transform for gans. arXiv preprint arXiv:
2018.
[38] Kai Sun, Jiangshe Zhang, Hongwei Yong, and Junmin Liu. Fpeanet: Fisher discrimination for principal component analysis network. Knowledge-Bast Based Systems, 166:108-117 2019.
[39] Lloyd N Trefethen and David Bau. Numerical linear algebra. SIAM, 2022.
[40] John von Neumann. Some matrix-inequalities and metrization of matrix-space. John von Neumann Collected Works, 1937.
[41] Wei Wang, Zheng Dang, Yinlin Hu, Pascal Fua, and Mathieu Salzmann. Backpropagation- friendly eigendecomposition. Advances in Neural Information Processing Systems, 32, 2019.
[42] Wei Wang, Yan Huang, Yizhou Wang, and Liang Wang. Generalized autoencoder: A neural network framework for dimensionality reduction. In Proceedings of the IEEE conference computer vision and pattern recognition workshops, pages 490-497, 2014.
[43] Zaiwen Wen and Wotao Yin. A feasible method for optimization with orthogonality constraints. Mathematical Programming, 142(1):397-434, 2013.
[44] John Wright, Arvind Ganesh, Shankar Rao, Yigang Peng, and Ma. Robust principal component analysis: Exact recovery of corrupted low-rank matrices via convex optimization. Advances in neural information procexsing systems, 22, 2009.
[45] Chenggang Yan, Zhisheng Li, Yongbing Zhang, Yutao Liu, Xiangyang Ji, and Yongdong Zhang. Depth image denoising using nuclear norm and learning graph model. ACM Transactions Multimedia Computing, Communications, and Applications (TOMM), 16(4):1-17, 2020.
[46] Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. Advances in neural information processing systems, 31, 2018.
[47] Yanqiao Zhu, Weizhi Xu, Jinghao Zhang, Qiang Liu, Shu Wu, and Liang Wang. Deep graph structure learning for robust representations: A survey. arXiv preprint 14:1-1,2021
13

---

Appendix A
More Experimental Results
A.1 Nuclear Norm Regularization
128
10-1
256
512
10-2
MSE of singular values
10
10
10-5
10-6
10
2000
4000
6000
8000
10000
Step
Figure 5: The convergence curve of the singular values.
In this experiment, we present the convergence analysis of our amortized cigendecomposition approach concerning the accurate estimation of singular values. We tackle the nuclear norm reg- ularization problem outlined in Eq. (17). In each iteration, we compare the diagonal elements of UToTou OU with the exact singular values of O. We conducted tests on O of sizes 128 128, 256 x 256 and 512 x 512. The mean square error (MSE) is depicted Figure 5.
The results indicate that initially, the singular values obtained by our approach do not match the correct singular values within small error. However, within just few iterations. our approach rapidly converges to the correct values, demonstrating the effectiveness of the loss function in accurately estimating singular values.
A.2 Latent-space Principle Component Analysis
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 linear projection UUT aligning with the direction of the principal components. We aim to extract the first two principal components the latent space.
The experimental outcomes depicted in Figure 6. Specifically, Figure 6 (a) displays the conver- gence trajectories for both the reconstruction loss and the eigenvalue loss as defined in Eq. (19). is observed that the reconstruction loss curves for the conventional eigh function and our amortized eigendecompossitin strategy are nearly indistinguishable. However, for the eigen-loss, our method initially registers lower values compared the eigh function but eventually converges to equivalent values. This demonstrates the efficacy the amortized optimization approach.
Figure 6(b) illustrates the distribution of features in the latent space along the two principal compo- nents for cases where 7) = and r) = 1. The eigenvalues correspond feature dimensions. Without regularization (when 7) 0), the scales of the principal component axes
total variance across the
14

---

(a)
(b)
Raconatuccion Lozs
xxx 10
1.00
0.75
D-256
10-1
0.50
60
0.25
cum
xx
0.00
49g
10a
Recunatuction cetonnrucionnoo Less
2 20
1.00
0.75
D-512
10-1-
0.50
1x-0
0.25
(c)
n=1
xx
0.00
10a
10a
Reconstruction LOSs
10
Lats
0.75
D-1024
10-3-
0.50
1x-2
0.25
0
1x
0.00
Figure 6: Experimental results of Latent-space PCA on MNIST dataset. (a) Convergence curves. Convergence curve reconstruction loss. Second column: Convergence curve eigen loss. (b) Principle components in latent space: The principle components of features in latent space for 0 and = 1. (c) The sparsity of the network structure. The weight matrices of the second layer of the encoder. The color indicates the scale of the absolute values of the weight matrix ranging from 0 to 1.
First column:
can vary significantly. However, when regularization applied (with v = 0), the distribution of the embedding becomes more uniform, indicating more even spread across the principal components.
The regularization applied the eigenvalues also influences the sparsity 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 77 = 0, the network prioritizes learning weights that optimally reconstruct the images, resulting in relatively dense weight configuration. Conversely, when 1j = 0. the eigenvalues corresponding 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.
A.3 Adversarial Attacks on Graph Convolutional Networks
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, introduce contamination rate, which represents the ratio of the number of modified edges to the total number of node pairs.
For each adversarial scenario, commence randomly determining 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 accuracy on the test set is then recorded.
with respective proportions of 60%, 20%, and 20%. The classification
The robustness of the GCN when subjected to graph contamination is illustrated in Figure 7. Each data point on the graph corresponds to specific attack instance and its resultant test accuracy. Additionally, we have applied polynomial regression with 5th-degree basis 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 applied to graph signals of lower rank.
15

---

GCN
Rank-8
Rank=16
Rank-32
Ronk-64
0.9
Test Accurar
0.8
Thes Accuracy
UISE sccurscy
0.8
CORA
0.7
O.Y
Du7
00
0.01
0.02
0.00
0.01
02
0:00
0.
5.00
0,01
0.52
0.00
0.01
0.
Contamination Bate
Contamination Rate
Contumination Pate
Contamination Rate
Contaminw Race
GCN
Rank-8
Rank-16
Rank-=32
Rank=64
0.8
0.0
Thi Accuncy
Test Accuroy
Tel:: AcCurey
Clesser
7
7-
0
0.5
0.5
0.6
Contamination
GCN
Rankc=
Rank=16
Rank:=32
Rank=66
Thst Acciracy
Tost Accuracy
PubMed
Tee
Tes
5 00000
B 0.0005
0:0010
0. 00000
2 0005
0.0010
5 0000
B 0.0005
0.0010
0.00
0.0005
0 0.000
0.0000
0,0005
3,0010
Contamiation Rate
Contamiation Rate
Comtaminat on Rake
Contamination Rale
Contamination Rall
GCN attack
low-rank GCN attack
GCN fitting curve
low-rank GCN fitting curve
Figure 7: Experimental results of graph structure adversarial attacking
16

---

Appendix B Proofs
In this section, we provide the proofs for theoretical findings.
B.1 Preliminaries
Lemma (Jensen's inequality for convex functions). For convex funcctio f. given a real numbers
Ii...., and non-negative weights Wi,...y Wn such that Σi=1 = 1, Jensen's 's inequality is expressed as:
Σ was
Σ w.f (zi).
(21)
The equality holds if and only if one of the following conditions is satisfied:
1. f is linear;
2. x; are equal;
3. one of w; is one and the rest are zero.
B.2 Proof of Theorem 1
Theorem (Trace inequality with weight matrix) Consider n--dimensional symmetric matrix A. Let M be k-dimensional diagonal matrix with elements 0 < m, .2 mg- For an arbitrary matrix U u with orthogonal columns, the following inequality always holds:
Σ Ei--4tie & tr(MUU AU) &
m
(22)
The equalities achieved if and only U contains the eigenvectors corresponding the k largest (smallest) eigenvalues of A.
Proof. We consider the case where A has all nonnegative eigenvalues Ej >> 0, without loss of generality, as we can always shift the diagonal elements A to meet this condition:
tr (MUTAU) = =t (MUT (A (A+aT)
(23)
Let uty be the i, j-th entry of U. The Right-hand-sin (RHS) of the leftmost inequality can be expressed as,
k RHS = (MU AU) = Σ
m,a,ru;:4i
(24)
1=1
==
We denote the actual eigenvectors of A as V with entries {1u)). and {Ae)s 1...., as the eigenvalues, we have:
0gt
Aelis
(25)
Substituting this into (24), we obtain:
n
n AoUsauil
RHS
m;
us:
(26)
M& 7 Σ Vjoulis
Σ 1m1
Vis ui
(27)
j=1
mzd9 (v) vi uu
(28)
2m 1=1
17

---

We define
Isi = vJ ui
(29)
then the RHS becomes,
RHS- =
my maDa
(30)
-=1 s=1
and,
0 < 221 1
(31)
It's worth noting that for any
Σ Jss Σ u u.u!u ui
(32)
s=1
ir(u u VVVVV: ui = 1
tr
(33)
Thus we have,
n
n Σ
RHS-
Σ m; Σ
Aages
gos
(34)
Applying the Sauchy-Schwarte inequality, we get,
RHS
m,
Σ Aogaa
(35)
1=1
The equality holds if and only if for all s and a fixed 4,
Agga
const.
(36)
or
g
(37)
This is only true when gss are either 0 or 1, with at most one non-zero value for each s and i. In other words, G = {g=+] must be a permutation matrix. This implies that U is equivalent to the eigenvectors of V, albeit with possible permutation.
Thus this inequality reaches a minimum if we rearrange the indices of the eigenvalues:
RHS >
m
(38)
The 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-Schward inequality to Eq. (30) and obtain,
LHS <
m
gx
(39)
1=1
s=1
1=1
mi
Xg
(40)
1=1
s=1
mrdo(s)-
(41)
18

---

Here again, the Ssuchy-Schwart, inequality achieves the maximum when U is equivalent to the eigenvectors of V, up to possible permutation. By applying the rearrangement inequality, we further obtain the final bound:
LHS
m,71-
(42)
This yields the desired upper bound for the second inequality of the main result.
B.3 Proof of Theorem 2
Theorem 2 (Trace inequality with convex function) Let A be given n-dimensional symmetric matrix and U be a matrix of size n. x that resides on the Stiefel manifold. Suppose f :I -+ be a monotonically increasing, convex function applied element-wise The following inequalities hold:
ls (11 (H(AAA)) < (trr)((((rrr) <t(((UTAUU)
ESOO.
(43)
The rightmost inequality becomes equality if and only U comprises the eigenvectors of A that correspond to the largest eigenvalues, The leftmost inequality is met with equality when U is such that all diagonal elements of the matrix U 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 consider the left-hand side (LHS) of the equation. Let's denote the diagonal elements of A as {ajz]- Consider the eigendecomposition of A as in Eq. (25), can rewrite ajj in terms of the eigenvalues and eigenvectors:
0gs
Σ
(44)
where vy is the () entry of its cigenvectors V. The diagonal elements of the matrix product UT AU can then be written as:
(UTAU)
(45)
1=1
Substituting this result back into the LHS, we have
RHS = Σ ((TTU)) AU)
.0)
(46)
s=1
(47)
Σ
(48)
U
(49)
We can express the summation of the elements in vector form as follows:
Σ u. V.uu u.
(50)
1=1
tr VVT u.) =1
(51)
19

---

For all indices j. I and s, all the addends satisfy:
n n
Uy 1
(52)
which allows us treat Σ وー= to the equation previously labeled as Eq. (49), obtain:
Mij as a set of weights. By applying Jensen's inequality
k
R
R
n
LHS Σ (AA)
ΣΣΣ =-
(53)
Equality holds if and only if all the cigenvalues {Ai} are identical, which is generally not the case for an arbitrary symmetric matrix A, or for each index 8, there exists an index such that,
(u. v2,,v)==11 =1,
(54)
with all other terms being zero. Therefore, U must be permutation of V:
U (iii)
(55)
Substituting this result back into Eq. (53), we obtain:
EINEE i=
Es (No(s) & (Aa)
(56)
With this, we conclude the derivation of the first inequality. The leftmost inequality. For the leftmost inequality,
& RHS Σ
(57)
i=1
j=1
Uft
(58)
> kl
k
a agmng uef
(Jensen)
(59)
ter (UT (UTAU)
kk
(60)
1 ks (r((A))
(61)
the equality holds if and only if E, uiy Uij equals for all i.
similar result for finding the & smallest eigenvalues can be obtained. The next corollary states the results
Corollary (Trace inequality with concave function) Let A be a given n-dimensional symmetric matrix and let U be matrix of size x k that resides the Stiefel manifold. Suppose f I be monotonically increasing, concave function applied element-wiss The following inequalities hold:
ΣSO) & E
(I(TTTA & <k<(1 <
(62)
The rightmost inequality becomes an equality and only U comprises the eigenvectors of A that correspond to the l smallest eigenvalues.
20

---

Appendix C Further Extensions
C.1
Related Work
Eigensolver in numerical linear algebra Large-scale cigensolvers in numerical linear algebra usually involve two stages [39]. In the first stage, depending on whether the matrix symmetric or not, either Amoldi or Lanczos iteration is applied to reduce it into a matrix of smaller scale. Then, classical iteration methods such as power, inverse, and Rayleigh quotient iterations are applied to extract the eigensystem of the small-scale matrix. Compared to iterative solvers, our method is more easy to be integrated into the network training procedure.
Figendeccopoostton via manifold optimization In the optimization community, especially man- ifold optimization [1], identifying the sigen-subspace or solving for eigendecomposition often achieved by optimizing the original trace loss function in Eq. 6 or the Brockett's 's cost function in Eq. respectively. Instead of parametrizing the matrix U via QR decomposition, one usually solves the constraint optimization via gradient descent with retraction. Namely, in each iteration, one firstly performs one step of gradient descent from on the Stiefel manifold and then uses a retraction to correct it back to the manifold. Some convergence guarantees to the critical point and linear convergence rate can be established under certain assumptions [1, 43]. Various methods exist, such as [24] introducing the splitting method and [43] proposing Trank-Nicolson-- like update scheme to preserve the orthogonal constraints. No matter which method is used, the computation bottleneck lies in preserving the orthogonal constraints similar
amortized method.
Spectral normalization Spectral normalization is frequently added to improve the generalization ability of the neural network, i.e. CNN [11, 19] and GAN [32]. Compared nuclear-now norm regular- ization, spectral regularization only penalizes the largest absolute value of the eigenvalue so that the network outputs depend continuously and the spectral norm can be more efficiently estimated based on the power iteration on a randomly initialized vector in [32, 11].
the input. Therefore no orthogonal constraint is enforced
C.2 Extensions the Amortized Eigendecomposition Approach
Singular Value Decompusition Problem The singular value decomposition of a rectangular matrix A E pmxn can be achieved by solving the following objective:
max UeUnxk
||AUM||
(63)
where ||. ||r represent the Frobenius norm. The i-th largest singular values can be obtained by 0, Auil2 with u, the i-th column of optimal U.
Generalized Eigendecomposition Problem Generalized eigendecomposition is an extension of conventional eigendecomposition to pair of matrices, rather than just single matrix. The generalized eigendecompositiinn can be formulated as follows:
Au Bu
(64)
where A and B are square matrices of same dimensions in the complex field Cuxn The vector represents the eigenvector, and the scalar denotes the corresponding eigenvalue. Together, (s, u) constitute the eigenpair of the generalized eigenvalue problem of the matrix pair (A, B). In manner akin to prior problems, the resolution to this generalized version can be achieved through the optimization of the following objective function:
u Au;
max UEU*XA
mi
(65)
Bu:
where miss are the weights as defined previously and u, 's are the orthonormal column vectors of U.
21

---

NeurIPS Paper Checklist
1. Claims
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. 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 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 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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
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 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 (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn 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 the scope of the claims made, e.g. if the approach was only tested on few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
22

---

Answer:
[Yes]
Justification: All the assumptions are stated in Thm. and Thm. 2 respectively and the proof is put in Sec. B.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
This is included in this paper in Sec.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
23

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
Our code will be released upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //niss. wttpcccppppppeeeeeeeccecccccceet tww///uiisee//ee Setees//ooeSuuiieeeeeee ssissssoiiiyyyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary 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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. 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, confi- dence intervals, statistical significance tests, 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 with given experimental conditions).

---

The method for calculating the error bars should be explained (closed form formula, call 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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 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 failed experiments that didn't I make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics pettps:/nneuripet Setublic/Ethcssuiddee...
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations in their jurisdiction).
10. 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 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 why the paper does not address societal impact.
25

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
Answer: [No]
Justification: Not relevant to our research.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification:
We have done this. CC-BY 4.0
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include 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 use in package should be provided. For popular datasets, paperswi twwode..ooo..... com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
26

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We have done this.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
27


Reviews:
[{"review_id": "sGxTsaTMSB", "reviewer": "Reviewer_6W8h", "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. \nTo do so the author's add an auxiliary (and differentiable) loss whose optima is obtained by the appropriate set of eigenvectors.\nA summary of the presented results:\n- The author's provide both theoretical (proofs) and empirical evidence that the proposed estimation algorithm converges to the true eigenvectors.\n- 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.", "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).\n\n- 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": "Yes.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Motivation: The paper benefits from a very clear goal: to reduce the computational load of neural network training schemes that require eigendecomposition in each iteration. While this may be a somewhat niche use case, I believe such approaches are becoming increasingly prevalent. For example besides the three situations considered in the paper multiple recent self-supervised learning methods for learning image representations could benefit from the proposed speedups [1, 2]. \n\n- Novelty: The idea of learning estimates of eigenvectors using loss functions from existing optimization methods for eigendecomposition in neural network training is novel to the best of my knowledge. The author's discuss existing solutions in the appendix C. (iterative solvers, manifold optimization, etc.) and it seems clear to me the proposed method fills a hole in the practitioner's toolbox. I think this section could potentially be promoted to the main text if space permitted. \n\n- Clarity of presentation: The theorems and their proofs, the description of the algorithm, and the experimental design are all clearly presented and easy to understand. \n\n- Strong empirical results: In all of the considered settings the proposed amortization method provides significant speedups relative to full eigendecomposition (or SVD) during each iteration. I would be curious to know how such benefits scale as larger scale networks are employed (see below). Additional results in the appendix also suggest the amortization method results in similar solutions to the traditional approach. \n\n\n[1] Ermolov, Aleksandr, et al. \"Whitening for self-supervised representation learning.\" International conference on machine learning. PMLR, 2021.\n\n[2] Yerxa, Thomas, et al. \"Learning efficient coding of natural images with maximum manifold capacity representations.\" Advances in Neural Information Processing Systems 36 (2023): 24103-24128.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer 6W8h for the constructive comments and valuable suggestions. Here are our responses.\n\n***Q1 Regarding the scalability.***\n\nFor additional experiments on large-scale setups, please refer to the general response provided to all reviewers.\n\nWe 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.\n\n***Q2 Regarding Eq. 19.***\n\nThank 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. \n\nIf 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.\n\nWe 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:\n\n$$ tr(\\mathbf{M}\\mathbf{U} \\text{StopGrad} (\\text{cov} (h_\\theta(X)) \\mathbf{U} ) $$\n\nIn 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E3xme0aDse", "reviewer": "Reviewer_5kBi", "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.", "questions": "How does amortized eigendecomposition scale to larger dimensions?\nHow 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": "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.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Excellent organization and presentation, reads like a textbook chapter. \nThe method is well motivated and clearly explained. The analyses are on-point and thorough. \nThe significance and applications of the method is made clear, and the limitations are adequately discussed.\nThe method addresses a concrete application and is of immediate benefit for a class of problems in ML.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer 5kBi for his acknowledgement and valuable comments of our paper. Here are our responses.\n\n***Q1 Regarding the scalability.***\n\nFor additional experiments on large-scale setups, please refer to the general response provided to all reviewers.\n\n***Q2 Comparison with randomized SVD.***\n\nWhile 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:\n- *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.\n- *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.\n- *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.\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jmAji50i1u", "reviewer": "Reviewer_WiHy", "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.", "questions": "I have no questions.", "limitations": "The author clearly state the limitation of this paper.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The overall method is easy to understand and soundness. Combining the designed eigen loss with normal training loss is an interesting and reasonable method to effectively speed up the training process.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer WiHy for the constructive comments on our paper. Here are our responses.\n\n***Q1 Whether the accelerate method can maintain the performance.***\n\nIn the manuscript, we try to answer this question through two experiments:\n- 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.\n- 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.\n\nBoth 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W2ss5dxkbu", "reviewer": "Reviewer_fsNU", "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.\n\nThe 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)$.\n\nFor 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.", "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?\n2. Can this framework be applied to other aspects such as improving the stability of SVD in network training?", "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.\n\n2. 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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality: Good. The idea of applying regularizers to simulate SVD/eigendecomposition has been proposed in previous works such as https://arxiv.org/abs/2303.10512. However, the framework proposed here is novel and detailed.\n\nQuality: Good. The authors combined theoretical analysis (such as Fig.2) with empirical results (Fig.3).\n\nClarity: Good. The paper is well-written and easy to follow.\n\nSignificance: Fair. The experiments showed a substantial speedup for multiple tasks involving eigendecomposition.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate reviewer fsNU for the constructive feedback on our paper. Here are our responses.\n\n***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?***\n\nOur 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.\n\nAdditionally, 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.\n\n***Q2 Can this framework be applied to other aspects such as improving the stability of SVD in network training?***\n\nYes, 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. \n\nFurthermore, 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.\n\n***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.***\n\nWe 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.\n\n***Q4 Limited to eigendecomposition which is only applicable to symmetric matrices.***\n\nAlthough 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.\n\n[1]. Demmel, James W. Applied numerical linear algebra. Society for Industrial and Applied Mathematics, 1997.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all reviewers for their in-depth comments and valuable suggestions, which have significantly improved the quality of our paper. As many of the reviewers mentioned the scalability of our approach, we would like to provide a general response on this aspect.

We conducted an additional study on the scalability, with the results presented in the additional PDF page. Using the Celeb-A-HQ (256x256) dataset, we examined the scaling of the latent PCA task by varying the depth and width of the backbone autoencoder. The average execution time per iteration is reported. Notably, the largest model tested, with an autoencoder of 64 layers and a dimension of 4096, comprises over 1.0 billion parameters.

From the results, we can draw two main conclusions:

- **Efficiency of amortized eigen loss**: Our amortized eigen loss does not significantly increase the computational cost of the backbone but greatly reduces the training time for eigendecomposition. This is evident from the close alignment of the red (backbone) and green (backbone + our approach) lines. Compared to the traditional eigendecomposition approach (shown in the blue line), which increases dramatically as the dimension increases, our approach scales much more slowly.

- **Bottleneck of the latent PCA**: The primary bottleneck in such neural network structures is the eigendecomposition, while the computation for fully-connected layers is comparatively minor, especially when the width is large (>2000). This is demonstrated by the increasing gap between the backbone (red line) and backbone + eigh (blue line) as the dimension increases. However, if we increase the depth of the backbone while keeping the hidden dimension fixed, the execution time remains relatively unchanged. This indicates that the cost of fully-connected layers is small compared to that of eigendecomposition, echoing the results shown in Figure 1.

It is important to note that the total execution time includes both the computation of the eigen solver and the backbone autoencoders. The speed-up ratios presented in Table 1 reflect the speed-up of the eigen solvers alone, excluding the computation time of the backbone autoencoders.

Meta-review: 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.
Decision: Accept (poster)