Analog Foundation Models

TL;DR

The paper tackles the barrier of running trillion-parameter LLMs on analog in-memory computing hardware by introducing a scalable hardware-aware training pipeline. It shows that analog foundation models can approach the performance of $4$-bit weight/$8$-bit activation baselines on challenging models like Phi-3-mini-4k-instruct and Llama-3.2-1B-Instruct, even under hardware- realism noise, and can be quantized post-training for digital accelerators. Additionally, the authors report that test-time compute scaling yields superior scaling behavior compared to models trained with heavy quantization, highlighting a path toward energy-efficient, high-capacity foundation models. The work combines synthetic data generation, distillation, and hardware-aware training to produce robust models that generalize across memory technologies and noise profiles, contributing a practical route to deployable analog AI with broad implications for edge and green AI systems.

Abstract

Analog in-memory computing (AIMC) is a promising compute paradigm to improve speed and power efficiency of neural network inference beyond the limits of conventional von Neumann-based architectures. However, AIMC introduces fundamental challenges such as noisy computations and strict constraints on input and output quantization. Because of these constraints and imprecisions, off-the-shelf LLMs are not able to achieve 4-bit-level performance when deployed on AIMC-based hardware. While researchers previously investigated recovering this accuracy gap on small, mostly vision-based models, a generic method applicable to LLMs pre-trained on trillions of tokens does not yet exist. In this work, we introduce a general and scalable method to robustly adapt LLMs for execution on noisy, low-precision analog hardware. Our approach enables state-of-the-art models $\unicode{x2013}$ including Phi-3-mini-4k-instruct and Llama-3.2-1B-Instruct $\unicode{x2013}$ to retain performance comparable to 4-bit weight, 8-bit activation baselines, despite the presence of analog noise and quantization constraints. Additionally, we show that as a byproduct of our training methodology, analog foundation models can be quantized for inference on low-precision digital hardware. Finally, we show that our models also benefit from test-time compute scaling, showing better scaling behavior than models trained with 4-bit weight and 8-bit static input quantization. Our work bridges the gap between high-capacity LLMs and efficient analog hardware, offering a path toward energy-efficient foundation models. Code is available at https://github.com/IBM/analog-foundation-models.

Figures (11)

• Figure 1: a. Linear layers in the transformer architecture can be offloaded to cores. b. In one implementation, given a digital weight matrix $\mathbf{W}$, each core calculates $\mathbf{y}=\mathbf{x} \mathbf{W}^T$ in the analog domain using voltages $\mathbf{v}$ and conductances $\mathbf{G}$, where $G_{i,j} \propto W_{i,j}$.
• Figure 2: a. Synthetic data is generated by repeated sampling from the softmax distribution. b. Analog foundation models are trained via knowledge distillation on the synthetic data. c. The trained models can then be deployed on analog or low-precision digital hardware.
• Figure 3: Average benchmark performance (y-axis) as a function of additive Gaussian noise magnitude (x-axis) expressed in percentage of per-channel maximum absolute weight. Results shown for a.Phi-3-mini-4k-instruct- and b.Llama-3.2-1B-Instruct-based models.
• Figure 4: Math-500 accuracy (y-axis) as a function of number of generations (x-axis). Results shown for the original models, LLM-QAT models, and analog foundation models based on a.Phi-3-mini-4k-instruct- and b.Llama-3.2-1B-Instruct-based models.
• Figure 5: Sweep over amount of noise injected during training.