Yunchong Song, Jushi Kai, Liming Lu, Kaixi Qiu, Zhouhan Lin
9 min
Abstract
Transformers face a quadratic bottleneck in attention when scaling to long contexts. Recent approaches introduce recurrent memory to extend context beyond the current window, yet these often face a fundamental trade-off between theoretical principles and practical scalability. To address this, we introduce Elastic Memory, a novel memory architecture grounded in the HiPPO framework for online function approximation. Elastic Memory treats historical sequence as samples from continuous signals, applying optimal online compression to encode them into a fixed-size memory state. For retrieval, we propose a flexible \textit{polynomial sampling} mechanism that reconstructs a history summary from this compressed state. Elastic Memory consistently outperformed baselines on long-context (32k+) datasets across three domains. With equal parameters, it beat Memorizing Transformer by 16x memory and outperformed Melodi at all memory sizes, even when Melodi had 30% more parameters. When scaling model size, Elastic Memory stayed ahead of all baselines and was significantly faster than Melodi at 4x size. Furthermore, its decoupled design allows for injecting inductive biases at test-time to boost performance.
Alex: A flow rule... you mean something continuous, like water level rising gradually?
Sam: Yes. To handle word-by-word text, they snapshot this flow into discrete steps with a stable trick called Zero-Order Hold, turning endless flow into quick jumps forward. For speed in batches, they roll it into block updates—precomputing blends across chunks of text in parallel. This scales to huge lengths, like 32,000 words, using about one-sixteenth the memory of similar systems.
Alex: Oh—so the update stays cheap, constant time per word, because it's just multiplying by those preset tables. But with text coming in chunks, how does it actually blend the old summary with a new chunk without redoing everything?
Sam: They group words into fixed-size blocks. For each block, there's a preset multiplier that carries the summary forward across the whole chunk, plus another tool that mixes in the new words evenly. Those tools—a state transition matrix and a HiPPO kernel—are calculated ahead for every possible block spot, so during use it's just quick multiplies. Precomputing them turns a step-by-step crawl into parallel blasts on hardware like GPUs.
Alex: Preset multipliers... so no on-the-spot math, just grab and go? Now for pulling info back out: how does that work?
Sam: The summary holds curve coefficients, and to rebuild past highlights, it samples points along that curve using another preset grid of weights—a reconstruction matrix. Pick spots evenly for full-history view, or cluster near the end for recent focus; that's polynomial sampling. They prepend those sampled keys and values to the current block's. Then a mask shapes attention like a trapezoid: every spot in the new block sees all samples fully, but only past spots within the block itself—no future leaks. This keeps causality while injecting long-range recall efficiently.
Alex: Trapezoid mask... makes sense for that balance. The paper notes trade-offs in sampling, like uniform versus recent-heavy?
Sam: Yes, uniform treats history evenly, exponential favors now—which fits language patterns. Both outperform baselines on long sequences with less memory.
Alex: Solid path for scaling... but what does that look like in actual tests? Like, on real long texts?
Sam: The paper tests it on long-document language modeling—predicting words in huge chunks of text, over 32,000 tokens each. They use three datasets: books from PG-19, math papers called Proof-Pile, and educational web pages from FineWeb-Edu. Metrics are perplexity—how puzzled the model gets predicting words, lower is sharper—and LongPPL, which zooms on spots needing far-back recall.
Alex: Okay, so controlled long chunks from varied sources. What do they measure against?
Sam: They compare to other memory systems: a queue of past info like a waiting line; associative links; block summaries; plus a plain Transformer baseline. All trained identically from scratch on 40 billion tokens.
Alex: Fair across the board. So the results?
Sam: Elastic Memory's uniform and exponential versions both score lower perplexity than baselines on most datasets—meaning clearer predictions. Notably, it beats the prior best on LongPPL for books and web texts, using the same memory size but zero extra parameters.
Alex: Zero extra parameters... so same model size, better long-range grasp. Does that hold across domains, like math too?
Sam: Yes, competitive on Proof-Pile too, where math needs precise distant links. The paper attributes it to the principled compression avoiding guesswork in rivals, plus efficient updates letting training stabilize.
Alex: Huh, principled over heuristic pays off there. Pushing contexts further... but does that hold up when they make the memory bigger or the whole model larger?
Sam: The paper checks that directly on math texts. They scale memory up to sixteen times the base size—and train everything from scratch. Elastic Memory's exponential version gets the lowest overall prediction sharpness across all sizes, while uniform does best on long-range spots. Notably, its base size matches or beats a sixteen-times larger rival, using far fewer extra parts to train. Benefits persist as the base model grows too, outpacing rivals consistently.
Alex: So same small memory, but sharper predictions than rivals bloated up sixteen-fold. Why does that happen mechanically?
Sam: It stems from the compression packing history tightly without waste or extra trainable pieces—rivals add millions of those as memory grows, diluting focus. Here, scaling just extends the fixed math rules, keeping updates efficient and stable.
Alex: Ah, no dilution from extras. Any other flexibility they test?
Sam: Yes—they train with one sampling style, say uniform, then switch to exponential at test time without retraining. It's like having a stored summary that rebuilds flexibly for different needs. Exponential tweaks often sharpen average predictions, uniform distant ones, across datasets and scales—showing the memory state decouples from retrieval.
Alex: Decoupled state... so tweak retrieval post-training for the task. That's a meaningful edge for real use. But to confirm it's really tapping that long history, not just leaning on recent words, what else did they check?
Sam: They ran a test where they scrambled parts of the recent words—replacing them with random junk—to force reliance on older info. Elastic Memory stayed sharper overall, especially under heavy noise, suggesting the stored history acts as a solid backup when nearby details fail. They also measured training speed as data processed per second: it matches the plain Transformer closely, holding competitive throughput even as memory or model scales up sixteen times.
Alex: Competitive even scaled up... no big drags. Any catches?
Sam: The precomputes tie to a fixed top context length, so adapting to wildly varying ones needs tweaks. Sampling choices trade recent focus for even history coverage, though both beat baselines. The paper notes these as areas for refinement.
Alex: Fair trade-offs, grounded in the tests. Overall, it shows a reliable way to extend context with less memory and solid speed.
Sam: Exactly—this design resolves the clash between theory-backed compression and scalable updates, hitting top marks on long texts across books, math, and web data. Context scrambling confirms real long-range use, and speeds support real-world training. It's a meaningful advance for handling massive inputs efficiently.
Alex: Well put. That's our look at compressive memory for long-context models. Thanks for joining ResearchPod.