[论文解读] Attention Scheme Inspired Softmax Regression

Yichuan Deng, Zhihang Li+1 more|arXiv (Cornell University)|Apr 20, 2023

Machine Learning and Algorithms被引用 6

一句话总结

本文提出了一种受大语言模型中注意力机制启发的新型Softmax回归框架，引入了一种具有可证明收敛保证的贪心优化算法。在温和条件下建立了线性收敛性，为使用迭代方法在实践中训练基于Softmax的模型提供了理论支持。

ABSTRACT

Large language models (LLMs) have made transformed changes for human society. One of the key computation in LLMs is the softmax unit. This operation is important in LLMs because it allows the model to generate a distribution over possible next words or phrases, given a sequence of input words. This distribution is then used to select the most likely next word or phrase, based on the probabilities assigned by the model. The softmax unit plays a crucial role in training LLMs, as it allows the model to learn from the data by adjusting the weights and biases of the neural network. In the area of convex optimization such as using central path method to solve linear programming. The softmax function has been used a crucial tool for controlling the progress and stability of potential function [Cohen, Lee and Song STOC 2019, Brand SODA 2020]. In this work, inspired the softmax unit, we define a softmax regression problem. Formally speaking, given a matrix $A \in \mathbb{R}^{n imes d}$ and a vector $b \in \mathbb{R}^n$, the goal is to use greedy type algorithm to solve \begin{align*} \min_{x} \| \langle \exp(Ax), {\bf 1}_n angle^{-1} \exp(Ax) - b \|_2^2. \end{align*} In certain sense, our provable convergence result provides theoretical support for why we can use greedy algorithm to train softmax function in practice.

研究动机与目标

• 形式化一个新颖的Softmax回归问题，通过指数函数实现归一化，受大语言模型中注意力机制的启发。
• 为在实践中使用贪心算法训练Softmax函数提供理论依据，特别是在大语言模型的背景下。
• 在矩阵范数和正则化的真实假设下，分析迭代算法在Softmax回归目标上的收敛行为。
• 建立一种每步具有多项式时间复杂度的可证明收敛算法，适用于大规模应用。
• 通过引入Softmax操作中固有的归一化因子，扩展先前关于双曲回归的工作。

提出的方法

• 将Softmax回归问题形式化为最小化 $Ax$ 的归一化指数与目标向量 $b$ 之间的平方L2距离，目标函数中包含 $\|\operatorname{diag}(w)Ax\|_2^2$ 的正则项。
• 提出一种贪心迭代算法（算法1），通过子采样对角矩阵近似Hessian矩阵，以在每一步更新解。
• 采用子采样技术（SubSample）降低每轮迭代的计算成本，同时保持收敛保证。
• 使用类似线搜索的更新规则，其中近似Hessian矩阵 $\widetilde{H} = A^T \widetilde{D} A$，$\widetilde{D}$ 为当前梯度对角矩阵的子采样子集。
• 通过在误差 $\|x_t - x^*\|_2$ 上进行归纳，证明收敛性，表明在特定条件下每轮迭代的收缩因子为0.4。
• 依赖于Hessian矩阵的Lipschitz连续性以及归一化项 $\langle \exp(Ax), \mathbf{1}_n \rangle$ 的下界，以控制误差传播。

实验结果

研究问题

• RQ1 在大语言模型的注意力机制中，Softmax回归目标具有核心作用，能否从理论上证明贪心算法的合理性？
• RQ2 数据矩阵 $A$、目标向量 $b$ 和正则化权重 $w$ 需满足何种条件，可确保迭代算法的线性收敛？
• RQ3 在高维设置下，如何高效计算Hessian近似，同时保持收敛保证？
• RQ4 收敛速率对 $A$ 的谱范数、最小奇异值 $\sigma_{\min}(A)$ 和正则化强度的依赖关系如何？
• RQ5 归一化因子 $\langle \exp(Ax), \mathbf{1}_n \rangle^{-1}$ 如何影响优化过程的稳定性和收敛性？

主要发现

• 在给定假设下，所提算法每轮迭代的收缩因子为0.4，实现线性收敛。
• 达到 $\epsilon$-精度所需的迭代次数为 $T = \log(\|x_0 - x^*\|_2 / \epsilon)$，确保对数收敛速率。
• 每轮迭代的时间复杂度为 $O((\operatorname{nnz}(A) + d^\omega) \cdot \operatorname{poly}(\log(n/\delta)))$，其中 $\omega \approx 2.373$ 为矩阵乘法指数。
• 在给定条件下，损失函数的Hessian矩阵被证明为正定且Lipschitz连续。
• 在归一化项 $\langle \exp(Ax), \mathbf{1}_n \rangle$ 上建立了 $\exp(-R^2)$ 的下界，这对误差控制至关重要。
• 该算法在 $T$ 轮迭代中以高概率 $1 - \delta$ 成功，利用联合界和失败概率可控的子采样技术。