Deep sequence models tend to memorize geometrically; it is unclear why

TL;DR

Transformers and Mamba, when asked to memorize the edges of a large path-star graph and then plan paths over it implicitly (no chain-of-thought, no intermediate supervision), succeed at scales (10k nodes, 10-hop paths) where the standard “memory as associative lookup” abstraction predicts failure. The reason: the model spontaneously learns a geometric embedding of the graph — non-co-occurring entities end up positioned to reflect global multi-hop distances — turning an L-fold compositional recall into a 1-step spatial query. The paper isolates why this is genuinely surprising (it isn’t forced by capacity bottlenecks, optimizer pressure, or supervision), connects it to a CE-loss-induced low-rank spectral bias visible already in Node2Vec, and frames “associative vs geometric memory” as the memorization puzzle in sequence modeling.

Key claims

On an in-weights path-star task with up to 10^4 nodes and 10-hop paths, decoder-only Transformers and Mamba reach ~100% held-out path-prediction accuracy without chain-of-thought or per-hop supervision [§2.1, Fig. 4].
The hardest token (the first one, requiring full L-fold composition) is learned in isolation when only the first-token loss is used — directly contradicting the predicted exponential-in-L hardness of strict associative recall [Observation 2.1.1, Fig. 4 right].
Tokens along the path are not learned in reverse order, weakening the “future-token gradients let the model find the right-to-left solution” hypothesis [Observation 2.1.1, Fig. 5, Fig. 24].
The trained models embed nodes such that leaf and first-hop embeddings of the same path are more aligned than across paths, exposing global multi-hop structure not present in any single training example [§2.3, Fig. 6].
The emergence of geometric memory cannot be attributed to architectural capacity bottlenecks: in setups where geometric memory naturally arises, freezing embeddings still permits a purely associative solution at the same architecture/optimizer [Observation 3.1, §C.3].
Geometric memory is not strictly more succinct than associative memory — there exist graphs where both representations have equal parameter cost, yet the model still picks geometric [§3.1].
The geometric bias maps onto a low-rank spectral bias of CE loss already present in shallow Node2Vec — top eigenvectors of the adjacency dominate, and powers of A (which encode multi-hop distance) are revealed for free; prior Node2Vec analyses needed extra pressures (bottlenecks, early stopping) that this work argues are unnecessary [§5, §5.6].
Empirically, Node2Vec produces more strongly geometric embeddings than the Transformer on the same graphs, indicating headroom to make Transformer memory more geometric in practice [§1.1].

Method

The setup is the in-weights path-star task: a single fixed graph (root + D disjoint length-L paths) is memorized into the weights through edge-memorization examples (input: node v; target: a neighbor w), interleaved with path-finding examples (input: a leaf; target: the full root→leaf path). Crucially, the same graph generates all examples, paths are disjoint (no train/test substring overlap), lengths are fixed (no implicit curriculum), and forward+reverse edges are provided to defuse the reversal curse for edge memorization only. Training is from-scratch GPT-mid (and Mamba for replication), with pause tokens for stability.

The analysis then dismantles each candidate explanation for why the hardest first token is learned (Clever Hans cheats, future-token gradients, capacity pressure, implicit compression). For the positive side, the authors visualize embedding geometries (UMAP, dot-product heatmaps between leaf-of-path-i and first-hop-of-path-j) showing same-path clustering, and reduce the phenomenon to a Node2Vec analogue where they show the cross-entropy loss naturally amplifies top eigenvectors of the graph adjacency matrix — eigenvectors that encode multi-hop relations through A^L.

Results

On a graph with ~10^4 nodes and 10-hop paths, both Transformer and Mamba achieve close to 100% held-out path accuracy from leaf-only conditioning [Fig. 4 left, Fig. 8 left].
First-token-only training (no path supervision past the first token) still solves the task to high accuracy [Fig. 4 right].
Heatmap of ⟨leaf_i, first-hop_j⟩ shows a sharp diagonal: within-path alignment is consistently higher than across-path [Fig. 6a, Fig. 9 for Mamba].
Even with only edge-memorization examples (no path-finding objective at all), a weaker but visible geometric structure still emerges [Fig. 6c].
Freezing token embeddings makes the same model+optimizer learn a clean associative solution — proving geometry is a learning bias, not an architectural necessity [§E.1, §C.3].

Why it’s interesting

This paper supplies a concrete, clean mechanism for “implicit in-weights reasoning” that is otherwise easy to dismiss as a curiosity of small toy graphs. For a generative-models team, the headline implication is that pretraining can — and apparently does — synthesize global geometric structure over facts that never co-occur, which is exactly the substrate one would want for combinational creativity, world-model emergence, and zero-shot multi-hop inference. The flip side, called out in the paper, is that this same interdependent geometry is what makes knowledge editing, unlearning, and faithful retrieval brittle — relevant to anyone trying to surgically modify a large model’s parametric knowledge. Finally, the Node2Vec → Transformer reduction reframes the long-standing low-rank spectral bias story as a property of CE loss itself rather than of bottlenecked architectures, which is a useful sharpening of the prevailing theoretical picture.