NATTEN: 为邻域带来注意力的高效神经网络扩展

NATTEN简介

NATTEN(Neighborhood Attention Extension)是一个为PyTorch提供快速滑动窗口注意力机制实现的开源项目。它的核心是邻域注意力(Neighborhood Attention),这是一种局部化的自注意力机制,可以在保持高效计算的同时捕获图像的局部上下文信息。

NATTEN的诞生源于研究人员对滑动窗口自注意力机制的长期关注。尽管这类方法在理论上很有效,但由于实现困难,它们一直被相对忽视。例如,早期提出这种方法的SASA论文就指出,虽然这些方法理论上很高效,但在实践中相对于卷积操作来说速度较慢。NATTEN的出现改变了这一局面,它为滑动窗口注意力机制提供了高效的实现,使其在实际应用中变得可行和有竞争力。

NATTEN的核心特性

NATTEN的主要特性包括:

1. 快速实现: NATTEN提供了邻域注意力机制的高效CUDA内核实现,大大提高了计算速度。

2. 多维支持: 支持1D、2D和3D数据的邻域注意力计算。

3. 灵活性: 提供了多种后端实现,包括naive、GEMM和FNA(Fused Neighborhood Attention)。

4. 兼容性: 支持PyTorch 2.0及以上版本,适配Python 3.8及以上版本。

5. 丰富的功能: 支持因果掩码、可变参数、相对位置偏置等高级功能。

NATTEN