【2024】QKFormer: Hierarchical Spiking Transformer using Q-K Attention

基本信息:

标题: QKFormer: Hierarchical Spiking Transformer using Q-K Attention

作者: Chenlin Zhou, Han Zhang, Zhaokun Zhou 等

机构: Peng Cheng Laboratory, Harbin Institute of Technology, Peking University

发表日期: 2024年3月25日

论文链接: arXiv:2403.16552

代码和模型: GitHub - QKFormer

论文摘要

尖峰变换器（Spiking Transformers）结合了尖峰神经网络（SNNs）和变换器（Transformer）架构的优势，旨在提高能效和性能。然而，现有模型在处理复杂数据集上的视觉任务时仍面临性能瓶颈。本文提出了QKFormer，一种基于Q-K注意力机制的层次化尖峰变换器，通过直接训练实现了显著的性能提升。

QKFormer的核心创新包括：

• 引入了一种新的尖峰形式Q-K注意力机制，该机制通过二进制向量以线性复杂度高效地建模令牌或通道维度的重要性。
• 采用了层次结构，以实现多尺度的尖峰表示，这在大脑和人工神经网络中都显示出显著的性能提升。
• 设计了一个多功能且强大的补丁嵌入模块，特别为尖峰变换器定制，以增强尖峰信息传输。
  QKFormer在ImageNet-1K等主流数据集上取得了优异的性能，特别是在与Spikformer相比时，QKFormer在ImageNet-1K上达到了85.65%的top-1准确率，这是直接训练SNNs首次在该数据集上超过85%的准确率。此外，QKFormer的代码和模型已公开，供研究社区使用。

总体而言，QKFormer的开发标志着在变换器基础上的SNNs研究迈出了重要一步，为未来在能效和性能上进一步提升SNNs提供了新的方向。

背景与动机

论文背景:

尖峰神经网络（SNNs）因其与生物神经网络的高度相似性、事件驱动特性和低功耗而在神经形态硬件上显示出巨大潜力。然而，现有的SNN模型在处理复杂数据集上的视觉任务时性能受限，尤其是在变换器架构的实现上。

过去方案

以往的SNN模型通常采用直接训练或将预训练的人工神经网络（ANN）转换为SNN的方法。这些方法虽然在性能上取得了一定的进展，但在处理大规模数据集时仍面临梯度爆炸/消失问题和计算复杂度高的挑战。

论文的Motivation

为了克服现有模型的局限性，作者提出了一种新的层次化尖峰变换器模型，通过引入Q-K注意力机制和层次化架构来提高模型的性能和能效。

核心方法:

a. 理论背景:

QKFormer基于尖峰变换器的概念，结合了尖峰神经网络的生物学可信度和变换器架构的强大表示能力。

b. 技术路线:

• 提出了一种新的尖峰形式Q-K注意力机制，适用于SNNs。
• 引入了层次结构，以实现多尺度的尖峰表示。
• 设计了一个强大的补丁嵌入模块，特别为尖峰变换器定制。

c.核心贡献：

• 开发了一种新的Q-K注意力机制，具有线性复杂度，适用于SNNs的时空尖峰模式。
• 设计了一个多功能的补丁嵌入模块，显著提高了尖峰变换器的性能。
• 构建了一个基于Q-K注意力和补丁嵌入模块的层次化尖峰变换器，名为QKFormer。
• 在多个静态和神经形态数据集上，QKFormer的性能超过了现有的最先进SNN模型。

实验结果

a. 详细的实验设置

• 使用ImageNet-1K、CIFAR10/100等数据集进行评估。
• 使用AdamW优化器，采用RandAugment、随机擦除和随机深度等数据增强技术。
• QKFormer模型在ImageNet-1K上进行了200个周期的训练，批量大小为512。

b. 详细的实验结果

QKFormer在ImageNet-1K上达到了85.65%的top-1准确率和97.74%的top-5准确率。
在CIFAR10和CIFAR100数据集上，QKFormer分别达到了96.18%和81.15%的top-1准确率。
在DVS128 Gesture和CIFAR10-DVS等时间神经形态数据集上，QKFormer也显示出优越的性能。

总结

QKFormer通过引入创新的Q-K注意力机制和层次化架构，显著提升了SNN在图像识别任务上的性能。这项工作不仅推动了SNN的研究，也为未来的SNN应用提供了新的可能性。

. 缺陷和不足

尽管QKFormer在多个数据集上取得了显著的性能提升，但作为一种新兴的模型，可能还存在一些未被发现的问题和局限性。例如，模型在处理更复杂或不同类型数据集时的表现尚未得到充分验证。

可以进一步改进的点：

• 模型泛化能力: 探索QKFormer在更多种类的数据集上的性能，包括不同领域的图像识别和视频处理任务。
• 计算效率: 进一步优化模型的计算效率，以适应资源受限的设备。
• 模型解释性: 研究QKAttention机制的内部工作机制，提高模型的可解释性。
• 硬件实现: 探索QKFormer在神经形态硬件上的实现和优化，以充分发挥其低功耗的优势。