Analog in-memory computing attention mechanism for fast and energy-efficient large language models

Nature Computational Science, vol. 5, pp. 813–824

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, graphics processing unit (GPU)-stored projections must be loaded into static random-access memory for each new generation step, causing latency and energy bottlenecks. Here 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 reduces attention latency and energy consumption by up to two and four orders of magnitude, respectively, compared with GPUs, marking a substantial step toward ultrafast, low-power generative transformers.

Authors 7

1. Nathan Leroux corresponding

   Forschungszentrum Jülich

   Affiliation as printed

   PGI-15, Forschungszentrum Jülich, Jülich, Germany. n.leroux@fz-juelich.de

2. Paul-Philipp Manea corresponding Aachen Faculty of Electrical Engineering

   Forschungszentrum Jülich · RWTH Aachen University

   Affiliation as printed

   Faculty of Electrical Engineering, RWTH Aachen, Aachen, Germany. p.manea@fz-juelich.de

   PGI-14, Forschungszentrum Jülich, Jülich, Germany. p.manea@fz-juelich.de

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