arXiv:2505.17967cs.LGcs.AI2025-05被引 2

用FFT加速低秩优化,训练大模型更快更省内存。

FFT-based Dynamic Subspace Selection for Low-Rank Adaptive Optimization of Large Language Models

  • 用DCT正交基替代SVD/QR,通过FFT快速计算梯度投影。
  • 动态选相关基向量,实现与高成本方法相当的性能。
  • 适用于大规模模型训练,速度提升最高达25%。

低秩优化通过将学习限制在低维空间,提升了大语言模型训练的运行效率并降低了自适应优化器的内存开销。现有方法通常使用奇异值分解(SVD)或QR分解对线性层梯度进行投影,但在大型模型中逐层应用代价高昂,且需存储投影矩阵。本文提出一种基于离散余弦变换(DCT)的预定义正交基两步法,通过快速傅里叶变换(FFT)以$O(n^2 \log n)$时间高效计算,动态选择与各层梯度对齐的列。投影仅需一次矩阵乘法和轻量排序,无需额外存储。实验表明,该方法在预训练与微调任务中均能逼近最优低秩投影性能,实现与模型规模无关的运行时间,比传统SVD/QR方法快25%,内存消耗显著降低。

原文摘要 · Abstract (English)

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}}.

低秩优化FFT大模型训练DCT

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。