用LSTM替代注意力机制,让Transformer更适配内存计算芯片
FAR: Function-preserving Attention Replacement for IMC-friendly Inference
- 用双向LSTM替换注意力模块,保持功能等价性
- 在ImageNet上精度接近原模型,延迟显著降低
- 适合资源受限的内存计算边缘设备部署
尽管变压器模型在视觉与语言任务中占据主导地位,但其注意力机制因大量激活间乘法和非局部内存访问,难以适配基于忆阻器的内存计算(IMC)设备,导致在ReRAM加速器上产生显著延迟与带宽开销。为此,我们提出FAR——一种功能保持型注意力替换框架,将预训练DeiT中的所有注意力模块替换为天然兼容IMC数据流的序列化模块。具体而言,FAR通过分块蒸馏将自注意力替换为多头双向LSTM架构,在保持功能等价的同时实现线性时间计算与局部权重复用。进一步结合结构化剪枝,使FAR模型可灵活适配资源受限的IMC阵列,同时维持功能保真度。在DeiT系列上的评估表明,FAR在ImageNet及多个下游任务上保持接近原始注意力模型的精度,参数量与延迟均减少。分析显示,FAR保留了注意力学习到的语义标记关系,同时提升计算效率,展现出在基于IMC的边缘加速器上实现节能推理的巨大潜力。
原文摘要 · Abstract (English)
While transformers dominate modern vision and language models, their attention mechanism remains poorly suited for in-memory computing (IMC) devices due to intensive activation-to-activation multiplications and non-local memory access, leading to substantial latency and bandwidth overhead on ReRAM-based accelerators. To address this mismatch, we propose FAR, a Function-preserving Attention Replacement framework that substitutes all attention in pretrained DeiTs with sequential modules inherently compatible with IMC dataflows. Specifically, FAR replaces self-attention with a multi-head bidirectional LSTM architecture via block-wise distillation to retain functional equivalence while enabling linear-time computation and localized weight reuse. We further incorporate structured pruning on FAR models, enabling flexible adaptation to resource-constrained IMC arrays while maintaining functional fidelity. Evaluations on the DeiT family demonstrate that FAR maintains comparable accuracy to the original attention-based models on ImageNet and multiple downstream tasks with reduced parameters and latency. Further analysis shows that FAR preserves the semantic token relationships learned by attention while improving computational efficiency, highlighting its potential for energy-efficient transformer inference on IMC-based edge accelerators.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。