Efficient Preconditioned Optimization Methods with Frequent Directions

The gist: The authors propose several efficient preconditioned optimization methods, including FTSL and S-ASGO, that accelerate adaptive optimization techniques by integrating frequent directions (FD) with various frameworks and exploiting matrix structures to achieve tighter regret bounds and reduced computational complexity.

How it works

1. The paper investigates the acceleration of preconditioned optimization methods through frequent directions (FD), a widely used matrix sketching technique, by proposing FTSL which integrates FD with adaptive subgradient methods under the primal-dual framework and adds back the cumulative information discarded by FD to obtain a tighter regret bound than existing state-of-the-art results.

2. To further reduce the time complexity of FTSL, they propose an accelerated variant named FTFSL which doubles the sketching size to lower the frequency of expensive SVD computations while retaining the same regret bound and space complexity.

3. For optimization over matrix-valued decision variables, they develop FTSL-Shampoo by adapting FD to Shampoo under the primal-dual framework, which attains a dimension-free regret bound of O(√r(tr(L1/4T) + (ρL:T)1/4)(tr(R1/4T) + (ρR:T)1/4)), while preserving the same space complexity.

4. To handle limitations of FTSL-Shampoo, they incorporate FD into ASGO and propose Sketchy ASGO (S-ASGO), which maintains a low-rank approximation of the preconditioner only along the smaller dimension to reduce space complexity from O(τ(m+n)) to O(τm).

5. The analysis is extended to offline stochastic convex optimization for both non-smooth and smooth objectives, achieving convergence rates that exploit smoothness for faster rates.

Key Algorithms and Variants

(Algorithm 2) Follow-the-Sketchy-Leader (FTSL)

This algorithm integrates FD with ADA-FULL under the primal-dual framework and adds back the cumulative discarded eigenvalues into the preconditioner. The regret of FTSL is upper bounded by O(tr(G1/2T) + √ρ1:T), avoiding dependence on d.

(Algorithm 3) Follow-the-Fast-Sketchy-Leader (FTFSL)

This variant employs a more efficient strategy for computing the SVD of Bt, representing it through two matrices Vt and Mt, which allows them to only perform an SVD decomposition on Mt every τ + 1 rounds, reducing the time complexity to O(τ 3) per round when τ ≤ O(√d).

(Algorithm 6) FTSL-Shampoo)

This method uses FD to approximate the left and right preconditioning matrices for Shampoo, maintaining a space complexity of O(τ(m+n)) while achieving a dimension-free regret bound.

(Algorithm 7) S-ASGO)

S-ASGO maintains only a single sketching matrix Lˆ t in R mxτ to approximate the matrix t i=1 GX i(GX i) T along the smaller axis, reducing space complexity from O(τ(m+n)) to O(τm).

Theoretical Guarantees

(Theorem 1) for FTSL

Under Assumption 1, by setting the learning rate η = √D squared, FTSL ensures R(T) ≤ O(tr(G1/2T) + √ρ1:T), where GT = Pt=1 gtg T.

(Theorem 7) for S-ASGO under the primal-dual framework

Under Assumption 1, by setting the learning rate η = D√F squared, S-ASGO under the primal-dual framework ensures R(T) ≤ O(DF tr(L1/2T) + qρL:T + ϵ), where DF = X∗F is the Frobenius norm of the optimal decision X.

(Theorem 8) for S-ASGO under the mirror descent framework

Under Assumption 1, by setting the learning rate η = √D 2/2, S-ASGO under the mirror descent framework ensures R(T) ≤ O(tr(L1/2T) + m qρL:T + ϵD 2F), where D2 = maxt∈[T] Xt - X∗ squared.

Empirical Results

Experiments on the Gisette and Epsilon datasets show that FTFSL outperforms all other methods in both loss and testing accuracy, aligning with its superior regret bound. Furthermore, FTFSL and Fast S-ADA exhibit significantly lower runtimes compared to S-ADA due to their superior time complexities. In language model pretraining experiments on GPT-2, S-ASGO (M) maintains a much smaller preconditioner than ASGO because it keeps only a low-rank preconditioning matrix rather than the full one, reducing optimizer memory while leaving the pretraining loss essentially unchanged.

Future Work

Two directions are particularly worth pursuing in future work: first, for sparse gradient matrices, the sparse frequent directions technique (Ghashami et al., 2016a) attains a better time complexity than standard FD. Second, our theoretical analysis is developed for convex loss functions in the offline stochastic optimization setting, whereas the loss functions of many practical machine learning applications are inherently non-convex.

