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

AI summary

Overview

Research area: Deep learning theory and mechanistic interpretability — specifically how deep sequence models store facts in their parameters, and what the learned embedding geometry reveals about memorization.

Technical level: Advanced. The paper combines empirical sequence-modeling experiments (Transformer, Mamba, small MLPs) with analysis framed around graph embeddings, spectral biases, and the Node2Vec objective. The concepts are explained accessibly, but the framing and the mathematical definitions assume familiarity with representation learning and gradient-descent training dynamics.

Scope (one sentence): The paper identifies a form of parametric memory it calls geometric memory, shows it arises in sequence models even when a simpler brute-force lookup would work, and argues that standard explanations (supervision, capacity limits, optimizer preferences) do not account for it, linking it instead to a spectral bias in Node2Vec-style models.

What This Paper Is About

Deep sequence models are usually assumed to store atomic facts as associative memory: the embeddings of entities are arbitrary keys, and co-occurrences are stored as a lookup in a weight matrix. The authors show that models instead often build a geometric memory, where the embeddings themselves encode global, multi-hop relationships between entities — including pairs that never appear together in training. The puzzle is why a model would prefer this geometry over the lookup table when both fit the training data equally well and the lookup is, by the authors' argument, easier to find.

Key Contributions

1. The authors construct a clean instance of implicit in-weights reasoning: models memorize a graph's edges in their weights and then answer path-finding queries, with no form of step-wise supervision. They present this as behavior that is hard to reconcile with the dominant associative view of parametric memory. (§2)
2. They propose an alternative, global geometric view of parametric memory in deep sequence models, demonstrated across Transformers, Mamba, and simple neural networks. (§2.4)
3. They argue the geometry is surprising by empirically refuting supervisory, explicit-capacity, and optimization-based explanations for it, and frame this as a memorization puzzle in sequence modeling. (§3)
4. They connect the observed geometry to the spectral bias of 2-layer Node2Vec models, empirically intuit how it emerges from cross-entropy loss minimization without typically assumed pressures, and identify headroom for making Transformer memory more strongly geometric. (§4)

Main Findings

• Geometry beats a hard composition problem. On a path-star graph — a topology designed to make next-token-trained models fail when the graph is supplied in context — making the model memorize the graph's edges in weights turns an ℓ-fold composition into what the authors describe as an easy 1-step navigation task. The authors state that learning the ℓ-fold composition should otherwise demand Ω(exp(ℓ)) compute.

• Success on large graphs, including learning the hardest token in isolation. With path lengths of 6 to 10 hops and graphs as large as 10⁴ nodes, a next-token-trained Transformer achieves perfect or highly non-trivial accuracy. On graphs of as many as 5×10⁴ nodes, when trained only on edge-memorization examples and first-token-training examples from 75% of all paths, both the Transformer and Mamba predict the first token on unseen paths with as much as 100% accuracy.

• The first-token setup removes every known aid. Prior positive results on implicit in-weights reasoning relied on curricula (paths of varying lengths), full path supervision, test-train path overlap, or 2-hop tasks. Here paths have fixed lengths, later-token gradients are eliminated, paths are disjoint, and composition can be as deep as 10-fold.

• Embeddings show a global geometry. Heatmaps of cosine similarity between the leaf embedding of one path and the first-hop embedding of another show a clear diagonal: embeddings within a path are more aligned with each other than with other paths. The authors report similar geometries for Mamba and for a simple 3-layer neural network on a variety of small graphs, and note their visualizations can show that a geometry exists but do not definitively establish how global it is (they recommend examining inner products).

• Geometry appears from purely local supervision. A global geometry emerges even in models trained only on local edge-memorization, without path-finding supervision. Such models can then be finetuned purely on the hardest-token task and reach high test accuracy on unseen paths. Test paths that never overlap training paths still exhibit the geometry.

• Explicit capacity pressure is not the cause. The models can express the associative lookup. The authors cite results roughly showing that with m frozen embedding dimensions and at least m² free parameters, an MLP layer can store m² associations; their models have m ≈ 400. They also show that in the very settings where geometries emerge, the same training setup can fit the data associatively when embeddings are frozen and one intermediate layer is trainable, and that a geometry still arises with weight decay and dropout switched off.

• Optimization pressure is not the cause either. Rather than a lookup being slow to find, the authors report the opposite ordering: for wide models on the tiny graphs, associative memory forms in 2 steps while geometric memory requires 100 steps, under the cross-entropy loss. They propose their setting as arguably the most minimal instance of grokking, and note that models fail at the task when embeddings are frozen.

• A Node2Vec connection explains the geometry's origin. By analyzing a connection to Node2Vec, the authors argue the geometry stems from a spectral bias that arises naturally from cross-entropy loss minimization, in contrast to prevailing theories (they cite Levy and Goldberg in this context). In the simplest geometric view, associations are factorized as Φ_geom(u)ᵀ Φ_geom(v), where Φ_geom corresponds to eigenvectors of the graph Laplacian, potentially only the top eigenvectors.

• Transformer geometry has measurable headroom. The embeddings learned by the more naive Node2Vec-style models are more strongly geometric than those of Transformers, which the authors present as a well-specified headroom for making Transformer memory more geometric in practice.

• Two competing data structures, compared. Associative memory stores a lookup table of pairwise edge associations whose description length is |E| (up to logarithmic factors) and yields precise retrieval; geometric memory stores a graph embedding of the vertices whose description length is |V|·m and yields approximate retrieval. The authors also discuss variant factorizations — full-rank, or of the form Φ_geom(u)ᵀ Λ Φ_geom(v) with a diagonal Λ containing both positive and negative values — which may not produce clean multi-hop logit values but can still be probed for global directions.

