Analog In-Memory Computing Attention Mechanism for Fast and Energy-Efficient Large Language Models
This addresses the problem of slow and energy-intensive generative Transformers for AI applications, representing a significant but incremental hardware optimization.
The paper tackles the latency and energy bottlenecks in self-attention mechanisms for large language models by proposing an analog in-memory computing architecture using gain cells, which reduces attention latency by up to two orders of magnitude and energy consumption by up to five orders of magnitude compared to GPUs.
Transformer networks, driven by self-attention, are central to Large Language Models. In generative Transformers, self-attention uses cache memory to store token projections, avoiding recomputation at each time step. However, GPU-stored projections must be loaded into SRAM for each new generation step, causing latency and energy bottlenecks. We present a custom self-attention in-memory computing architecture based on emerging charge-based memories called gain cells, which can be efficiently written to store new tokens during sequence generation and enable parallel analog dot-product computation required for self-attention. However, the analog gain cell circuits introduce non-idealities and constraints preventing the direct mapping of pre-trained models. To circumvent this problem, we design an initialization algorithm achieving text processing performance comparable to GPT-2 without training from scratch. Our architecture respectively reduces attention latency and energy consumption by up to two and five orders of magnitude compared to GPUs, marking a significant step toward ultra-fast, low-power generative Transformers.