自注意力机制与工作记忆容量限制：基于Transformer的模型研究

近期关于基于Transformer的大语言模型（LLMs）的研究发现，其工作记忆容量存在显著限制，这与人类行为学研究的发现相似。具体而言，随着N-back任务中N的增大，模型性能显著下降。然而，对于这一现象产生的原因，目前仍缺乏机制层面的可解释性。受行为科学中执行注意理论的启发，我们假设基于Transformer的模型中的自注意力机制可能是其工作记忆容量限制的原因。为验证这一假设，我们训练了仅含解码器的Transformer模型来执行N-back任务，并发现随着训练进行，注意力分数逐渐聚集到N-back位置，这表明模型通过学会一种策略——关注当前位置与N-back位置之间的关系——来掌握任务。关键的是，我们发现注意力分数矩阵的总熵随N的增大而增加，这表明注意力分数的分散可能是N-back任务中观察到的容量限制的原因。因此，我们的研究为注意力在人类和人工智能中的共同作用提供了见解。此外，本研究所揭示的自注意力机制的局限性，可为未来设计具有增强工作记忆容量和认知能力的模型架构提供参考。