lil-lab/sps

TL;DR: SPS (from Artzi's lab at Cornell, lil-lab) is the training/analysis code for a paper arguing that transformers overload one computation stream with two jobs — predicting the next token and storing state for future positions to read via the KV cache — and that splitting these roles helps. The SPS Transformer interleaves a persistent "input" stream with an ephemeral "prediction" stream (kept only in a sliding window), and reportedly improves validation loss, OOD generalization, and zero-shot accuracy from 53M to 1.68B params at matched inference cost. It's a genuinely different information-flow design, not an attention tweak — but note there are no released weights, and the README quotes no numbers.

The idea

In a standard decoder, position t's hidden states must simultaneously (a) be a good basis for predicting token t+1 and (b) be a good long-lived memory that positions t+100 will attend to. These objectives can conflict: prediction wants position-specific, soon-to-be-stale features; state wants durable, reusable ones. The SPS hypothesis says decoupling them yields better language modeling.

The realization is simple: after every input token, insert a dedicated prediction token. This creates two interleaved streams. The input stream is persistent — it stays in the KV cache forever and carries state forward, but never emits predictions. The prediction stream reads from the input stream and emits the next-token prediction, but is ephemeral: its KV entries are kept only within a sliding window w (the paper uses w=64), so it never accumulates as long-range state.

SPS: interleaved streams x₁ p₁ x₂ p₂ x₃ p₃

↓ predicts x₂ ↓ predicts x₄

input stream: persistent state, stays in KV cache prediction stream: ephemeral, kept only in sliding window w (paper: w=64)

The ablation set is what makes this convincing as a hypothesis test rather than an architecture hack: reverse_sps (swap which stream predicts), delayed_state, and a 2× memory control (SPS with w=4096, i.e., the prediction stream also persists) — the last isolates whether gains come from the separation itself rather than extra KV capacity. There are also mechanistic probes: a gradient-ratio analysis (how much gradient each stream gets from future vs. present positions) and multi-token-ahead linear probes measuring how far ahead each stream predicts — a nice interpretability angle on why the separation helps.

What's actually here

Full pretraining pipeline: Hydra recipes for 5 scales (53M–1.68B) × 4 families × 10B/20B FineWeb-Edu token budgets, custom Triton FlashAttention kernels for the interleaved masks, lm-eval-harness evaluation (FineWeb-Edu val loss; corpus NLL on WikiText, C4, Books3, GovReport; zero-shot ARC-E, HellaSwag, PIQA, SciQ, LAMBADA), and scripts reproducing every paper figure and table — including DeepMind-convention per-architecture FLOP accounting for compute-matched comparisons, and a throughput/peak-memory benchmark protocol.

What's not here: pretrained checkpoints, and any concrete numbers in the README. All claims ("lowers validation loss, improves held-out generalization, raises zero-shot accuracy at matched inference cost") point to the paper (arXiv 2607.01218). Every figure script reads your own W&B runs or local checkpoints — verifying anything means pretraining yourself, though the 53M/10B-token recipe should be cheap on a single node.

Caveats

• Training compute is not free: doubling sequence length with interleaved prediction tokens roughly doubles attention/FFN work per real token. The repo's compute-matched plots (loss vs. training FLOPs with per-architecture frontiers) exist precisely to address this — worth scrutinizing in the paper before taking "better" at face value. The "matched inference cost" claim is more plausible, since the ephemeral stream adds only ~w KV entries.
• 1.68B is small by current standards; whether the gap persists or shrinks at scale is the open question.
• CUDA + Triton required for the fused kernels (CPU fallbacks only for tests). Research code, but tested, MIT-licensed, and from a credible lab (Artzi, Brantley).

This sits in the same conceptual neighborhood as pause/filler tokens and multi-token prediction, but the framing — KV cache as state store vs. forward pass as predictor, with the separation as the explanatory variable — is cleaner and the ablations directly test it. If it holds up at scale, "give state its own stream" is cheap enough to become a standard recipe ingredient.

Try it

git clone https://github.com/lil-lab/sps && cd sps
uv sync
uv run python src/data/prepare.py system.data_root=/data
# smallest recipe: 53M params, SPS w=64
uv run python scripts/train.py +experiment=xs_sps_w64_10b system.data_root=/data
uv run python scripts/evaluate.py +experiment=xs_sps_w64_10b +checkpoint=final system.data_root=/data
# renders attention-mask schematics with no training or data
uv run python scripts/figures/plot_attention_masks.py