ME-MoD: Algorithm-Hardware-Dataflow Co-Design for Memory Efficient Mixture-of-Depth-Based Vision Transformer Accelerator

Abstract

Transformers outperform traditional neural networks but face high computational and memory costs, limiting edge device deployment. Although many hardware accelerators aim to address this, the original Transformer structure still restricts the optimization effect. A recent breakthrough, mixture-of-depths (MoDs), employs conditional computation and effectively reduces the computational complexity of large language models, providing a valuable opportunity for designing an efficient hardware accelerator. However, when applied to vision transformers, MoD suffers from accuracy degradation and excessive external memory access (EMA). Therefore, this article presents ME-MoD, the first memory-efficient MoD-based vision transformer inference accelerator, leveraging the idea of reordering and algorithm-hardware-dataflow codesign. Algorithmically, distribution adjustment forward (DAF) and routing decision forward (RDF) techniques restore accuracy and alleviate memory access costs through token reordering. Architecturally, a LayerNorm-Routing (L-R) fusion module and a token reordering and sequential recording module enhance computational efficiency while minimizing memory overhead. In addition, a token-stationary layer fusion dataflow and an on-chip dynamic memory module are designed, which further optimizes the EMA caused by the intermediate results of interlayer computation of valid tokens routed by MoD. With negligible accuracy loss, our ME-MoD accelerator achieves $1.62\times $ inference speed up, eliminates 46.5% of the external memory bandwidth requirement and 45.2% of energy consumption compared with standard MoD. It achieves 23.6 TOPS/W energy efficiency, which is $4.02\times $ improvements compared with state-of-the-art (SOTA) designs.