The paper concludes by demonstrating that their methods attain performance comparable to existing approaches while consistently reducing computational cost through empirical validation across various tasks. The authors also provide a practical variant of S-ASGO tailored for large-scale model pretraining. The overall findings suggest that sketched preconditioners retain the benefit of axis-wise preconditioning at a substantially lower memory cost.

The paper is structured to provide a systematic investigation into computationally efficient preconditioned optimization methods via frequent directions (FD) and their extensions for both vector-valued and matrix-valued decision variables. The remainder of the paper is structured as follows: Section 2 briefly reviews related work, Section 3 presents the preliminaries, including notation, assumptions, as well as the technical tools that underpin our analysis, Section 4 illustrates our main results and their extension to offline stochastic optimization, Section 5 establishes results for optimization with matrix structure, Section 6 provides the analysis of all theorems and Appendix A provides the proof of Theorem 4. The authors also conduct extensive experiments to validate the effectiveness and efficiency of their methods on large language model pretraining tasks. The paper concludes by discussing future work.

The authors also provide a practical variant of S-ASGO tailored to large-scale model pretraining. The final results show that FTSL-Shampoo attains performance comparable to that of the full Shampoo while substantially reducing memory consumption and running time, in line with their theoretical guarantees. The work provides a comprehensive analysis for both non-smooth and smooth objectives in the offline stochastic setting. It is also noted that the sketching size of 0.2 can capture the principal components of the gradients in language modeling. The paper is a significant contribution to making adaptive optimization methods practical for large-scale machine learning tasks.

1. Bold Preconditioning for Large-Scale Neural Networks: The FTSL-Shampoo method allows for optimization over matrix-valued decision variables with a space complexity of O(τ (m + n)) while achieving a dimension-free regret bound, making it suitable for training very large models where the full matrix preconditioner is computationally prohibitive.

2. Accelerated Optimization Time Complexity: The FTFSL variant reduces the time complexity to O(τ d) when sketching size τ is sufficiently small, and this acceleration can be applied to S-ADA, leading to faster per-round computation during training on high-dimensional tasks.

3. Efficient Matrix Preconditioning for Structured Gradients: By exploiting the low-rank structure of gradients in deep neural networks, FTSL-Shampoo constructs preconditioning matrices by maintaining two low-rank matrices Lˆ t ∈ R m×τ and Rˆ t ∈ R n×τ to approximate the preconditioning matrices Lt and Rt, which keeps optimizer memory governed by the sketching size rather than the full preconditioner dimension.

4. Robust Convergence in Offline Stochastic Optimization: The extension of S-ASGO to offline stochastic optimization provides convergence rates for both non-smooth and smooth objectives, ensuring that the sketching error term scales linearly with d or achieving a dimension-free regret bound, which is critical for training on real-world datasets where only noisy gradient estimates are available.

5. Practical Implementation for LLM Pretraining: The practical variant of S-ASGO, Algorithm 8, incorporates an exponential moving average and momentum, enabling its use in large language model pretraining while maintaining performance comparable to state-of-the-art methods like Muon and Shampoo with a significantly smaller preconditioner memory footprint.

Adaptive optimization methods have received considerable attention in recent years due to their strong empirical performance. However, they generally require maintaining a preconditioning matrix, which incurs prohibitive computational overhead in large-scale applications. To address these limitations, we investigate the acceleration of preconditioned optimization methods through frequent directions (FD), a widely used matrix sketching technique. We first propose FTSL by integrating FD with adaptive subgradient methods under the primal-dual framework and adding back the cumulative information discarded by FD. FTSL obtains a tighter regret bound than existing state-of-the-art results. We then apply a doubling trick that further reduces the time complexity of FTSL, while maintaining the same regret guarantee. We also consider optimization over matrix-valued decision variables and exploit their structural properties to construct sketched preconditioners. We develop FTSL-Shampoo by adapting FD to Shampoo and give a novel analysis under the primal-dual framework. To further mitigate its computational complexity, we design one-sided variants of FTSL-Shampoo under both mirror descent and primal-dual frameworks. Our algorithms not only reduce the computational cost of FTSL-Shampoo, but also achieve tighter regret bounds. We extend our methods to offline stochastic convex optimization and exploit smoothness to achieve faster convergence rates. Finally, we introduce practical variants of our algorithms and empirically demonstrate their efficiency and effectiveness across a broad range of tasks.