Processing-in-memory DRAM chips deliver huge bandwidth by putting tiny compute units next to the memory banks, but those units are so weak that they can only help the memory-bound scraps of a neural network — under 15% of the runtime — while the compute-heavy matrix multiplications that dominate still need a powerful host GPU; PIM-DL (Peking University + Microsoft Research) sidesteps the weakness instead of fighting it, replacing every matrix-multiply in a Transformer's linear layers with table lookups (a LUT-NN, where inputs are matched to a few precomputed "centroids" and results are read from tables rather than multiplied), which turns compute-heavy layers into the bandwidth-heavy, multiplication-light work PIM is actually good at — offloading most of the network to commodity PIM and beating GEMM-on-PIM by 22–37×, a CPU by up to 3.5×, and even a GPU by up to 1.2×.
The argument in one glance
-
Problem
The memory-wall cure that can only touch the easy parts
Moving data between memory and processor is the dominant cost and energy drain in modern computing — the "memory wall." Commodity DRAM processing-in-memory (DRAM-PIM) attacks it directly: put small compute units right next to the DRAM banks, and you get up to 8× the memory bandwidth with far less data-movement energy. Products now ship (UPMEM's PIM-DIMM, Samsung's HBM-PIM, SK-Hynix's AiM). But there is a catch buried in the physics: because those compute units are fabricated in the DRAM process, they are feeble — and a neural network's heaviest work slips right through their fingers.
-
What people did
Offload the memory-bound operators to PIM, keep a big host for the rest
The deployed approach: run the operators PIM is good at — the memory-bound ones like ReLU, residual adds, layer-norm, and matrix-vector products (GEMV) — on the PIM, and send the compute-heavy matrix multiplications (GEMM) back to a powerful host CPU or GPU. This "PIM-enabled system" does speed things up, and it is the natural division of labor given PIM's limits.
-
Why it fell short
The easy parts are a sliver; the hard parts still need the host
DRAM-PIM compute units are built with DRAM transistors — about 3× slower than CMOS, with several times lower logic density and fewer metal layers for routing — so they cannot host powerful arithmetic. UPMEM tops out at 43.8 GOP/s per DIMM; even HBM-PIM and AiM, with dedicated vector units, reach only ~1 TFLOP/s, while a Transformer's GEMM layers demand 10+ TOP/s for adequate throughput. And the operators PIM CAN accelerate — element-wise ops — are under 15% of end-to-end latency; GEMV helps only single-batch inference, whereas cloud serving is batched and GEMM-dominated. So the expensive host still does most of the work, the PIM sits underutilized, and the manufacturing cost went UP without proportional benefit. Making the DNN lighter by brute force (binary networks) wrecks accuracy on models like BERT and ViT.
-
This paper's idea
Change the workload so the hard parts become PIM-friendly
PIM-DL (Peking University + Microsoft Research) stops trying to make weak compute units do heavy math and instead removes the heavy math. Its philosophy: replace the GEMM in every linear layer with LOOKUP TABLES. A LUT-NN observes that a layer's input activations cluster into a few typical patterns (centroids); precompute each centroid's product with the weights and store it in a table; then at inference, just match each input to its nearest centroid and READ the answer from the table — reducing multiplications by up to 18× and leaving a workload that is bandwidth-heavy and multiplication-light, exactly what PIM does well. Three pieces make it work: eLUT-NN, a calibration algorithm that converts EVERY linear layer to LUTs at near-original accuracy using under 1% of the training data; an inference engine that offloads LUT operators to PIM and the rest to a wimpy host; and an auto-tuner that maps LUT kernels efficiently onto each PIM platform. Result: 22–37× over GEMM-on-PIM, up to 3.5× over a CPU, and up to 1.2× over a GPU.
Why this problem is worth solving
The single most expensive thing a computer does is move
data, and neural-network inference is drowning in it —
weights and activations shuttling between memory and
compute, over and over. Putting compute inside the memory
is the most direct answer, and
DRAM-PIM: memory chips
with small processing units placed right next to the DRAM
banks, so computation happens where the data already lives
instead of hauling it across the bus. This yields up to 8×
the effective memory bandwidth and slashes data-movement
energy. Commodity products ship today — UPMEM’s DDR4-based
PIM-DIMM (programmable RISC cores near banks), Samsung’s
HBM-PIM and SK-Hynix’s AiM (with vector MAC units). The
fundamental compromise is that these compute units are
fabricated in the DRAM manufacturing process, whose
transistors are ~3× slower than logic-process CMOS with far
lower density — so the units are powerful at bandwidth but
weak at arithmetic
an
operator is memory-bound when its speed is limited by moving
data (few arithmetic operations per byte) and compute-bound
when limited by the arithmetic itself (many operations per
byte). Element-wise ops (ReLU, add, normalize) and
matrix-VECTOR products are memory-bound — and PIM’s huge
bandwidth helps them. But matrix-MATRIX multiplication
(GEMM) in a network’s linear layers is compute-bound,
needing 10+ TOP/s where PIM offers ~1, and GEMM is where
modern batched cloud inference spends most of its time.
Memory-bound operators are under 15% of runtime; the
compute-bound majority is exactly what PIM cannot touch
The paper’s move is a co-design inversion. Rather than
accept the workload and struggle to run it on weak hardware,
it changes the workload so the hardware’s strength becomes
the right tool. The vehicle is the
a lookup-table neural network. Its insight:
within a layer, the many input activation sub-vectors are
not all different — they cluster into a small set of typical
patterns called CENTROIDS. So instead of multiplying every
input by the weight matrix, you (1) precompute, ONCE, the
product of each centroid with the weights and store it in a
lookup table; then (2) at inference, match each input
sub-vector to its nearest centroid and simply READ the
precomputed result from the table, accumulating across
sub-vectors. Multiplication — the expensive part — nearly
vanishes: it survives only in the cheap distance
calculations that find the nearest centroid, which are a
tiny fraction of the work. What remains is table reads and
adds: bandwidth-heavy, multiply-light
What people did before, and where it broke
- PIM-enabled systems (offload memory-bound ops, keep the host for GEMM). The deployed norm. Accelerates the under-15% of runtime that is element-wise; the compute-bound GEMM majority still runs on an expensive host, so PIM utilization stays low and the added PIM cost is barely justified.
- GEMV offloading (HBM-PIM, AiM for GPT/LSTM). Genuinely helps single-batch inference, where linear layers are matrix-VECTOR products (memory-bound). But cloud serving batches requests, making the layers matrix-MATRIX (GEMM, compute-bound), and the benefit evaporates.
- Brute-force lightening — binary neural networks on PIM. Cut computation by making weights and activations 1-bit, which fits UPMEM — but the accuracy loss is severe, unacceptable for large models like BERT and ViT.
- Prior LUT-NNs. The right idea, but earlier algorithms could not convert a whole network without wrecking accuracy: one replaced only the final classifier layer; the best prior work could substitute GEMM in just 6 of BERT-base’s 12 layers. Replace ALL layers and accuracy collapses (44% average drop on NLP tasks) — so the computation reduction could never be fully realized.
- The gap. No system converted an entire DNN to LUTs at production accuracy, supported PIM as a real inference backend, AND mapped the resulting kernels efficiently onto commodity PIM hardware. PIM-DL couples all three.
How PIM-DL works
Convert the model, offload the lookups, tune the mapping
-
LUT-NN conversion: GEMM becomes table lookup new
For each linear layer, cluster the input activation sub-vectors (via K-means over calibration data) into a small set of centroids, organized into codebooks; precompute each centroid's inner product with the weight matrix into lookup tables. The weight matrix is gone — replaced by codebooks (to find the nearest centroid) plus tables (to read the answer). Multiplications drop up to 18×, and what survives is 2.9–14.3% of total operations.
-
eLUT-NN calibration: full-network conversion at accuracy new
The hard part is keeping accuracy when EVERY layer is converted. eLUT-NN adds a RECONSTRUCTION LOSS — the sum of per-layer approximation errors, added to the model loss — so training directly drives centroids to represent activations faithfully, avoiding the gradient-vanishing of layer-by-layer methods. A STRAIGHT-THROUGH ESTIMATOR passes gradients through the non-differentiable nearest-centroid step (treating it as identity in the backward pass). Together: all linear layers converted, ~2.3% average accuracy drop, using under 1% of the training data (100× less than prior LUT-NN).
-
Inference engine: split the network across PIM and host new
A frontend of host operators and PIM operators over backend libraries (UPMEM SDK / HBM-PIM / AiM on the PIM side; MKL, cuBLAS on the host side). In a Transformer, the four linear layers (QKV projection, O projection, FFN1, FFN2) become LUTs and go to PIM; attention, which needs real GEMM and cannot be a lookup, stays on the host; element-wise ops (add, norm, GELU) go wherever the target PIM supports them.
-
Two PIM sub-operators: CCS and LUT new
A LUT-NN layer on PIM runs in two steps. CCS (closest-centroid search): each input tile is compared against its codebook by inner product, and the argmin gives the centroid index. LUT: that index selects a column from the precomputed table, and the columns are accumulated into the output. Both are dominated by memory reads and adds — PIM-friendly by construction.
-
Auto-tuner: map kernels onto each PIM platform new
Turning fewer FLOPs into real speedup is non-trivial given PIM's constrained host-PIM and inter-PE communication. The auto-tuner abstracts commodity DRAM-PIMs into a common model, analytically models the dataflow of the CCS and LUT steps, and searches mapping parameters (tile shapes over N and F, LUT-load granularity) for the fastest legal mapping on the given hardware — avoiding illegal mappings and the large gap between best and worst.
-
Tunable accuracy-performance knobs
Two parameters trade accuracy for speed: sub-vector length V (larger V means fewer codebooks and smaller LUTs — faster, up to a point where PIM transfer efficiency drops) and centroid count CT (fewer centroids mean smaller LUT memory footprint — faster). PIM-DL sweeps these to place a workload where it wants on the curve.
The idea
Do not fit the workload to the hardware — fit the hardware by changing the workload. The deployed approach accepts GEMM as immovable and tries to squeeze it onto weak PIM units, which is hopeless: the arithmetic gap (1 vs 10+ TOP/s) is physics, not tuning. PIM-DL’s inversion is to ask what computation PIM is GOOD at — high-bandwidth, low- arithmetic — and transform the network into that shape. LUT- NN is the transformation: it replaces multiplication with memory lookup, precomputing the expensive part offline so inference is reads and adds. The compute-bound operator that PIM could never accelerate is swapped for a memory-bound one it excels at. This is algorithm-hardware co-design in its purest form — not optimizing an algorithm for a target, but choosing a DIFFERENT algorithm whose bottleneck matches the target’s strength, the same manufactured-affinity move the site sees whenever a system reshapes a problem to suit the substrate rather than the reverse.
The algorithm is only useful if the WHOLE network converts. LUT-NN’s promise had always foundered on accuracy: convert a few layers and it works, convert them all and it collapses, because clustering errors compound and layer-by-layer gradient updates vanish through the non-differentiable lookup. eLUT-NN’s two techniques attack exactly this. The reconstruction loss makes the network optimize its own approximation error directly — centroids are trained to represent activations well, with gradients that flow straight to them instead of vanishing. The straight-through estimator lets gradients pass through the nearest-centroid selection (undifferentiable in principle) by pretending it is the identity in the backward pass — a standard trick for training through hard decisions, here enabling end-to-end learning of a fully-converted network. The payoff is the difference between a curiosity (6 of 12 layers) and a system (all layers, ~2.3% drop, under 1% of the data): the hardware idea only pays off if the algorithm can convert enough of the model to matter.
Fewer FLOPs is not free speedup — the mapping must be tuned per platform. LUT-NN reduces operations, but commodity PIM has awkward constraints — limited host-PIM bandwidth, no direct inter-PE communication — so a naive mapping squanders the reduction. The auto-tuner is the bridge from algorithmic advantage to real speedup: it models the dataflow of the two PIM sub-operators (closest-centroid search, table lookup) and searches tile shapes and load granularities for the fastest legal mapping on each device. The paper’s mapping-space illustration shows the gap between best and worst mappings is large and the auto-tuner lands near optimal — a reminder that co-design is three-legged: the algorithm must reduce work, the system must express it, and the mapping must realize it. Drop any leg and the speedup does not appear.
Converting a Transformer’s FFN layer
From GEMM to lookup, and onto the PIM
A feed-forward linear layer of BERT: activations (N×H) times weights (H×F). Numbers from the paper's UPMEM evaluation, V=4, CT=16.
Step 1 of 5
The GEMM that PIM cannot run
The layer multiplies a large activation matrix by the weight matrix — 2·N·H·F operations, half of them multiplications, needing 10+ TOP/s for throughput. On the weak PIM units this crawls; a PIM-enabled system therefore ships this whole layer back to the host CPU or GPU, which is where most of the network's time and cost go.
PIM peak: ~1 TFLOP/s (HBM-PIM/AiM) or 43.8 GOP/s (UPMEM); GEMM wants 10+ TOP/s these linear layers are the GEMM-dominated majority of batched cloud inference
Why it works
- GEMM is compute-bound and PIM units are arithmetically weak → replacing GEMM with table lookup removes the multiplication PIM cannot do (18× fewer FLOPs, multiplies down to 2.9–14.3% of ops) → the compute-bound bottleneck becomes memory-bound work PIM excels at.
- LUT-NN is bandwidth-heavy and multiply-light → its arithmetic intensity (~0.2–0.29 ops/byte) lands squarely in the memory-bound region → PIM’s 8× bandwidth advantage is the exactly-right resource, and the weak compute units suffice.
- Converting every layer normally destroys accuracy → reconstruction loss trains centroids to minimize approximation error directly and the straight-through estimator flows gradients through the hard lookup → full-network conversion holds at ~2.3% drop with under 1% of the data.
- Fewer FLOPs do not auto-translate on constrained PIM hardware → an auto-tuner models the CCS and LUT dataflow and searches the mapping space per platform → the algorithmic reduction becomes real 22–37× speedup instead of being lost to bad mappings.
- Most of a Transformer’s time is in convertible linear layers → offloading those LUTs to PIM and leaving only attention on the host → PIM runs 73–79% of the network, so a wimpy host suffices and the memory-centric system is cost-efficient.
What it costs
- An accuracy tax, small but real. Full LUT conversion drops accuracy ~2.3% on average versus the original model — fine for many uses, but a hard limit for accuracy-critical deployment, and the tunable knobs (V, CT) trade still more accuracy for speed.
- Not everything converts. Attention needs genuine GEMM and cannot become a lookup, so it stays on the host; the method targets the LINEAR layers of Transformers, and models dominated by non-convertible operators benefit less.
- Calibration is required per model. Converting a network needs an eLUT-NN calibration pass (up to 100k iterations on a data sample); cheap relative to training, but not zero-shot — you cannot drop an arbitrary pretrained model onto PIM without it.
- Partly simulated. Real results are on UPMEM PIM-DIMM; the HBM-PIM and AiM numbers (including the GPU-beating ones) come from simulators, so those headline comparisons are projections on hardware not yet shipping.
- A three-legged dependency. The win needs the algorithm (eLUT-NN), the engine, AND the auto-tuner all working; it is more moving parts than a drop-in library, and the benefit is contingent on the mapping being well-tuned for each specific PIM platform.
The evaluation
Full-network LUT conversion on BERT-base/large and ViT-base/huge (GLUE, CIFAR for accuracy) and throughput on BERT and ViT-huge, on real UPMEM PIM-DIMM plus simulated HBM-PIM and AiM, against GEMM-on-PIM, a dual-Xeon CPU (FP32/INT8), and a V100 GPU.
Three ways to run a Transformer's linear layers
| Property | GEMM on the host (PIM-enabled system) | GEMM on PIM (naive) | PIM-DL (LUT-NN on PIM) |
|---|---|---|---|
| Where the compute-heavy layers run | |||
| Arithmetic character | Compute-bound GEMM (wrong fit for PIM) | ||
| Accuracy | Original (no approximation) | Original | |
| Speedup | Reference (host-bound) | Far slower than PIM-DL | |
| Energy & utilization | Poor |
Check yourself
Check your understanding
1. Commodity DRAM-PIM offers 8× the memory bandwidth, yet it can only accelerate under 15% of a neural network's runtime. Predict the fundamental reason, and why PIM-DL's response is to change the WORKLOAD rather than improve the hardware.
2. LUT-NN replaces matrix multiplication with table lookups. Predict how this eliminates most multiplication, and why the result is a good fit for PIM specifically.
3. Prior LUT-NN work could convert only about half of a network's layers before accuracy collapsed. eLUT-NN converts all of them at ~2.3% drop. Predict why full conversion is the make-or-break requirement, and what makes it possible.
4. PIM-DL reduces FLOPs by up to 18× via LUT-NN, but the paper stresses that an auto-tuner is essential to turn that reduction into real speedup. Predict why fewer operations do not automatically mean faster execution on PIM.