[论文解读] Zero-th Order Algorithm for Softmax Attention Optimization

Yichuan Deng, Zhihang Li+2 more|arXiv (Cornell University)|Jul 17, 2023

Topic ModelingComputer Science被引用 3

一句话总结

该论文提出了一种专为大语言模型（LLMs）中的softmax注意力机制设计的零阶优化算法，仅通过前向传播实现梯度估计。其建立了全局收敛性，样本复杂度为 $ O(M \cdot (1 + d^{1.5} \kappa^2(A)/k) \cdot \mu^{-2} B^{-1} \log((L(x_0)-L^*)/\epsilon)) $，其中 $ M = \exp(O(R^2 + \log n)) $，在高维设置下实现了无需反向传播的高效训练。

ABSTRACT

Large language models (LLMs) have brought about significant transformations in human society. Among the crucial computations in LLMs, the softmax unit holds great importance. Its helps the model generating a probability distribution on potential subsequent words or phrases, considering a series of input words. By utilizing this distribution, the model selects the most probable next word or phrase, based on the assigned probabilities. The softmax unit assumes a vital function in LLM training as it facilitates learning from data through the adjustment of neural network weights and biases. With the development of the size of LLMs, computing the gradient becomes expensive. However, Zero-th Order method can approximately compute the gradient with only forward passes. In this paper, we present a Zero-th Order algorithm specifically tailored for Softmax optimization. We demonstrate the convergence of our algorithm, highlighting its effectiveness in efficiently computing gradients for large-scale LLMs. By leveraging the Zeroth-Order method, our work contributes to the advancement of optimization techniques in the context of complex language models.

研究动机与目标

• 解决因参数量巨大而导致的大规模LLM中反向传播计算成本过高的问题。
• 在不依赖反向传播计算梯度的前提下，实现对softmax注意力机制的高效优化。
• 开发一种零阶方法，仅通过前向传播评估近似梯度，从而降低内存和计算开销。
• 在模型参数和损失曲面的合理假设下，为所提算法建立理论收敛保证。
• 提供一种适用于大规模LLM的可扩展优化框架，其中传统反向传播不可行。

提出的方法

• 基于同时扰动随机逼近（SPSA）提出一种仅依赖前向损失评估的零阶梯度估计器。
• 通过扰动与未扰动参数间损失差值推导无偏梯度估计：$ \widehat{g}(x_0)_i = \frac{1}{2\epsilon p_i}(L(x_0 + \epsilon p) - L(x_0 - \epsilon p)) $。
• 将梯度估计应用于投影梯度下降（GD）框架，以优化softmax注意力目标。
• 定义损失函数 $ L(x) $ 为归一化指数输出与目标向量之间的平方误差，将注意力建模为softmax回归问题。
• 利用集中不等式和矩阵范数界控制梯度估计中的方差，确保在光滑性和PL条件下收敛。
• 通过包含每次迭代损失下降量的递归期望界建立收敛性，利用 $ \mu $-PL条件和Hessian协方差的迹。

实验结果

研究问题

• RQ1 零阶优化方法能否在大语言模型的softmax注意力中实现全局收敛？
• RQ2 在模型参数的合理假设下，注意力机制的零阶优化样本复杂度是多少？
• RQ3 与传统反向传播相比，该方法在计算效率和收敛速度方面表现如何？
• RQ4 在高维、非凸的注意力优化问题中，梯度估计可提供哪些理论保证？
• RQ5 模型参数如条件数 $ \kappa(A) $、维度 $ d $ 和批量大小 $ B $ 如何影响收敛性？

主要发现

• 所提出的零阶算法在 $ O(M \cdot (1 + d^{1.5} \kappa^2(A)/k) \cdot \mu^{-2} B^{-1} \log((L(x_0)-L^*)/\epsilon)) $ 次迭代内实现全局收敛至 $ \epsilon $-最优解，其中 $ M = \exp(O(R^2 + \log n)) $。
• 该方法在 $ \mu $-PL、$ l $-光滑性和有界Hessian有效秩条件下确保收敛，方差通过SPSA扰动得到控制。
• 收敛速率取决于PL常数 $ \mu $、批量大小 $ B $ 和条件数 $ \kappa(A) $ 的倒数，且 $ B $ 越大，样本效率越高。
• 该算法完全避免反向传播，仅依赖前向传播评估，适用于黑箱或梯度不可见的大语言模型。
• 理论分析表明，梯度估计器的方差有界于 $ \operatorname{tr}[\Sigma(x)] \leq \alpha (L(x) - L^*) $，从而实现稳定下降。
• 最终收敛界与 $ \exp(O(R^2)) $ 成比例，其中 $ R $ 为参数范数上界，表明在最坏情况下对模型规模呈指数依赖。