Paper 2505.17967
FFT-Based Dynamic Subspace Selection for Low-Rank Adaptive Optimization of Large Language Models
- Published
- May 2025
- Research lab
- Together AI
- Citations
- 2
- GitHub
- 14 stars
01 In brief
Summary
The paper introduces a computationally efficient method for low-rank adaptive optimization of large language models (LLMs) by replacing SVD/QR-based gradient projections with a dynamic column selection from a fixed Discrete Cosine Transform (DCT) matrix.
The approach computes alignments between gradients and DCT columns, selects the top-r columns per layer, and uses them as projection matrices.
This reduces computational complexity to O(n^2 log n) via Makhoul's FFT-based algorithm and avoids storing per-layer projection matrices, saving memory.
The authors propose two optimizers: Trion, which improves Dion by replacing Power-Iteration with DCT-based projection and Newton-Schulz orthogonalization, and DCT-AdamW, which replaces SVD in low-rank AdamW variants with optional quantized error feedback.
Experiments on pretraining Llama models (350M-1.3B) and fine-tuning Llama-2-7B and Qwen-2.5-7B show that Trion and DCT-AdamW match or outperform baselines in perplexity/accuracy while reducing memory usage by up to 25% and runtime by up to 18% (Trion) or 25.75% (DCT-AdamW vs LDAdamW).
Theoretical analysis proves the optimality of norm-based column selection for minimizing reconstruction error and justifies DCT as a linear approximation of gradient eigenbases.
Limitations include testing only up to 1.3B parameters and bfloat16 constraints for Makhoul's algorithm.
02 From the paper
Abstract
Low-rank optimization has emerged as a promising direction in training large language models (LLMs) to improve running time and reduce the memory usage of adaptive optimizers by constraining learning to a lower-dimensional space. Prior work typically projects gradients of linear layers using approaches based on Singular Value Decomposition (SVD) or QR-decomposition. Applying these techniques individually to each layer in large models is computationally expensive and incurs additional memory costs due to storing the projection matrices. In this work, we propose a computationally efficient and conceptually simple, two-step procedure to approximate SVD/QR-based gradient projections into lower-dimensional spaces by using a predefined orthogonal matrix of the Discrete Cosine Transform (DCT). We dynamically select columns from the DCT matrix based on their alignment with the gradient of each layer. The effective projection matrices are obtained via a simple matmul with the DCT matrix in $O(n^3)$ time, followed by a lightweight sorting step to identify the most relevant basis vectors. For large layers, DCT can be computed via Makhoul's $N$-point algorithm based on Fast Fourier Transform (FFT) in $O(n^2 \log(n))$ time. Due to the predefined nature of the orthogonal bases, they are computed once at the start of training. Our numerical experiments on both pre-training and fine-tuning tasks demonstrate the effectiveness of our dual strategy in approximating optimal low-rank projections, obtaining an approach with rank-independent running time that matches the performance of costly SVD/QR-based methods while achieving faster runtime and reduced memory usage by up to $25\%$ across different model sizes. Our code is available at \href{https://github.com/IST-DASLab/ISTA-DASLab-Optimizers}{\texttt{https://github.com/IST-DASLab/ISTA-DASLab-Optimizers}}.