基于结构化状态空间对偶性的贝叶斯注意力机制设计与实现

当前的大型语言模型在处理长序列文本时面临挑战。主要的瓶颈在于注意力机制，它将文本处理为单词（或 tokens）序列。注意力计算的复杂度随序列长度 T 呈平方增长，导致处理长文本的成本显著增加。为了降低计算成本，研究人员积极探索注意力的替代方案，包括递归模型（如 Mamba [1] 和 xLSTM [2]）、卷积模型（如 Hyena [3]）以及基于稀疏性的模型（如 Longformer [4] 和 BigBird [5]）。其中状态空间模型作为一种有效的注意力替代方案，受到了越来越多的关注。例如基于状态空间模型的大型语言模型（LLM），如 Mamba 2 [6]，在多项任务上表现出与 transformers 相当甚至更优的性能 [7]。状态空间模型还为理解注意力与半可分离矩阵和状态空间模型之间的关系提供了新的视角 [6]。

然而当前基于注意力和状态空间的模型通常需要大量的训练数据。当训练数据有限，或者需要将领域知识融入模型时，贝叶斯方法是一种有效的选择。与标准神经网络训练不同，贝叶斯模型不易受到过度自信的影响 [8]，并且支持利用未标记数据进行训练。此外贝叶斯模型能够提供不确定性估计，这在金融和医疗等高风险领域具有重要价值。

本文介绍了一种贝叶斯风格的注意力机制，用于序列预测。我们将详细阐述如何使用马尔可夫链蒙特卡罗法（MCMC）训练该模型。

贝叶斯注意力适用性评估

以下是一些建议，帮助您快速评估贝叶斯注意力是否适合您的应用场景：

如果满足以下一个或多个条件，可以考虑使用贝叶斯注意力：

• 数据记录具有顺序依赖性（如文本或时间序列）。
• 需要对一个或多个连续值进行预测（回归任务）。
• 数据集规模较小，例如仅有少量样本。
• 团队拥有希望融入模型的领域知识或启发式规则。
• 需要对预测结果和模型参数进行不确定性评估。
• 每个输出 y 都有对应的输入变量 x。
• 部分输出 y 可能存在缺失值。

注意力机制回顾

 

https://avoid.overfit.cn/post/2af685d939434a668f734f33ee2c3388