基于FFT的动态子空间选择用于大型语言模型的低秩自适应优化

于红博

摘要: 低秩优化已经成为训练大型语言模型（LLMs）的一个有前途的方向，可以通过将学习限制在低维空间中来改善运行时间和减少自适应优化器的内存使用。先前的工作通常使用基于奇异值分解（SVD）或QR分解的方法来投影线性层的梯度。将这些技术分别应用于大型模型中的每一层是计算昂贵的，并且由于存储投影矩阵而产生额外的内存成本。在这项工作中，我们提出了一个计算效率高、概念简单的两步过程，通过使用离散余弦变换（DCT）的预定义正交矩阵，将基于SVD/QR的梯度投影近似到低维空间中。我们根据每一层梯度与DCT矩阵的对齐性动态选择DCT矩阵中的列。通过在$O(n^3)$时间内使用DCT矩阵简单地进行矩阵相乘，然后通过轻量级的排序步骤来识别最相关的基础向量，得到有效的投影矩阵。对于大型层，可以通过基于快速傅里叶变换（FFT）的Makhoul的$N$点算法在$O(n^2 \log(n))$时间内计算DCT。由于正交基的预定义性质，它们在训练开始时只需计算一次。我们在预训练和微调任务上的数值实验表明，我们的双重策略在近似最佳低秩投影方面的有效性，获得了一个运行时间与秩无关的方法，可与昂贵的SVD/QR方法的性能相匹配，同时在不同模型大小上实现了更快的运行时间和减少了高达25%的内存使用。我们的代码可在\href{https://github.com/IST-DASLab/ISTA-DASLab-Optimizers}{\texttt{https://github.com/IST-DASLab/ISTA-DASLab-Optimizers}}找到。

更新时间: 2025-10-08 15:33:25

领域: cs.LG,cs.AI

下载: http://arxiv.org/abs/2505.17967v4