Analog In-Memory Computing Attention Mechanism for Fast and Energy-Efficient Large Language Models

JuSER (Forschungszentrum Jülich)

Abstract

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.

Authors 7

1. Nathan Leroux

   Forschungszentrum Jülich

   Affiliation as printed

   PGI-15 , Forschungszentrum Jülich , Jülich , Germany

2. Paul-Philipp Manea Aachen Faculty of Electrical Engineering

   Forschungszentrum Jülich · RWTH Aachen University

   Affiliation as printed

   Faculty of Electrical Engineering RWTH Aachen , Aachen , Germany

   PGI-14 , Forschungszentrum Jülich , Jülich , Germany

3. Chirag Sudarshan

   Forschungszentrum Jülich

   Affiliation as printed

   PGI-14 , Forschungszentrum Jülich , Jülich , Germany

4. Jan Finkbeiner Aachen Faculty of Electrical Engineering

   Forschungszentrum Jülich · RWTH Aachen University

   Affiliation as printed

   Faculty of Electrical Engineering RWTH Aachen , Aachen , Germany

   PGI-15 , Forschungszentrum Jülich , Jülich , Germany

5. Sebastian Siegel

   Forschungszentrum Jülich

   Affiliation as printed

   PGI-14 , Forschungszentrum Jülich , Jülich , Germany

6. John Paul Strachan Aachen Faculty of Electrical Engineering

   Forschungszentrum Jülich · RWTH Aachen University

   Affiliation as printed

   Faculty of Electrical Engineering RWTH Aachen , Aachen , Germany

   PGI-14 , Forschungszentrum Jülich , Jülich , Germany

7. Emre Neftci Aachen Faculty of Electrical Engineering

   Forschungszentrum Jülich · RWTH Aachen University

   Affiliation as printed

   Faculty of Electrical Engineering RWTH Aachen , Aachen , Germany

   PGI-15 , Forschungszentrum Jülich , Jülich , Germany