Methodology in Plain English

The authors build a controlled task around graphs. Instead of handing the model a graph in its prompt, they make the model memorize the graph's edges in its weights by training on edge examples that cover all edges. Then they ask the model to answer path-finding queries on path-star graphs — trees where several disjoint paths of the same length branch out from a root — where the input is a leaf node and the target is the root-to-leaf path.

The adversarial part of the design is that almost all of the path can be predicted trivially from the revealed context, which means the very first token carries the whole reasoning burden. In the original in-context version of this task, that first token is where models fail. The authors make two key alterations. First, all examples come from one fixed graph held in weights rather than a fresh graph in context. Second, they train path-finding on only a subset of leaves and test on leaves whose root-to-leaf paths are never seen end to end, so the constituent edges are known individually but the full path is not. In the sharpest experiments, they remove the later-token losses entirely and train only on the first-token loss.

To probe what the model has learned, they visualize the learned token embeddings — for instance, a heatmap of cosine distances between the leaf embedding of one path and the first-hop embedding of another — and look for structure that only makes sense if the model has organized entities globally. They then compare these geometries across architectures (a from-scratch decoder-only Transformer they call GPT-mid, Mamba, and a simple 3-layer neural network), and contrast them against a 2-layer Node2Vec-style model in which associative storage is explicitly prohibited by removing intermediate layers. Finally, they test whether the geometry survives when the supposed causes are removed: training on local edge supervision only, disabling weight decay and dropout, checking whether the fixed graph could be stored associatively in the available parameters, and measuring how many gradient steps each memory type takes to form. Details such as the exact hyperparameters, formatting, and further analyses in the paper's appendices are referenced but their specific values are not reported in the content available here.

Why This Matters

Impact on research. The paper argues that the associative view of parametric memory is a default, often unstated intuition guiding work on knowledge acquisition, discovery, unlearning, reasoning, storage capacity, and scaling laws. If memorized facts are stored in an interdependent geometry rather than as independent lookups, predictions in those areas may need revisiting. The authors also offer their minimal setups, which have closed-form solutions, as tractable sandboxes for studying grokking, the emergence of "world models," linear representations and superposition in interpretability, and the Platonic representation hypothesis.

Real-world applications (as the paper discusses them):

• Knowledge editing and unlearning: interdependencies in geometric storage may impose limits on editing facts, for example causing effects such as representation shattering, and may complicate unlearning and accurate retrieval.
• Hallucination and spurious association: the authors suggest geometric storage may hallucinate associations, drawing a comparison to "illusory correlations."
• Retrieval systems: the gap between Transformer and Node2Vec geometries is presented as relevant to choosing between modern generative retrieval models and traditional dual-encoder models.
• Natural implicit reasoning and discovery: if the geometric bias can be improved on natural language tasks, the authors argue it could benefit implicit reasoning tasks where results have so far been mixed, and could help surface novel connections between information scattered across large pretraining sets.

Industry relevance. The paper points to a concrete engineering target: the finding that naive Node2Vec-style embeddings are more strongly geometric than Transformer embeddings suggests a measurable direction for redesigning Transformer memory. It also implies that systems relying on precise fact retrieval, targeted editing, or removal of specific knowledge should account for the fact that stored facts are entangled in a shared geometry rather than held as separate lookup entries.

Future Directions

• How do the two memories compete under gradient descent? The authors state this as their foremost open question, along with which aspects of the data or of the optimization make the geometry arise.
• A foundational open question for deep sequence models. The spectral-bias account is given for a minimal two-layer, Node2Vec-style architecture; the authors explicitly leave open whether it explains geometric memory in full deep sequence models.
• Closing the geometry gap in practice. The headroom between Transformer and Node2Vec embedding geometry is identified but how to exploit it — and whether it transfers to natural language — is left open.
• Re-examining capacity and scaling laws through a geometric lens. The authors suggest capacity and scaling-law analyses could yield very different answers under the geometric view, and that the examples here can be used to study grokking, world models, superposition, and the Platonic representation hypothesis in a tractable setting.

Target Audience

Researchers and graduate students in machine learning theory, mechanistic interpretability, and representation learning who are interested in how networks store and reuse facts; practitioners working on retrieval systems, knowledge editing, or unlearning who need to reason about how parametric memory behaves; and anyone following debates about implicit reasoning, grokking, or the relationship between memorization and generalization. The paper is best suited to readers comfortable with embedding-based models, cross-entropy training, and graph spectral concepts; it is not an introductory tutorial.

Authors’ abstract

Deep sequence models are said to store atomic facts predominantly in the form of associative memory: a brute-force lookup of co-occurring entities. We identify a dramatically different form of storage of atomic facts that we term as geometric memory. Here, the model has synthesized embeddings encoding novel global relationships between all entities, including ones that do not co-occur in training. Such storage is powerful: for instance, we show how it transforms a hard reasoning task involving an $\ell$-fold composition into an easy-to-learn $1$-step navigation task. From this phenomenon, we extract fundamental aspects of neural embedding geometries that are hard to explain. We argue that the rise of such a geometry, as against a lookup of local associations, cannot be straightforwardly attributed to typical supervisory, architectural, or optimizational pressures. Counterintuitively, a geometry is learned even when it is more complex than the brute-force lookup. Then, by analyzing a connection to Node2Vec, we demonstrate how the geometry stems from a spectral bias that -- in contrast to prevailing theories -- indeed arises naturally despite the lack of various pressures. This analysis also points out to practitioners a visible headroom to make Transformer memory more strongly geometric. We hope the geometric view of parametric memory encourages revisiting the default intuitions that guide researchers in areas like knowledge acquisition, capacity, discovery, and unlearning.