Source-linked AI summary
Robust High-dimensional Memory-augmented Neural Networks
Geethan Karunaratne, Manuel Schmuck, Manuel Le Gallo, Giovanni Cherubini, Luca Benini, Abu Sebastian, Abbas Rahimi
TL;DR
Soft access to explicit memory can burden conventional hardware, while real-valued key vectors complicate in-memory cosine similarity. The paper trains a controller to use robust HD vectors with computational memory, and reports large-scale PCM experiments for few-shot Omniglot classification.
Problem
Soft reads and writes involve every memory entry, creating a bottleneck for explicit-memory access on conventional von Neumann hardware.
Source: §I · p. 1Method
The approach trains a controller to produce robust HD representations for an analog computational memory, using binary or bipolar vectors and dimensionality-preserving transformations.
Source: §I · p. 2, §III · pp. 7–8Results
The architecture was demonstrated on 100-way 5-shot Omniglot using 256,000 PCM devices, and spatial variability caused a further 0.11% classification-accuracy loss versus PCM model simulations.
Source: §I · p. 2, Supplementary Note 6: Spatial… · p. 38Takeaways & Limitations
The work combines deep neural-network representations with HD computing for analog in-memory computation using robust bipolar or binary vectors.
Source: §III · pp. 7–8
Abstract
from arXiv · showhide
Traditional neural networks require enormous amounts of data to build their complex mappings during a slow training procedure that hinders their abilities for relearning and adapting to new data. Memory-augmented neural networks enhance neural networks with an explicit memory to overcome these issues. Access to this explicit memory, however, occurs via soft read and write operations involving every individual memory entry, resulting in a bottleneck when implemented using the conventional von Neumann computer architecture. To overcome this bottleneck, we propose a robust architecture that employs a computational memory unit as the explicit memory performing analog in-memory computation on high-dimensional (HD) vectors, while closely matching 32-bit software-equivalent accuracy. This is achieved by a content-based attention mechanism that represents unrelated items in the computational memory with uncorrelated HD vectors, whose real-valued components can be readily approximated by binary, or bipolar components. Experimental results demonstrate the efficacy of our approach on few-shot image classification tasks on the Omniglot dataset using more than 256,000 phase-change memory devices. Our approach effectively merges the richness of deep neural network representations with HD computing that paves the way for robust vector-symbolic manipulations applicable in reasoning, fusion, and compression.
I. INTRODUCTION
MANNs move storage into explicit memory, but soft access over every entry burdens conventional hardware. The proposed CNN–HD-memory architecture uses softabs attention to form robust vectors and improves few-shot accuracy over softmax.
I. INTRODUCTION
- MANNs separate processing from storage, using explicit key-value memory to offload new information and support few-shot learning, subject to memory capacity.The key memory stores and compares patterns; the value memory stores labels.Source: §I · p. 1
- Soft reads and writes compare each query with every key entry, creating a memory-intensive bottleneck on conventional von Neumann hardware.TCAM-based approaches also constrain search precision and supported similarity metrics, and typically return only the first nearest neighbor.Source: §I · pp. 1–2
- The algorithmic-hardware codesign trains a deep neural network controller to generate HD vectors and implements a bipolar or binary key memory for few-shot inference.The reported mixed hardware/software experiments use 256,000 PCM devices for 100-way 5-shot Omniglot classification.Source: §I · p. 2
A. Proposed MANN architecture using in-memory computing
- The CNN maps unrelated images to quasi-orthogonal HD vectors, while differentiable memory-interface functions and episodic training let classification errors update the controller.The approach replaces designing a task-specific encoder with end-to-end training of the controller.Source: §I · p. 2, A. Proposed MANN architecture using… · p. 3
B. A new attention mechanism appropriate for HD geometry
- Softabs sharpening is designed to direct unrelated classes toward uncorrelated vectors, unlike softmax, which can drive them toward anticorrelation during training.HD vectors are quasi-orthogonal with high probability, supporting robust representation of unrelated items.Source: §I · p. 2, B. A new attention mechanism… · p. 4
- Up to 5.0%, 9.6%, and 19.6% higher accuracy than softmax is reported for 5-way 1-shot, 20-way 5-shot, and 100-way 5-shot, respectively.The softabs-trained representations retain robustness when real-valued HD vectors are transformed to bipolar.Source: B. A new attention mechanism… · p. 4
C. Bipolar key memory: transforming real-valued HD vectors to bipolar
Real-valued key vectors demand analog storage and require extra computation for cosine similarity. The proposed dimensionality-preserving transformations produce bipolar vectors whose similarities can be computed as a dot product.
C. Bipolar key memory: transforming real-valued HD vectors to bipolar
- Real-valued key vectors require analog storage and cosine-similarity computations that need norms and vector-dependent scaling, complicating crossbar implementation.The absolute-value norm approximation still requires additional circuitry for scaling.Source: C. Bipolar key memory: transforming… · p. 4
- The proposed dimensionality-preserving transformations convert real-valued vectors directly into dense bipolar or binary vectors, without additional quantization, mapping, or coding schemes.The described transition transforms real-valued vectors to bipolar, then can transform bipolar vectors to binary.Source: C. Bipolar key memory: transforming… · pp. 4–5
- At inference, the sign function clips controller outputs into bipolar support and query vectors, and their constant norm makes cosine similarity a scaled dot product.This removes cosine-similarity normalization and attention-vector normalization during inference.Source: C. Bipolar key memory: transforming… · p. 5
D. Binary key memory: transforming bipolar HD vectors to binary
Binary and bipolar HD representations let the proposed memory architecture perform similarity search on PCM hardware with accuracy close to 32-bit software. The section also describes the binary transformation, robustness to device variability, and potential uses of HD key-value memories beyond classification.
D. Binary key memory: transforming bipolar HD vectors to binary
- The bipolar-to-binary transformation approximately preserves cosine similarity, and regularizing the controller reduced fixed-norm deviation to 0.91% and improved accuracy by as much as 0.74%.Without the regularizer, transformed vectors at d = 512 deviated 2.08% from the fixed norm.Source: D. Binary key memory: transforming… · pp. 5–6
E. Experimental results
- Moving from real to bipolar representations caused at most a 0.45% accuracy drop; binary caused a further 0.11% drop with cosine similarity, versus as much as 1.13% with dot product.The software inference used a regular absolute function in place of softabs sharpening; this was its only approximation.Source: E. Experimental results · p. 6
- PCM experiments matched the PCM model within ±0.2%; binary accuracy dropped at most 1.12% from the ideal crossbar in 100-way 5-shot, compared with 0.41% for bipolar.For smaller problems, accuracy differences were within 0.58%.Source: E. Experimental results · p. 6
- The binary architecture achieved the highest accuracy across the compared multibit CAM works, while PCM search energy was 25.9 fJ per bit versus 5 pJ in 65 nm digital.The compared CAM works cannot support cosine similarity, according to the passage.Source: E. Experimental results · p. 7
III. DISCUSSION
- Noisy cosine-similarity standard deviation scales inversely with the square root of vector dimensionality; at d = 512, σrel = 31.7%, and α = 0.5, σ(Λ) was ≈0.015.The discussion attributes tolerable similarity deviations at higher dimensions to this scaling.Source: §III · p. 7
- The learned HD key-value patterns are proposed as representations for reasoning or multimodal fusion, and as prototype stores for unsupervised and contrastive learning.The discussion also identifies compression and reasoning as potential applications.Source: §III · pp. 7–8
- For the largest Omniglot problem tried, binary key memory on 256,000 noisy PCM devices had less than 2.7% accuracy drop versus 32-bit real-valued software vectors (94.53% vs. 91.83%); bipolar key memory lost less than 1%.Source: §III · p. 8
METHODS
The Omniglot evaluation uses episodic few-shot classification with augmentation and validation-based checkpoint selection. A CNN maps grayscale images to feature vectors used by the memory architecture.
Omniglot dataset: evaluation and symbol augmentation
- Episodes form support sets from randomly selected classes and shots, then classify query batches sampled from those same classes.Accuracy is the proportion of correctly classified queries, averaged across episodes.Source: Omniglot dataset: evaluation and… · p. 9
- Testing accuracy is averaged over 1,000 testing episodes per run, with the checkpoint chosen for highest validation accuracy; results can be averaged across runs.Training runs use up to 50,000 episodes, with validation every 500 episodes on 250 validation episodes.Source: Omniglot dataset: evaluation and… · p. 9
- Training and inference use randomly shifted and rotated 32 × 32 Omniglot images to augment support sets and query batches.The supplied passages specify shift and rotation standard deviations of 2.5 pixels and π/12, respectively.Source: Omniglot dataset: evaluation and… · p. 9
The CNN as a controller for the MANN architecture
- A CNN embedding maps grayscale 32 × 32 images into d-dimensional feature vectors, with random shifting and rotation before mapping.The last dense layer sets feature dimensionality; its output is activated with sign or step functions during inference, depending on the representation.Source: The CNN as a controller for the MANN… · pp. 9–10
- The CNN uses four convolutional layers, two 2 × 2 stride-2 max-pooling layers, and a fully connected layer with d units.The first two convolutional layers each have 128 filters of shape 5 × 5; the next two each have 128 filters of shape 3 × 3.Source: The CNN as a controller for the MANN… · pp. 9–10
Details of attention mechanism for the key-value memory
The key-value memory compares each query with stored support vectors, sharpens and normalizes similarities into attention weights, and reads out labels. PCM simulations model measured device variability and conductance drift.
Details of attention mechanism for the key-value memory
- A query is compared with stored support vectors, and normalized attention weights combine their one-hot labels into a probability distribution for prediction.The authors report that ranking by summing probabilities within each class improves PCM inference accuracy over global-argmax.Source: Details of attention mechanism for the… · p. 10
- The attention mechanism applies a similarity metric and sharpening function to query-key comparisons before normalization.Cosine similarity and an exponential sharpening function are given as examples.Source: Details of attention mechanism for the… · p. 10
PCM model and simulations
- PCM conductance is modeled with Gaussian variations and drift, with device parameters fitted from measurements of 10,000 SET devices over five orders of magnitude.The fitted means are 22.8 µS for G0 and 0.0715 for ν; programming, drift, and read-out variability are also estimated from measurements.Source: PCM model and simulations · p. 11
- Simulations process query batches together because the estimated 1 µs between successive queries is negligible relative to the estimated 20 s before the first query after programming.The batch simulation computes currents from query read-out voltages and conductance at the evaluation time.Source: PCM model and simulations · p. 11
Experimental details
The experimental setup connects a host-controlled platform to a PCM chip and emulates support-vector storage and query reads. The chip uses a 512-word-line by 2,048-bit-line array.
- The PCM chip contains doped-Ge2Sb2Te2 cells integrated in 90 nm CMOS, with cell-addressing circuitry, an 8-bit ADC, and voltage- or current-mode programming.The host Matlab environment coordinates experiments through Ethernet with two FPGAs and an analog front end.Source: Experimental details · p. 11
- The array has 512 word lines and 2,048 bit lines; cell reads use a 300 mV bias and take 1 µs.SET programming yields an average conductance of 22.8 × 10^-6 S with 31.7% variability, while RESET yields 0 S conductance.Source: Experimental details · pp. 11–12
- For Omniglot tests, support vectors are programmed onto PCM devices and query vectors are applied as read-out voltages; attention values are reduced digitally after sequential device reads.Bipolar support vectors use adjacent bit-line pairs to store positive and negative components.Source: Experimental details · p. 12
Energy estimation and comparison
The section estimates search energy for PCM and a throughput-matched digital CMOS baseline, and describes robustness comparisons of bipolar and binary key memories. The supplementary comparison varies device conductance variation across two few-shot tasks.
Energy estimation and comparison
- 5 pJ per bit is the estimated search energy for the digital CMOS baseline, while the PCM binary key memory uses 2.5 fJ per device plus 23.4 fJ for analog/digital peripheral circuits.The CMOS baseline is synthesized for throughput equivalent to the PCM-based counterpart.Source: Energy estimation and comparison · p. 12
SUPPLEMENTARY FIGURES
- The bipolar PCM architecture stores complementary support-vector versions on twice as many crossbar columns as the binary architecture.The bipolar design also uses a sign activation and approximates softabs sharpening with an absolute-value function.Source: Supplementary Figure 1: Architecture… · p. 20
- The robustness comparison varies device conductance variation for 5-way 1-shot and 100-way 5-shot tasks, comparing bipolar d = 512 with binary d = 512 and d = 1024.At zero or low variation, binary d = 1024 outperforms binary d = 512; the passage attributes this to information loss from reducing representation dimensionality.Source: Figure 2 · p. 21
- The bipolar architecture is reported as more robust to in-memory-computing non-idealities than binary architectures under the same power and area constraint.The accuracy distributions use eight test runs of 1,000 episodes each, with error bars of one standard deviation.Source: Supplementary Figure 2: Robustness of… · p. 21, Figure 2 · p. 21
SUPPLEMENTARY TABLES
The controller’s embeddings approach the high-dimensional occupancy target at d ≥ 512, and a regularizer further reduces occupancy-ratio deviation while improving classification accuracy. The supplementary tables cover PCM parameters and accuracy comparisons with and without the regularizer.
Supplementary Note 1: The CNN Controller in Conformity with Dense High-dimensional Representations
- At d ≥ 512, the controller’s vectors closely follow the equiprobable property; d = 512 was selected for PCM resiliency and accuracy comparable to real-valued vectors.At d = 512, the unregularized controller’s occupancy-ratio standard deviation was 2.08%, corresponding to a 1.02% accuracy drop for binary 100-way 5-shot classification.Source: Supplementary Note 1: The CNN… · p. 25
Introducing a regularizer to optimize the occupancy ratio
- 2.08% to 0.91% standard deviation from the target occupancy ratio at d = 512; accuracy increased across all three problems, up to 0.74%.The regularized binary dot-product architecture was at most 0.28% lower in accuracy than cosine similarity without regularization.Source: Introducing a regularizer to optimize… · p. 26, Table II · p. 28
SUPPLEMENTARY TABLES
- The supplementary tables list PCM model parameters, their derived values, and average classification accuracy with versus without the regularizer.The accuracy comparison covers 10,000 episodes across 10 runs; the captions do not provide cell values.Source: Table 1 · p. 24, Table I · p. 24, Table II · p. 28
Supplementary Note 2: Training and Inference for Key-Value Memory Network
Training matures the controller through repeated few-shot episodes, while the key-value memories remain fixed during each episode’s query evaluation and backpropagation. Validation and early stopping help select a controller that generalizes to unseen classes.
Supplementary Note 2: Training and Inference for Key-Value Memory Network
- The supplementary table summarizes the learning and inference phases, which the note describes in detail.Learning begins with random initialization of the controller’s embedding parameters; the mature controller is selected based on validation performance.Source: Supplementary Note 2: Training and… · p. 29, Table III · p. 29, §1 · pp. 29–30
1. Learning Phase
- Training episodes repeatedly load support sets, evaluate query batches, and update only the controller by backpropagating supervised log loss while memories remain fixed.Episodes use randomly drawn training classes; repeating them matures the embedding function, and early stopping can use a validation set with non-overlapping classes.Source: §1 · pp. 29–30
2. Inference Phase
At inference, the mature controller stays frozen while support examples program key-value memory for a new few-shot problem and queries are classified against those stored examples. The sharpening-function analysis motivates softabs and compares its properties with softmax.
2. Inference Phase
- During inference, the frozen mature controller encodes support examples into key-value memory, then encodes queries for classification against the stored support vectors.The loaded support examples represent classes never exposed during learning; inference skips the backpropagation step.Source: §2 · p. 30
Supplementary Note 3: Proof of the optimality of the sharpening function
- The sharpening analysis favors zero cosine similarity between dissimilar classes, with monotonic behavior on either side of zero and a zero derivative at zero.The text connects orthogonal high-dimensional vectors to accommodating more classes at a given embedding dimensionality.Source: Supplementary Note 3: Proof of the… · p. 33
- For β = 10, softabs has ε(0) = 0.0134, a 27× better approximation than softmax, and meets the stated optimal conditions more closely.Softabs has inflection points at -0.5 and 0.5; softmax has none and fails to meet the conditions cited in the proof.Source: Supplementary Note 3: Proof of the… · p. 33
Supplementary Note 4: Comparison of classification accuracy between sum-argmax vs argmax as the ranking criteria
Sum-argmax reduces the accuracy drop caused by PCM noise relative to global argmax in nearly all tested problems. Its class-wise aggregation can average out variation among programmed classes.
Supplementary Note 4: Comparison of classification accuracy between sum-argmax vs argmax as the ranking criteria
- Sum-argmax sums attention components within each class before choosing the largest class total, while global argmax selects the highest individual support-vector probability.For an m-way n-shot problem, sum-argmax uses O(m) comparisons and O(mlog(n)) additions; global argmax uses O(nm) comparisons.Source: Supplementary Note 4: Comparison of… · p. 35
- In all but bipolar 20-way 5-shot, sum-argmax yields lower accuracy drop than global argmax; binary 100-way 5-shot gains up to 1% better accuracy mitigation.The figure compares accuracy drop from the ideal crossbar to the PCM experiment across problems.Source: Supplementary Note 4: Comparison of… · p. 35, Figure 11 · p. 35
Supplementary Note 5: Effect of approximating sharpening function and similarity function
Approximation effects depend on the vector representation: regular absolute sharpening simplifies binary inference, while replacing cosine similarity with a dot product penalizes binary vectors.
Effect of approximating sharpening function
- The regular absolute sharpening function simplifies inference and slightly improves accuracy for binary representations.It can be bypassed for binary representations.Source: Supplementary Note 5: Effect of… · p. 36, Figure 12 · p. 36
- Softabs consistently outperforms regular absolute sharpening for real and bipolar representations, but reduces accuracy for binary representations.The comparison covers 5-way 1-shot, 20-way 5-shot, and 100-way 5-shot problems.Source: Supplementary Note 5: Effect of… · p. 36, Figure 12 · p. 36
Effect of approximating similarity function
- Approximate dot-product similarity significantly reduces binary accuracy, while binary vectors outperform bipolar vectors when precise cosine similarity is used.The reported results attribute the binary accuracy drop to similarity approximation; bipolar vectors have constant norm.Source: Supplementary Note 5: Effect of… · pp. 36–37, Figure 13 · p. 36
Supplementary Note 6: Spatial Variability on PCM Crossbar
PCM spatial variability is associated with a small accuracy loss, while state-of-the-art comparisons report strong results for the proposed model in low-precision and noisy-memory settings, with a qualification for software comparisons.
Supplementary Note 6: Spatial Variability on PCM Crossbar
- For 100-way 5-shot, spatial SET-conductance variability is 31.7% across the chip and 5.38% across bitline averages, and is associated with a 0.11% accuracy loss.The loss compares PCM experiments with PCM-model simulations using the same parameters.Source: Supplementary Note 6: Spatial… · p. 38
Supplementary Note 7: Accuracy comparison with the state-of-the-art few-shot learning models
- Matching Networks achieve up to 0.49% higher accuracy than the model for 20-way 5-shot, but use an unconventional split with 4,720 additional training examples.For 100-way 5-shot, the model’s accuracy remains within the range of smaller-way problems; no other models are reported at that size.Source: Supplementary Note 7: Accuracy… · p. 39
- Using 1-bit key-memory vectors, the model with softabs achieves the highest accuracy across all problems among the compared ideal low-precision models.The comparison models use 3-bit or 4-bit representations.Source: Supplementary Note 7: Accuracy… · p. 39
- With 30% conductance variation, the model achieves higher accuracy than the compared multibit CAM model across all problems in noisy, low-precision simulations.The comparison uses 30% variation for the proposed model and the multibit CAM model’s lowest state.Source: Supplementary Note 7: Accuracy… · p. 39
- In 100-way 5-shot, software-to-noisy-hardware accuracy drops 4.09% with softmax and 2.46% with softabs.Across comparisons, the softmax-to-softabs accuracy gap reaches 19.26% in the software model.Source: Supplementary Note 7: Accuracy… · p. 39
Supplementary Note 8: Proof of robustness of noisy cosine similarity with binary high-dimensional vectors
The analysis relates noisy cosine similarity to vector dimensionality and SET-state variability. It also shows that non-unit mean conductance scales the expected similarity.
Noise scaling
- The noisy cosine-similarity standard deviation decreases inversely with the square root of vector dimensionality.This scaling diminishes the influence of SET-state variability.Source: Supplementary Note 8: Proof of… · p. 41
Expected similarity
- When the mean conductance factor is ε, the expected noisy similarity is εα, where α is the noise-free cosine similarity.A deviation of the expected conductance factor from 1 directly affects the expected noisy similarity.Source: Supplementary Note 8: Proof of… · p. 41
Conductance condition
- The crossbar devices should have the same mean SET-state conductance; otherwise, compensation depends on representation quality.Source: Supplementary Note 8: Proof of… · p. 41