用于表示学习的神经元注意力电路

《Neural Networks》：Neuronal Attention Circuit (NAC) for Representation Learning

　　摘要：注意力机制相较于RNNs能提升表示学习能力，但其离散特性限制了其对连续时间过程的建模能力。我们提出了神经元注意力电路（NAC），这是一种受生物学启发的连续时间注意力机制，它将注意力logit的计算转化为求解一阶线性常微分方程的问题，其中的非线性互连门控结构则是通过重新利用秀丽隐杆线虫的神经元电路策略（NCPs）的连接方式得到的。NAC用稀疏的感觉门控替代密集投影来处理查询-键值对，并引入具有两个分支的稀疏主干网络，用于计算内容-目标门控和可学习的时间常数门控，从而实现高效的自适应动态处理。为进一步提升效率并降低内存消耗，我们设计了可适配的、稀疏的次二次级Top-K成对连接机制，用以有选择地筛选查询-键值对之间的交互。我们提供了严格的理论保障，包括状态稳定性及有界的近似误差。在实证测试中，我们将NAC应用于多种领域，如不规则时间序列分析、长距离预测、自动驾驶车辆的车道保持以及工业预测任务。实验结果表明，NAC在准确率上可与现有的多种连续时间基准模型相媲美甚至更优，同时其在神经元层面也具备良好的可解释性。

引言：在时间或时空域中学习序列数据的表示形式，对于捕捉模式及实现精准预测而言至关重要。离散时间循环神经网络（DT-RNNs），例如RNN Jordan（1997）、Rumelhart等人（1985）提出的模型、长短期记忆网络（LSTM）Hochreiter和Schmidhuber（1997）以及门控循环单元（GRU）Cho等人（2014）的设计，都是通过迭代更新隐藏状态来表示或预测序列中的后续元素，这类方法在数据采样规律时效果良好，但由于假设时间间隔均匀，因此在处理采样不规则的数据时会面临挑战。此外，梯度消失问题也会使得捕捉长期依赖关系变得困难Hochreiter（1998）。

连续时间循环神经网络（CT-RNNs）Rubanova等人（2019）则将隐藏状态建模为常微分方程，这使得它们能够处理以任意或不规则时间间隔到达的输入数据。混合记忆循环神经网络（mmRNNs）Lechner和Hasani（2022）在此基础之上进一步将记忆模块与时间连续状态分开，从而在捕捉连续时间动态的同时保持稳定的误差传播。液态神经网络（LNNs）Hasani、Lechner、Amini、Liebenwein、Ray、Tschaikowski、Teschl、Rus，2022；Hasani、Lechner、Amini、Rus、Grosu，2021则采用受生物学启发的设计，为隐藏状态赋予不同的时间常数，以此提升模型的适应性和鲁棒性，不过训练过程中仍可能面临梯度消失的问题。

缩放点积注意力（SDPA）机制Vaswani等人（2017）通过平等对待所有时间步，并让模型能够聚焦于最相关的观测值，从而缓解了上述限制。该机制通过按键值维度进行缩放来计算查询项与键项之间的相似度，以此保持梯度的稳定性。多头注意力（MHA）Vaswani等人（2017）则在SDPA的基础上进一步拓展，允许模型同时关注不同的表示子空间。诸如稀疏注意力Zhang等人（2021）、BigBird Zaheer等人（2020）以及Longformer Beltagy等人（2020）等变体则通过仅关注选定的位置而非所有位置对，来调整注意力分配模式，从而降低计算成本，尤其适用于长序列处理。尽管有了这些改进，这类方法中的计算路径本质上仍是离散的，这限制了它们模拟连续轨迹的能力，而这类连续轨迹往往是其连续时间对应模型所能够处理的。

近期有一些研究试图通过神经常微分方程Chen等人（2018）的形式来弥合这一差距。mTAN Shukla和Marlin（2021）则学习连续时间嵌入，并利用基于时间的注意力机制将不规则的观测值映射为固定长度的表示形式，以便后续的编码器-解码器建模。连续时间注意力（CTA）Chien和Chen（2021）则将连续时间注意力机制嵌入到神经常微分方程中，使得注意力权重和隐藏状态能够随时间共同演化。不过，这类方法仍然存在计算成本较高的问题，且对常微分方程求解器的精度较为敏感。ODEFormer d’Ascoli等人（2023）则通过在合成轨迹上训练序列到序列的变换器，试图直接从带有噪声的不规则数据中推断出符号化的常微分方程系统，但它在使用时难以处理混沌系统，且在泛化能力方面也存在局限。ContiFormer Chen等人（2023）则通过将常微分方程定义的潜在轨迹与具有时间感知能力的注意力机制相结合，构建出了连续时间Transformer，用以模拟数据中的动态关系。

尽管取得了这些进展，但在开发一种能够将连续时间动态与大脑连接组抽象结构无缝整合的、受生物学启发的注意力机制方面，仍存在一个长期未被充分探索的空白。基于此，我们提出了一种名为神经元注意力电路（NAC）的新颖注意力机制，该机制将注意力logit的计算转化为求解一阶常微分方程的问题，而该方程则通过源自秀丽隐杆线虫神经系统中的神经元电路策略（NCPs）所衍生的非线性互连门控结构来进行调节（更多细节见附录A.1）。与通过线性层对查询、键值向量进行投影的SDPA不同，NAC采用感觉门控来转换输入特征，并借助主干网络来模拟非线性交互作用，同时通过多个分支生成用于计算注意力logit的有序输出。根据常微分方程的解，我们定义了三种计算模式：（i）精确模式，即使用常微分方程的封闭形式解；（ii）欧拉模式，即通过显式欧拉积分来近似求解；（iii）稳态模式，即仅使用稳态解，其方式类似SDPA的得分计算。为降低计算复杂度，我们设计了稀疏的次二次级Top-K成对连接算法，用以有选择地筛选查询-键值对输入。我们在多个领域对NAC进行了测试，包括采样不规则的时间序列、长距离预测、自动驾驶车辆的车道保持功能以及工业4.0相关应用，并将其与现有的最先进模型进行了对比。实验结果表明，NAC在性能上往往能与这些模型相媲美甚至更优，而在运行时间和峰值内存占用方面，它的表现则介于CT-RNNs与连续时间注意力机制之间。此外，我们还强调了NAC在细胞层面的可解释性。

