SpMem: Sparse Attention via Memory Tokens

Most sparse attention methods store the entire key–value (KV) cache and are sparse only in how they read it. A selector must keep every entry in case a later query wants it, so the cache still grows with the context. These methods also pay for sparsity with machinery outside attention itself: a KV block compressor, a top-k selector, and a training recipe that departs from the base model's. We present SpMem, a sparse attention that needs neither a compressor nor a selector and keeps the KV cache compressed throughout inference. On the global-attention layers, SpMem appends m memory tokens after every C real input tokens. Memory tokens are processed through a transformer forward pass, generating KV features that carry the core information of the prior input chunks. To this end, we train a model using next-token prediction with modified attention masking. During inference, real and memory tokens share one fixed retention rule, Tapered Window Attention (TWA), which keeps the most recent W tokens and halves the number retained with each further window of age. At 128K context a SpMem layer holds 16K entries against dense attention's 128K, and that number does not grow with context. On Gemma 4 12B and 31B, against a dense model trained on the same data, SpMem stays close on knowledge and long-context reasoning and the cost falls mostly on retrieval. Keeping one of the eight global layers at full attention recovers most of that retrieval with a 4.3× smaller KV cache at 12B, and puts that model at or above dense on five of seven benchmarks, MRCR-v2 among them. At 1M tokens, 8× its training length, it finds every needle in single-needle NIAH once its memory cache is widened at inference, where the dense model finds 85%.