Autoencoders as Associative Memory

No WebGPU detected. This demo runs faster with WebGPU enabled.

An overparameterized autoencoder, trained only to reconstruct its inputs, does something interesting when you iterate it: feed it a novel image and loop the output back as input, and it converges to the nearest training example. Reconstruction becomes retrieval. Radhakrishnan et al. (2020) showed that these networks develop stable attractors at their training points, despite never being trained for convergence.

The visualization below trains an overparameterized autoencoder (~466k parameters) on just 16 MNIST digits with heavy augmentation, implemented in jax-js in the browser.

graph LR
    X[Input 28×28] --> C[Conv2D 7×7
stride 7] --> F[Flatten
256 dims]
    F --> E[GLU Block
+ LayerNorm]
    E --> Z[Latent z
256 dims]
    Z --> D[GLU Block
+ LayerNorm]
    D --> P[Linear
256→784]
    P --> Y[Output 28×28]
    Y -.Iterate.-> X

Training loss decreases over 1k steps as the network learns to reconstruct MNIST digits

Weight norm

Gradient norm

The attractor doesn't exist statically throughout training. It forms gradually and continues to shift as the network converges. The landscape stabilizes as we drop the learning rate throughout training. This temporal evolution shows the network developing a shifting landscape where which inputs fall into which basins changes throughout training. You can almost see catastrophic forgetting playing out in real time.

Latent space (256-d) projected onto three orthogonal planes. Colored dots: training attractors. Paths: test images converging through iteration.

Distance from each iteration to the nearest attractor (first test image)

Without heavy augmentation, intermediate outputs are out of distribution, the landscape is jagged, and the network develops spurious attractors. Basically fixed points that aren't our training examples as represented by the decoder. Smoothing the input distribution through augmentation creates cleaner basins around the real data.

Test 1

Rollout

Attractor

Test 2

Rollout

Attractor

Test 3

Rollout

Attractor

Test 4

Rollout

Attractor

References

Radhakrishnan, A., Belkin, M. & Uhler, C. Memorization in overparameterized autoencoders. arXiv (2018).
Radhakrishnan, A., Belkin, M. & Uhler, C. Overparameterized neural networks implement associative memory. PNAS 117(44), 27162–27170 (2020).