受注意力机制启发的 Softmax 回归

大语言模型（LLMs）给人类社会带来了变革性改变。LLMs 中的关键计算之一是 softmax 单元。该操作在 LLMs 中十分重要，因为它使模型能在给定输入词序列时，对可能的下一个词或短语生成分布。该分布随后被用于根据模型分配的概率选择最可能的下一个词或短语。softmax 单元在训练 LLMs 中起着关键作用，因为它使模型能通过调整神经网络的权重与偏置从数据中学习。在凸优化领域，如用中心路径法求解线性规划，softmax 函数已被用作控制势函数进展与稳定性的关键工具 [Cohen, Lee and Song STOC 2019, Brand SODA 2020]。本工作中，受 softmax 单元启发，我们定义了 softmax 回归问题。形式化地说，给定矩阵 $A \in \mathbb{R}^{n \times d}$ 与向量 $b \in \mathbb{R}^n$，目标是使用贪婪型算法求解 \begin{align*} \min_{x} \| \langle \exp(Ax), {\bf 1}_n \rangle^{-1} \exp(Ax) - b \|_2^2. \end{align*} 在某种意义下，我们可证明的收敛结果为实践中为何能用贪婪算法训练 softmax 函数提供了理论支撑。