基于探针稀疏注意力机制的门控Transformer模型

赵婷婷; 丁翘楚; 马冲; 陈亚瑞; 王嫄

doi:10.13364/j.issn.1672-6510.20230145

天津科技大学学报 ›› 2024, Vol. 39 ›› Issue (03) : 56 -63. DOI: 10.13364/j.issn.1672-6510.20230145

基于探针稀疏注意力机制的门控Transformer模型

赵婷婷, 丁翘楚, 马冲, 陈亚瑞, 王嫄

作者信息 +

Author information +

文章历史 +

PDF (967K)

摘要

在强化学习中，智能体对状态序列进行编码，根据历史信息指导动作的选择，通常将其建模为递归型神经网络，但其存在梯度消失和梯度爆炸的问题，难以处理长序列。以自注意力机制为核心的Transformer是一种能够有效整合长时间范围内信息的机制，将传统Transformer直接应用于强化学习中存在训练不稳定和计算复杂度高的问题。门控Transformer-XL(GTr XL)解决了Transformer在强化学习中训练不稳定的问题，但仍具有很高的计算复杂度。针对此问题，本研究提出了一种具有探针稀疏注意力机制的门控Transformer(PS-GTr)，其在GTrXL中的恒等映射重排和门控机制的基础上引入了探针稀疏注意力机制，降低了时间复杂度和空间复杂度，进一步提高了训练效率。通过实验验证，PS-GTr在强化学习任务中的性能与GTrXL相当，而且训练时间更短，内存占用更少。