我们的贡献可总结为以下几点：
•我们提出了神经元注意力电路（NAC），这是一种受生物学启发的新型注意力机制，它通过由非线性互连门控结构调控的、与输入相关的线性常微分方程来重新定义注意力logit的计算方式。
•我们重新利用了秀丽隐杆线虫的神经元电路策略（NCPs）的连接机制，以此构建出非线性门控结构，从而为该注意力机制带来稀疏性特征以及细胞层面的可解释性。
•我们在多种任务中对NAC进行了测试，这些任务包括：（i）不规则时间序列处理；（ii）长距离依赖关系分析；（iii）自动驾驶车辆的车道保持控制；以及（iv）工业4.0相关应用。

本文的其余结构如下：第2节回顾相关研究文献并阐述本研究的贡献；第3节详细介绍所提出算法的架构细节；第4节在多种学习任务中对该算法进行全面评估；第5节探讨该算法的局限性及未来的研究方向；第6节对全文进行总结。

相关研究：受生物学启发的模型：液态神经网络Hasani、Lechner、Amini、Liebenwein、Ray、Tschaikowski、Teschl、Rus，2022；Hasani、Lechner、Amini、Rus、Grosu，2021这类模型通过允许循环隐藏状态的积分时间常数通过非线性门控结构发生变化，从而实现输入依赖型的隐藏状态动态变化。这类模型通常被集成在NCPs Lechner、Hasani、Amini、Henzinger、Rus、Grosu，2020；Lechner、Hasani、Grosu框架之中，充分利用秀丽隐杆线虫网络结构的稀疏连接特性。

神经元注意力电路（NAC）：我们提出了一种更为简单的注意力logit计算公式（更多细节见附录A.2），将该公式理解为由非线性互连门控结构调控的一阶线性常微分方程的解：$\frac{d\mathbf{a}}{dt}=-f_{\omega_t}(\mathbf{[q;k]},t,\theta_{\omega_t})\mathbf{u}+\frac{f_{\phi}(\mathbf{[q;k]},t,\theta_{\phi})}{\omega_t}\mathbf{u}$，其中$\mathbf{u}=\mathbf{[q;k]}$代表经过稀疏Top-K连接的查询-键值对输入。$\omega_t$表示可学习的时间常数门控分支，而$\phi$则表示非线性的内容-目标门控分支。这两个门控结构都由主干网络进行参数化处理。

评估：我们对NAC与多种基准模型进行了对比测试，这些基准模型包括CT-RNNs、（离散时间与连续时间）注意力机制以及多种经过改进的版本。实验覆盖了多种不同的应用场景，包括不规则时间序列分析、长距离依赖关系识别、自动驾驶车辆的车道保持功能以及工业4.0领域的预测任务。所有实验结果均通过5折交叉验证获得，即在每折数据上使用BPTT算法（详见附录C.2）对模型进行训练，然后再在所有折的数据上对模型性能进行评估。我们报告的是各项指标的平均值（μ）以及...

讨论：本文属于关于具有生物学合理性的注意力机制的持续研究工作之一，可视为该领域的一项开创性成果，当然未来仍需解决一些存在的局限性。

架构改进：目前，NAC采用预先设定的连接结构（AutoNCP），该结构需要三个输入参数：单元数量（包括中间神经元和运动神经元）、运动神经元数量以及系统的稀疏程度。为了在保留原有连接结构的同时将其与注意力机制整合在一起，我们设定神经网络中的感觉单元数量为$unit_{sensory}=\lceil d_{model}-0.5\times0.6\rceil$，而主干网络的单元数量则...

结论：在本文中，我们提出了神经元注意力电路（NAC），这是一种受生物学启发的注意力机制，它通过重新利用秀丽隐杆线虫的NCPs所衍生的非线性互连门控结构，将注意力logit的计算转化为求解一阶常微分方程的问题。NAC成功实现了离散注意力机制与连续时间动态处理的结合，使得模型能够实现自适应的时间处理能力，且不存在传统缩放点积注意力机制所固有的诸多限制。基于常微分方程的解，我们还进一步提出了三种...

可复现性声明：用于保证研究可复现性的代码可在https://github.com/itxwaleedrazzaq/neuronal_attention_circuit地址获取。

伦理审批声明：本研究是按照相应的伦理标准开展的。

CRediT作者贡献说明：Waleed Razzaq：负责原始稿件的撰写、模型验证、方法设计、实验实施、正式分析以及数据整理工作；Yun-Bo Zhao：负责项目整体监督、资源协调、项目管理以及概念构思工作。

利益冲突声明：作者们声明自己不存在任何可能影响本文研究结果的已知财务利益关系或个人关系。

致谢：本研究得到了中国科学技术协会-国际科学组织联盟奖学金的支持。我们同时感谢中国科学技术大学以及国际科学组织联盟在知识与物质资源方面的支持。在本文的撰写过程中，还使用了人工智能/大型语言模型工具，但所有相关工作都在严格的人工监督下完成。

Waleed Razzaq|Yun-Bo Zhao