FAR: Function-preserving Attention Replacement for IMC-friendly Inference

TL;DR

FAR introduces a function-preserving attention replacement for transformers, enabling efficient IMC-compatible inference with minimal accuracy loss and reduced latency.

Contribution

It proposes a novel attention replacement framework that retains model performance while optimizing for in-memory computing hardware.

Findings

FAR maintains comparable accuracy to original models on ImageNet.

FAR reduces model parameters and latency significantly.

Structured pruning enables resource adaptation without accuracy loss.

Abstract

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.