[论文解读] In-Context Learning for Attention Scheme: from Single Softmax Regression to Multiple Softmax Regression via a Tensor Trick

Yeqi Gao, Zhao Song+1 more|arXiv (Cornell University)|Jul 5, 2023

Tensor decomposition and applicationsMathematics被引用 3

一句话总结

本文提出了一种基于张量技巧的注意力机制少样本学习框架，通过向量化变换重新表述基于矩阵的softmax回归，以提升稳定性和泛化能力。该研究建立了损失函数的Lipschitz连续性，并证明了模型权重或参数的扰动会导致输出分布有界偏移，从而在归一化与缩放的softmax公式下实现鲁棒的少样本自适应，且其依赖关系呈指数形式，与模型容量和输入维度成正比。

ABSTRACT

Large language models (LLMs) have brought significant and transformative changes in human society. These models have demonstrated remarkable capabilities in natural language understanding and generation, leading to various advancements and impacts across several domains. We consider the in-context learning under two formulation for attention related regression in this work. Given matrices $A_1 \in \mathbb{R}^{n imes d}$, and $A_2 \in \mathbb{R}^{n imes d}$ and $B \in \mathbb{R}^{n imes n}$, the purpose is to solve some certain optimization problems: Normalized version $\min_{X} \| D(X)^{-1} \exp(A_1 X A_2^ op) - B \|_F^2$ and Rescaled version $\| \exp(A_1 X A_2^ op) - D(X) \cdot B \|_F^2$. Here $D(X) := \mathrm{diag}( \exp(A_1 X A_2^ op) {\bf 1}_n )$. Our regression problem shares similarities with previous studies on softmax-related regression. Prior research has extensively investigated regression techniques related to softmax regression: Normalized version $\| \langle \exp(Ax) , {\bf 1}_n angle^{-1} \exp(Ax) - b \|_2^2$ and Resscaled version $\| \exp(Ax) - \langle \exp(Ax), {\bf 1}_n angle b \|_2^2 $ In contrast to previous approaches, we adopt a vectorization technique to address the regression problem in matrix formulation. This approach expands the dimension from $d$ to $d^2$, resembling the formulation of the regression problem mentioned earlier. Upon completing the lipschitz analysis of our regression function, we have derived our main result concerning in-context learning.

研究动机与目标

• 通过softmax函数的矩阵表述形式，形式化基于注意力机制的回归任务中的少样本学习。
• 分析在权重空间与输入空间中，参数扰动下softmax损失函数的Lipschitz连续性。
• 建立模型权重或输入变化对归一化与缩放的softmax回归输出分布影响的理论边界。
• 证明小幅度的参数更新可导致模型输出产生受控的、有界的改变，从而实现稳定的少样本自适应。
• 通过向量化与张量技术，将先前关于softmax回归的研究统一并推广至矩阵形式的注意力机制。

提出的方法

• 本文采用向量化技术，将矩阵回归问题转化为高维向量空间，将参数维度从 $ d $ 扩展至 $ d^2 $。
• 提出两种公式：归一化版本 $ \min_X \|D(X)^{-1}\exp(A_1 X A_2^\top) - B\|_F^2 $ 和缩放版本 $ \|\exp(A_1 X A_2^\top) - D(X)\cdot B\|_F^2 $，其中 $ D(X) $ 为指数激活行求和的对角矩阵。
• 分析依赖于对关键函数 $ u(x), \alpha(x), f(x), h(x), c(x), q(x) $ 及其逆函数的有界性，利用Lipschitz连续性控制敏感度。
• 推导出损失函数梯度的Lipschitz界，表明输出变化与参数变化成正比，比例因子为 $ M = \exp(O(R^2 + \log n)) $。
• 该框架证明了对 $ x $ 或 $ \mathsf{A} $ 的扰动会导致输出分布有界偏移，满足 $ \|\delta_c\|_2 \leq M \cdot \|x_{t+1} - x_t\|_2 $。
• 通过递归有界中间函数的范数并利用矩阵指数与对角缩放的性质，推导出理论结果。

实验结果

研究问题

• RQ1 将矩阵参数向量化为 $ \mathbb{R}^{d^2} $ 如何影响基于softmax的注意力回归的稳定性和泛化能力？
• RQ2 在参数扰动下，归一化与缩放的softmax损失函数梯度的Lipschitz常数是多少？
• RQ3 模型权重或输入的微小变化在多大程度上会导致注意力机制输出分布的有界变化？
• RQ4 能否通过矩阵注意力回归的张量技巧重构，理论上解释大语言模型的少样本学习行为？
• RQ5 边界 $ M $ 中对 $ R^2 $ 与 $ \log n $ 的指数依赖关系，如何影响少样本自适应的鲁棒性？

主要发现

• 归一化softmax损失函数在参数向量 $ x $ 与矩阵 $ \mathsf{A} $ 上均表现出Lipschitz连续性，其梯度变化受 $ M \cdot \|x_{t+1} - x_t\|_2 $ 有界，其中 $ M = \exp(O(R^2 + \log n)) $。
• 对于基于 $ x $ 的更新，输出偏移 $ \|\delta_c\|_2 $ 受限于 $ M \cdot \|x_{t+1} - x_t\|_2 $，确保在小参数变化下实现稳定的少样本自适应。
• 对于基于 $ \mathsf{A} $ 的更新，输出偏移同样受 $ M \cdot \|A_{t+1} - A_t\|_2 $ 有界，证实对权重扰动具有鲁棒性。
• 理论分析表明，从 $ x_t $ 到 $ x_{t+1} $ 或从 $ A_t $ 到 $ A_{t+1} $ 的转换，会生成一个具有扰动目标 $ \widetilde{b} $ 或 $ \widehat{b} $ 的新稳定回归问题，且满足 $ \|\widetilde{b} - b\|_2 \leq M \cdot \|x_{t+1} - x_t\|_2 $。
• 边界通过递归应用中间函数（如 $ c(x) = \alpha(x)^{-1} u(x) $）的Lipschitz性质推导得出，利用指数与多项式项对范数增长实现紧密控制。
• 该框架通过证明模型更新仅引起输出的受控、可预测变化，为注意力机制中的少样本学习提供了理论基础，从而支持可靠的 few-shot 泛化。