EigenGame Unloaded: When playing games is better than optimizing

1 INTRODUCTION . Large , high-dimensional datasets containing billions of samples are commonplace . Dimensionality reduction to extract the most informative features is an important step in the data processing pipeline which enables faster learning of classifiers and regressors ( Dhillon et al. , 2013 ) , clustering ( Kannan and Vempala , 2009 ) , and interpretable visualizations . Many dimensionality reduction and clustering techniques rely on eigendecomposition at their core including principal component analysis ( Jolliffe , 2002 ) , locally linear embedding ( Roweis and Saul , 2000 ) , multidimensional scaling ( Mead , 1992 ) , Isomap ( Tenenbaum et al. , 2000 ) , and graph spectral clustering ( Von Luxburg , 2007 ) . Numerical solutions to the eigenvalue problem have been approached from a variety of angles for centuries : Jacobi ’ s method , Rayleigh quotient , power ( von Mises ) iteration ( Golub and Van der Vorst , 2000 ) . For large datasets that do not fit in memory , approaches that access only subsets—or minibatches—of the data at a time have been proposed . Recently , EigenGame ( Gemp et al. , 2021 ) was introduced with the novel perspective of viewing the set of eigenvectors as the Nash strategy of a suitably defined game . While this work demonstrated an algorithm that was empirically competitive given access to only subsets of the data , its performance degraded with smaller minibatch sizes , which are required to fit high dimensional data onto devices . One path towards circumventing EigenGame ’ s need for large minibatch sizes is parallelization . In a data parallel approach , updates are computed in parallel on partitions of the data and then combined such that the aggregate update is equivalent to a single large-batch update . The technical obstacle preventing such an approach for EigenGame lies in the bias of its updates , i.e. , the divide-and-conquer EigenGame update is not equivalent to the large-batch update . Biased updates are not just a theoretical nuisance ; they can slow and even prevent convergence to the solution ( made obvious in Figure 4 ) . In this work we introduce a formulation of EigenGame which admits unbiased updates which we term µ-EigenGame . We will refer to the original formulation of EigenGame as α-EigenGame.3 µ-EigenGame and α-EigenGame are contrasted in Figure 1 . Unbiased updates allow us to increase the effective batch size using data parallelism . Lower variance updates mean that µ-EigenGame should converge faster and to more accurate solutions than α-EigenGame regardless of batch size . In Figure 1a ( top ) , the density of the shaded region shows the distribution of steps taken by the stochastic variant of each algorithm after 100 burn-in steps . Although the expected path of α-EG is 3µ signifies unbiased or unloaded and α denotes original . 1The trajectory when updating with E [ X > t Xt ] . 2Overestimation is expected by Jensen ’ s : E [ 1 X ] ≥ 1E [ X ] . slightly more direct , its stochastic variant has much larger variance . Figure 1a ( bottom ) shows that with increasing iterations , the µ-EG trajectory approaches its expected value whereas α-EG exhibits larger bias . Figure 1b further supports µ-EigenGame ’ s reduced bias with details in Sections 3 and 4 . Our contributions : In the rest of the paper , we present our new formulation of EigenGame , analyze its bias and propose a novel unbiased parallel variant , µ-EigenGame with stochastic convergence guarantees . µ-EigenGame ’ s utilities are distinct from α-EigenGame and offer an alternative perspective . We demonstrate its performance with extensive experiments including dimensionality reduction of massive data sets and clustering a large social network graph . We conclude with discussions of the algorithm ’ s design and context within optimization , game theory , and neuroscience . 2 PRELIMINARIES AND RELATED WORK . In this work , we aim to compute the top-k right singular vectors of dataX , which is either represented as a matrix , X ∈ Rn×d , of n d-dimensional samples , or as a d-dimensional random variable . In either case , we assume we can repeatedly sample a minibatch Xt from the data of size n′ < n , Xt ∈ Rn ′×d . The top-k right singular vectors of the dataset are then given by the top-k eigenvectors of the ( sample ) covariance matrix , C = E [ 1n′X > t Xt ] = E [ Ct ] . For small datasets , SVD is appropriate . However , the time , O ( min { nd2 , n2d } ) , and space , O ( nd ) , complexity of SVD prohibit its use for larger datasets ( Shamir , 2015 ) including when X is a random variable . For larger datasets , stochastic , randomized , or sketching algorithms are better suited . Stochastic algorithms such as Oja ’ s algorithm ( Oja , 1982 ; Allen-Zhu and Li , 2017 ) perform power iteration ( Rutishauser , 1971 ) to iteratively improve an approximation , maintaining orthogonality of the eigenvectors typically through repeated QR decompositions . Alternatively , randomized algorithms ( Halko et al. , 2011 ; Sarlos , 2006 ; Cohen et al. , 2017 ) first compute a random projection of the data onto a ( k + p ) -subspace approximately containing the top-k subspace . This is done using techniques similar to Krylov subspace iteration methods ( Musco and Musco , 2015 ) . After projecting , a call to SVD is then made on this reduced-dimensionality data matrix . Sketching algorithms ( Feldman et al. , 2020 ) such as Frequent Directions ( Ghashami et al. , 2016 ) also target learning the top-k subspace by maintaining an overcomplete sketch matrix of size ( k + p ) × d and maintaining a span of the top subspace with repeated calls to SVD . In both the randomized and sketching approaches , a final SVD of the n× ( k + p ) dataset is required to recover the desired singular vectors . Although the SVD scales linearly in n , some datasets are too large to fit in memory ; in this case , an out-of-memory SVD may suffice ( Haidar et al. , 2017 ) . For this reason , the direct approach of stochastic algorithms , which avoid an SVD call altogether , is appealing when processing very large datasets . Algorithm 1 µ-EigenGameR 1 : Given : data stream Xt ∈ Rn ′×d , vectors v̂0i ∈ Sd−1 , step sequence ηt , and iterations T . 2 : v̂i ← v̂0i for all i 3 : for t = 1 : T do 4 : parfor i = 1 : k do 5 : rewards← 1n′X > t Xtv̂i 6 : penalties ← 1 n′ ∑ j < i〈Xtv̂i , Xtv̂j〉v̂j 7 : ∇̃µi ← rewards− penalties 8 : ∇̃µ , Ri ← ∇̃ µ i − 〈∇̃ µ i , v̂i〉v̂i 9 : v̂′i ← v̂i + ηt∇̃ µ , R i 10 : v̂i ← v̂ ′ i ||v̂′i|| 11 : end parfor 12 : end for 13 : return all v̂i A large literature on distributed approaches to PCA exists ( Liang et al. , 2014 ; Garber et al. , 2017 ; Fan et al. , 2019 ) . These typically follow the pattern of computing solutions locally and then aggregating them in a single round ( or minimal rounds ) of communication . The modern distributed machine learning setting which has evolved to meet the needs of deep learning is fundamentally different . Many accelerators joined with fast interconnects means the cost of communication is low compared to the cost of a single update step , however existing approaches to distributed PCA can not take full advantage of this . Notation : We follow the same notation as Gemp et al . ( 2021 ) . Variables returned by an approximation algorithm are distinguished from the true solutions with hats , e.g. , the column-wise matrix of eigenvectors V̂ approximates V . We order the columns of V such that the ith column , vi , is the eigenvector with the ith largest eigenvalue λi . The set of all eigenvectors { vj } with λj larger than λi , namely vi ’ s parents , will be denoted by vj < i . Similarly , sums over subsets of indices may be abbreviated as ∑ j < i = ∑i−1 j=1 . The set of all parents and children of vi are denoted by v−i . Let the ith eigengap gi = λi − λi+1 . We assume the standard Euclidean inner product 〈u , v〉 = u > v and denote the unit-sphere and simplex in ambient space Rd with Sd−1 and ∆d−1 respectively . α-EigenGame . We build on the algorithm introduced by Gemp et al . ( 2021 ) , which we refer to here as α-EigenGame . This algorithm is derived by formulating the eigendecomposition of a symmetric positive definite matrix as the Nash equilibrium of a game among k players , each player i owning the approximate eigenvector v̂i ∈ Sd−1 . Each player is also assigned a utility function , uαi ( v̂i|v̂j < i ) , that they must maximize : uαi ( v̂i|v̂j < i ) = Var︷ ︸︸ ︷ v̂ > i Cv̂i− ∑ j < i Align-penalty︷ ︸︸ ︷ 〈v̂i , Cv̂j〉2 〈v̂j , Cv̂j〉 . ( 1 ) These utilities balance two terms , one that rewards a v̂i that captures more variance in the data and a second term that penalizes v̂i for failing to be orthogonal to each of its parents v̂j < i ( these terms are indicated with Var and Align-penalty in equation ( 1 ) ) . In α-EigenGame , each player simultaneously updates v̂i with gradient ascent , and it is shown that this process converges to the Nash equilibrium . We are interested in extending this approach to the data parallel setting where each player i may distribute its update computation over multiple devices . 3 A SCALABLE UNBIASED ALGORITHM . We present our novel modification to α-EigenGame called µ-EigenGame along with intuition , theory , and empirical support for critical lemmas . We begin with identifying and systematically removing the bias that exists in the α-EigenGame updates . We then explain how removing bias allows us to exploit modern compute architectures culminating in the development of a highly parallelizable algorithm . 3.1 α-EIGENGAME ’ S BIASED UPDATES Consider partitioning the sample covariance matrix Ct into a sum ofmmatrices as Ct = 1n′X > t Xt = 1 m ∑ m m n′X > tmXtm = 1 m ∑ m Ctm . For sake of exposition , we drop the additional subscript t on C in what follows . We would like α-EigenGame to parallelize over these partitions . However , the gradient of uαi with respect to v̂i does not decompose cleanly over the data partitions : ∇αi ∝ Var︷︸︸︷ Cv̂i − ∑ j < i Align-penalty︷ ︸︸ ︷ v̂ > i Cv̂j v > j Cv̂j Cv̂j = 1 m ∑ m [ Cmv̂i − ∑ j < i v̂ > i Cv̂j v̂ > j Cv̂j Cmv̂j ] . ( 2 ) We include the superscript α on the EigenGame gradient to differentiate it from the µ-EigenGame direction later . The nonlinear appearance of C in the penalty terms makes obtaining an unbiased gradient difficult . The quadratic term in the numerator of equation ( 2 ) could be made unbiased by using two sample estimates ofC , one for each term . But the appearance of the term in the denominator does not have an easy solution . Cm is likely singular for small n′ ( n′ < d ) which increases the likelihood of a small denominator , i.e. , a large penalty coefficient ( boxed ) , if we were to estimate the denominator with samples . The result is an update that emphasizes penalizing orthogonality over capturing data variance . Techniques exist to reduce the bias of samples of ratios of random variables , but to our knowledge , techniques to obtain unbiased estimates are not available . This was conjectured by Gemp et al . ( 2021 ) as the reason for why α-EigenGame performed worse with small minibatches . 3.2 REMOVING α-EIGENGAME ’ S BIAS It is helpful to rearrange equation ( 2 ) to shift perspective from estimating a penalty coefficient ( in red ) to estimating a penalty direction ( in blue ) : ∇αi ∝ 1 m ∑ m [ Cmv̂i − ∑ j < i v̂ > i Cmv̂j Cv̂j v̂ > j Cv̂j ] . ( 3 ) The penalty direction in equation ( 3 ) is still difficult to estimate . However , consider the case where v̂j is any eigenvector of C with associated ( unknown ) eigenvalue λ′ . In this case , Cv̂j = λ′v̂j and the penalty direction ( in blue ) simplifies to v̂j because ||v̂j || = 1 . While this assumption is certainly not met at initialization , α-EigenGame leads each v̂j towards vj , so we can expect this assumption to be met asymptotically . This intuition motivates the following µ-EigenGame update direction for v̂i with inexact parents v̂j ( compare orange in equation ( 4 ) to blue in equation ( 3 ) ) : ∆µi = Cv̂i − ∑ j < i ( v̂ > i Cv̂j ) v̂j = 1 m ∑ m [ Cmv̂i − ∑ j < i ( v̂ > i Cmv̂j ) v̂j ] . ( 4 ) We use ∆ instead of ∇ because the direction is not a gradient ( discussed later ) . Notice how the strictly linear appearance of C in µ-EigenGame allows the update to easily decompose over the data partitions in equation ( 4 ) . The µ-EigenGame update satisfies two important properties . Lemma 1 ( Asymptotic equivalence ) . The µ-EigenGame direction , ∆µi , with exact parents ( v̂j = vj ∀ j < i ) is equivalent to α-EigenGame . Proof . We start with α-EigenGame and add a superscript e to its gradient to emphasize this is the gradient computed with exact parents ( v̂j = vj ) . Then simplifying , we find ∇α , ei ∝ Cv̂i − ∑ j < i v̂ > i Cvj v > j Cvj Cvj = Cv̂i − ∑ j < i v̂ > i Cvj v > j λjvj λjvj = Cv̂i − ∑ j < i ( v̂ > i Cvj ) vj = ∆ µ i . ( 5 ) Therefore , once the first ( i − 1 ) eigenvectors are learned , learning the ith eigenvector with µEigenGame is equivalent to learning with α-EigenGame . Lemma 2 ( Zero bias ) . Unbiased estimates of ∆µi can be obtained with samples from p ( X ) . Proof . Let X ∼ p ( X ) where X ∈ Rd and p ( X ) is the uniform distribution over the dataset . Then E [ ∆µi ] = E [ XX > ] v̂i − ∑ j < i ( v̂ > i E [ XX > ] v̂j ) v̂j = Cv̂i − ∑ j < i ( v̂ > i Cv̂j ) v̂j . ( 6 ) where all expectations are with respect to p ( X ) . These two lemmas provide the foundation for a performant algorithm . The first enables convergence to the desired solution , while the second facilitates scaling to larger datasets . Algorithm 1 presents pseudocode for µ-EigenGame where computation is parallelized over the k players . 3.3 MODEL AND DATA PARALLELISM In our setting we have a number of connected devices . Specifically we consider the parallel framework specified by TPUv3 available in Google Cloud , however our setup is applicable to any multi-host , multi-device system . The α-EigenGame formulation ( Gemp et al. , 2021 ) considers an extreme form of model parallelism ( Figure 2a ) where each device has its own unique set of eigenvectors . In this work we further consider a different form of model and data parallelism which is directly enabled by having unbiased updates ( Figure 2b ) . This enables µ-EigenGame to deal with both high-dimensional problems as well as massive sample sizes . Here each set of eigenvectors is copied on M devices . Update directions are computed on each device individually using a different data stream and then combined by summing or averaging . Updates are applied to a sin- gle copy and this is duplicated across theM−1 remaining devices . In this way , updates are computed using an M× larger effective batch size while still allowing device-wise model parallelism . This setting is particularly useful when the number of samples is very large . This form of parallelism is not possible using the original EigenGame formulation since it relies on combining unbiased updates . In this sense , the parallelism discussed in this work generalizes that introduced by Gemp et al . ( 2021 ) . Note that we also allow for within-device parallelism . That is , each vi in Figure 2 is a contiguous collection of eigenvectors which are updated independently , in parallel , on a given device ( for example using vmap in Jax ) . We provide pseudocode in Algorithm 2 in the appendix which simply augments Algorithm 1 with an additional parallelized for-loop and aggregation step over available devices . We also provide detailed Jax pseudo-code for parallel µ-EigenGame in Appendix F. We compare the empirical scaling performance of µ-EigenGame against α-EigenGame on a 14 billion sample dataset in section 5 .

Summary: The authors study the problem of finding the top $k$ right singular vectors of a data matrix $X$. They propose a modification of an established game theoretic gradient based algorithm $\alpha$-EigenGame. They observe that the gradients of $\alpha$-EigenGame are biased when implemented stochastically by subsampling the data matrix $X$. Their proposed modification guarantees unbiased updates while converging to the true singular vectors. Their modified updates in combination with their data parallel distrusted algorithm leads to significantly improved convergence rates.