Topkima-Former: Low-Energy, Low-Latency Inference for Transformers Using Top-k In-Memory ADC

Transformer has emerged as a leading architecture in neural language processing (NLP) and computer vision (CV). However, the extensive use of nonlinear operations, like softmax, poses a performance bottleneck during transformer inference and comprises up to 40% of the total latency. Hence, we propose innovations at the circuit, algorithm and architecture levels to accelerate the transformer. At the circuit level, we propose Topkima—combining top-k activation selection with in-memory ADC (IMA) to implement efficient softmax without any sorting overhead. Only the k largest activations are sent to softmax calculation block, reducing the huge computational cost of softmax. At the algorithmic level, a modified training scheme utilizes top-k activations only during the forward pass, combined with a sub-top-k method to address the crossbar size limitation by aggregating each sub-top-k values as global top-k. At the architecture level, we introduce a fine pipeline for efficiently scheduling data flows and an improved scale-free technique for removing scaling cost. The combined system, dubbed Topkima-Former, enhances $1.8\times -84\times $ speedup and $1.2\times -36\times $ energy efficiency (EE) over prior In-memory computing (IMC) accelerators. Compared to a conventional softmax macro and a digital top-k (Dtopk) softmax macro, our proposed Topkima softmax macro achieves about $15\times $ and $8\times $ faster speed respectively. Experimental evaluations demonstrate minimal (0.42% to 1.60%) accuracy loss for different models in both vision and NLP tasks.