Presentation Video

Official Example (IBM AIHWKit): examples/36_gpt2_on_wikitext.py
My Fork: github.com/gyulab/aihwkit (GPT-2, GPT-2 notebook, BERT noise analysis)

Update) The GPT-2 example I started in summer 2024 is now part of the official IBM AIHWKit examples as Example 36 (“This example was initiated by Gyujun Jeong”), released with AIHWKit 1.1.0 and later versions.

Motivation: Transformers on Analog In-Memory Computing

Analog in-memory computing (AIMC) performs matrix-vector multiplications directly inside non-volatile memory crossbars, avoiding the data movement that dominates the energy of large neural networks. However, analog weights are noisy: programming noise, read noise, conductance drift, and limited ADC/DAC precision all perturb the computation.

IBM Analog Hardware Acceleration Kit (AIHWKit) is an open-source PyTorch toolkit that simulates these non-idealities, so we can ask: how do modern transformer models behave on analog hardware, and which layers are most sensitive to noise? In this project, I explored two transformer families:

GPT-2 (decoder): an end-to-end example for analog conversion, hardware-aware (HWA) fine-tuning, and text generation.
BERT (encoder): a layer-by-layer noise sensitivity analysis on question answering.

Part 1. GPT-2 on AIHWKit

Model & Dataset

Model: DistilGPT2 (88.2M parameters), the smallest GPT-2 variant, for demonstration.
Dataset: wikitext-2-raw-v1 (~18 MB on disk).
Workflow: load the pre-trained model from Hugging Face → convert it to analog with convert_to_analog() → HWA fine-tuning → evaluate perplexity at different weight noise levels and inference times (drift), logged to TensorBoard.

Analog Layers

After convert_to_analog(), the LM head (768 → 50257) becomes an AnalogLinear layer mapped onto two simulated crossbar tiles. In this demonstration, the GPT-2 blocks use Hugging Face Conv1D modules (attention c_attn/c_proj, MLP c_fc/c_proj), which were not replaced by the converter, so they stay digital along with the embeddings and LayerNorms.

Results

Model	Validation Loss	Perplexity (PPL)
Digital (reference, Intel/distilgpt2-wikitext2)	3.326	27.8
Analog, HWA fine-tuned (AIHWKit)	4.059	57.9

The analog model still generates coherent text from a prompt, showing that the full pipeline (conversion → HWA training → generation) works end-to-end:

Input: “Once upon a time”
Output: “Once upon a time of war, the German army was able to defeat the German army in the Battle of the Rhine”

I opened a pull request to IBM/AIHWKit (#664). The example was later integrated by the maintainers into the official repository as Example 36.

Part 2. BERT Noise Analysis on SQuAD

Setup

Task: SQuAD v1.1 (100k+ question-answer pairs). BERT takes [CLS] question [SEP] passage [SEP] and predicts start/end logits of the answer span.
Metrics: F1 score (token overlap between predicted and ground-truth spans) and Exact Match (EM).
Analog configuration: PCM-like noise model, 512×512 tiles, 256-level DAC/ADC resolution, output noise, and global drift compensation.
Method: For each layer, one at a time, inject Gaussian weight noise (mean = 0, std = 0.1) and re-evaluate F1/EM. This isolates how much each layer contributes to the accuracy loss.

Baseline

Without additional noise, the analog BERT achieves F1 ≈ 88 and EM ≈ 81, stable over drift times up to 10⁶ s. This is the reference for the layer-wise analysis.

Layer-wise Sensitivity

Layer Group	F1 / EM under noise	Observation
Embeddings (word, position, token type)	F1 0.59–4.19, EM < 1	Most vulnerable: noise corrupts the input representation itself
Attention (query/key/value)	F1 86.75–87.47, EM 78.67–79.67	Robust: attention tolerates noise to a certain extent
Intermediate / output dense	Mostly ~87, worst F1 77.13, EM 63.69	Complex output dense layers can be highly sensitive
LayerNorm	F1 ~87, EM ~79	Strong resistance; normalization stabilizes small perturbations
QA outputs	F1 ~87, EM ~79	Robust: directly reads the contextualized embeddings

Noise propagation: Noise injected in earlier layers propagates and can be amplified downstream. For example, the same output dense tile gives F1 87.23 in encoder layer 0 but drops to F1 77.13 in layer 3.

Why Are Some Layers Vulnerable?

Hypothesis 1, Weight magnitude: “Heavier weights dominate, so they are more noise-sensitive.” I correlated the mean and std of each tile’s weights with its F1 score. The correlation coefficients were only −0.005 (mean) and 0.090 (std), so weight magnitude is not the explanation.

Hypothesis 2, Distribution spread: Comparing the vulnerable layer (output.dense ... array.3.0, F1 77.13) with an average layer (F1 88.24), the vulnerable layer shows a wider weight distribution and less consistent neuron activations. This higher variability makes it less able to generalize under perturbation, which likely explains its lower F1.

Summary

GPT-2 on AIHWKit: Built an end-to-end analog GPT-2 example (conversion, HWA fine-tuning, drift-aware evaluation, text generation), now included in the official IBM AIHWKit as Example 36.
Layer importance: In BERT, the first layers (embeddings) and complex final transformations (output dense) are the most critical, since they handle the initial encoding of raw input and the final prediction.
Layer functionality: Layers with complex transformations and wider weight/activation distributions are more susceptible to analog noise, while attention, LayerNorm, and QA outputs are comparatively robust.
Takeaway for hardware design: Not all layers need the same analog precision. Keeping sensitive layers (e.g., embeddings) digital or giving them more robust mapping is a practical path for transformer inference on AIMC.

IBM Analog Hardware Accelerator Kit Research Archive