[ARCH-PROP] Elastic Memory

[yunchong-song]

Architecture Name

Elastic Memory

Parent issue

#1

Motivations

Long-context attention has compute and memory costs that grow quadratically with sequence length. Existing recurrent-memory approaches either couple capacity to model dimensions or rely on heuristic caching, retrieval, or learned summaries.

We hypothesize that pre-RoPE K/V representations at a deep attention layer contain sufficiently smooth, low-frequency structure for compression using HiPPO-LegS under its weighted $L_2$ polynomial-projection objective. Polynomial reconstruction may then improve long-context modeling without additional trainable parameters, while the HiPPO order $N$ independently controls memory capacity.

Proposed Architecture

Keep the parent OLMo3 tokenizer, embeddings, layer count, width, normalization, MLPs, QKV/output projections, and other attention layers unchanged. Replace one designated deep attention layer with Elastic Memory:

• Divide each sequence into blocks of length $L=2048$. Compute Q/K/V normally for the current block. Apply RoPE only to the Q/K used by current attention; compress the pre-RoPE K and V representations with HiPPO.
• Maintain states $C^{(K)}\in\mathbb{R}^{N\times D_K}$ and $C^{(V)}\in\mathbb{R}^{N\times D_V}$. For each block, update them using $C_i=P_iC_{i-1}+\bar{K}_iF_i$, where $P_i$ and $\bar{K}_i$ come from ZOH-discretized HiPPO-LegS dynamics.
• Before processing block $i$, reconstruct historical representations using $K_{\mathrm{mem}}=R_{i-1}C^{(K)}{i-1}$ and $V{\mathrm{mem}}=R_{i-1}C^{(V)}_{i-1}$.
• Let $R_i$ support uniform and exponential sampling. Reconstructed memory tokens receive no additional positional encoding.
• Prepend the reconstructed K/V to the current K/V and apply a trapezoidal causal mask: every current query can access all memory tokens, while attention within the current block remains causal.
• Precompute $P_i$, $\bar{K}_i$, and $R_i$ up to a configured maximum context length. Explicitly handle the $k=0$ ZOH boundary instead of directly evaluating the singular expression $\log((k+1)/k)$.
• Add no trainable parameters. Recurrent-state storage is $O(N(D_K+D_V))$, while precomputed-bank storage, update/reconstruction compute, and attention width grow with $N$, $L$, and $L_{\mathrm{mem}}$.

Preliminary Results (if any)

Paper: Towards Compressive and Scalable Recurrent Memory

We evaluate custom Llama 3 implementations with 100M–400M parameters, 32,768-token contexts, 2,048-token blocks, and 40B training tokens per model:

• At 100M parameters, exponential sampling achieves the best PPL on PG-19, Proof-Pile, and FineWeb-Edu, while uniform sampling achieves the best LongPPL, with no additional trainable parameters.
• On Proof-Pile, $1\times$ Elastic Memory obtains PPL 2.92, outperforming a $16\times$ Memorizing Transformer at 2.94.
• As $N$ increases from 540 to 8,640, reconstruction MSE decreases from 0.45 to 0.27 and PPL decreases from 2.92 to 2.75.
• At 400M parameters, exponential sampling obtains PPL 2.737 versus 2.771 for Melodi. Its normalized training throughput is 0.557 versus 0.366 for Melodi.

Experiments Plan

• Baselines: OLMo3; direct token sampling with equal state capacity; and, where implementations are available, Memorizing Transformer, Infini-Transformer, and Melodi.
• Scale and protocol: Screen configurations on Proof-Pile using the smallest supported OLMo3 model, 5B tokens, and three random seeds. Validate successful configurations at the next model scale with 40B training tokens. Keep data order, optimizer, token budget, hardware, and software stack identical across comparisons.
• Data: PG-19, Proof-Pile, and FineWeb-Edu. Retain documents longer than 32,768 tokens and split them into non-overlapping 32,768-token samples. Add a controlled long-range retrieval task and a long-document QA task.
• Metrics: Sliding-window PPL, LongPPL, retrieval/QA accuracy, K/V reconstruction MSE, training and inference throughput, peak memory, precomputed-bank size, and trainable parameter count.
• Ablations: Uniform versus exponential sampling; train-time versus test-time sampling switches; $N\in{540,1080,2160}$, followed by $N\in{4320,8640}$ if successful; HiPPO versus equal-cost direct sampling; layer 9 versus an OLMo3 layer at approximately two-thirds depth; and explicit sweeps for $L_{\mathrm{mem}}$ and $\alpha$, frozen before test evaluation.