Engram: O(1) Hash-Based Associative Memory via Locality-Sensitive Hashing.

Engram implements a fast key-value memory that stores and retrieves values in amortised O(1) time using Locality-Sensitive Hashing (LSH). Multiple independent hash tables reduce collision probability, and exponential moving-average (EMA) writes allow smooth interpolation of stored values.

Motivation

Classical associative memories (NTM, MemoryNetwork) perform O(N) attention over all N memory slots per query. Engram instead hashes each query into a small number of buckets, making both reads and writes O(hash_bits × key_dim) independent of the number of stored memories.

Architecture

Query key [batch, key_dim]
     |
     v
+---------------------------+
| LSH Hash (per table t):   |
|   project = W_t @ key     |  W_t ∈ R^{hash_bits × key_dim}
|   bits = sign(project)    |  {0,1}^hash_bits
|   bucket = bits → int     |  0..num_buckets-1
+---------------------------+
     |
     v  (for each of num_tables tables)
+---------------------------+
| Memory Slots              |
|   [num_tables, num_buckets, value_dim]
+---------------------------+
     |
     v
Retrieve slot per table → average across tables
     |
     v
Retrieved value [batch, value_dim]

Hashing

For num_buckets = 256 = 2^8:

• hash_bits = 8 random projection vectors per table
• W shape: [num_tables, hash_bits, key_dim]
• bucket = sum(sign(W @ key) >= 0) * [1, 2, 4, ..., 2^{hash_bits-1}]

Write (EMA update)

memory[t, hash(key)] ← decay × memory[t, hash(key)] + (1 − decay) × value

Usage

# Create memory state
mem = Engram.new(key_dim: 32, value_dim: 64)

# Read
result = Engram.engram_read(mem, query)   # [batch, value_dim]

# Write (returns updated memory)
mem = Engram.engram_write(mem, key, value, decay: 0.99)

# Build an Axon model for differentiable reads
model = Engram.build(key_dim: 32, value_dim: 64)

References

• Andoni & Indyk, "Near-Optimal Hashing Algorithms for Approximate Nearest Neighbor in High Dimensions" (FOCS 2006)
• Locality-Sensitive Hashing for associative memory is explored in "Reformer: The Efficient Transformer" (Kitaev et al., 2020)