论文
Kalman Delta Networks:不确定性感知联想记忆
Kalman Delta Networks: Uncertainty-aware Associative Memory
摘要
线性注意力越来越多地应用于前沿语言模型中,以实现高效的长上下文推理和常量内存解码。然而,其固定大小的循环内存需要对每个标记进行在线决策:在知道未来查询需要哪些信息之前,要写入什么以及覆盖现有关联的强度。 Delta 规则模型从当前的词元嵌入中学习这种优势,但不跟踪内存估计的置信度,从而阻止每次写入适应累积的证据。为了明确地表示这种不确定性,我们将循环联想记忆重新表述为线性高斯状态空间模型,其中卡尔曼滤波器是最佳递归估计器,并引入了一个新的模型系列:卡尔曼增量网络(KDN)。在 KDN 内,转换会传播内存状态及其不确定性,从而允许卡尔曼增益通过积累的证据和观察可靠性来加权每个残留写入。在这种表述下,Delta 式更新作为一种特殊情况出现,它用 token-wise 各向同性代理替代预测协方差并省略协方差跟踪。然而,精确跟踪需要密集的、状态相关的 Riccati 递归,该递归不太适合 GPU 并行的线性注意力扫描。为了解决这个问题,我们引入了两种扫描兼容的 KDN 近似。对角 KDN 通过在线平均场变分推理将每个后一步投影到对角高斯族上,而各向同性 KDN 使用每个头具有单个不确定性标量的各向同性近似。它们的不确定性重现是莫比乌斯图,可以实现具有对数平行深度的关联扫描。在 750M 和 1.3B 参数的受控预训练中,与最先进的线性注意力模型相比,KDN 变体持续改善了困惑度和平均下游准确性。