Graph Memory Transformer (GMT)

Abstract: We investigate whether the Feed-Forward Network (FFN) sublayer in a decoder-only transformer can be replaced by an explicit learned memory graph while preserving the surrounding autoregressive architecture. The proposed Graph Memory Transformer (GMT) keeps causal self-attention intact, but replaces the usual per-token FFN transformation with a memory cell that routes token representations over a learned bank of centroids connected by a learned directed transition matrix. In the base GMT v7 instantiation studied here, each of 16 transformer blocks contains 128 centroids, a 128 * 128 edge matrix, gravitational source routing, token-conditioned target selection, and a gated displacement readout. The cell therefore returns movement from an estimated source memory state toward a target memory state, rather than a retrieved value. The resulting model is a fully decoder-only LLM with 82.2M trainable parameters and no dense FFN sublayers, compared with a 103.0M-parameter dense GPT-style baseline used in the evaluation. The base v7 model trains stably and exposes centroid usage, transition structure, and source-to-target movement as directly inspectable quantities of the forward computation. It remains behind the larger dense baseline in validation loss and perplexity (3.5995/36.58 vs. 3.2903/26.85), while showing close zero-shot benchmark behavior under the evaluated setting. These results are not intended as a state-of-the-art claim; they support the viability and structural interpretability of replacing dense within-token transformation with graph-mediated memory navigation. Broader scaling, optimized kernels, and more extensive benchmark evaluation are left for subsequent work.

• The paper replaces standard transformer FFNs with a graph-structured memory cell, providing a more interpretable token transformation mechanism.
• It introduces learned centroids, a directed edge matrix, and gravitational source routing to compute displacement vectors that replace nonlinear MLP updates.
• Empirical evaluations reveal competitive performance on language benchmarks while enabling detailed insight into memory utilization and routing dynamics.

Graph Memory Transformer: Replacing Transformer FFNs with Graph-Structured Memory

Introduction and Motivation

The Graph Memory Transformer (GMT) (2604.23862) proposes a structural alternative to the standard feed-forward sublayer in decoder-only transformers. Instead of the conventional position-wise dense multilayer perceptron (MLP), GMT introduces a learned graph memory cell that routes token representations over a layer-local bank of centroids connected by a directed transition matrix. This architectural modification aims to expose internal computation in a form that is more structurally interpretable, addressing opacity in the FFN's contribution to token-wise transformation.

GMT retains standard causal self-attention as the context aggregator but replaces the subsequent nonlinear MLP with a memory cell composed of persistent centroids and directed edges. At each transformer block, the memory cell estimates a source distribution over centroids for token states, propagates this through the transition matrix, adjusts the result via token-conditioned targeting, and returns a gated displacement vector to the residual stream.

Architecture and Mechanisms

Memory Cell: Centroids and Routing

Each GMT block contains:

• Centroid bank: 128 learned centroids ($\mathbb{R}^{H}$ vectors), per block.
• Edge matrix: A $128\times128$ learned matrix, made row-stochastic after masking diagonals.
• Gravitational source routing: Tokens are softly assigned to centroids by inverse-distance similarity.

The memory cell first computes a soft source routing distribution for each token, identifying its location in centroid space. This is performed via normalized cosine similarity, with a reciprocal (gravitational) temperature-controlled weighting to ensure sharp but differentiable slot assignment. This process is visually summarized for the earliest block in the stack as follows:

Figure 1: Slot-routing flow at Block~00, showing concentrated early intake routes.

Graph Traversal and Displacement Readout

After source assignment, the token's routing weight vector is diffused via the row-stochastic transition matrix, yielding an edge-propagated distribution. This is further refined using token-specific key-query scoring, and the resultant target distribution over centroids is computed.

Given source and target distributions, the memory cell aggregates centroid vectors at both endpoints and computes the displacement vector—the difference between the target and source memory state. This displacement is layer-normalized and injected into the block residual path, modulated by a learnable scalar gate.

This explicit displacement mechanism encodes the within-block transformation as a geometric "move" in the learned memory space, contrasting with the implicit nonlinear update in a traditional MLP. The routing and displacement dynamics evolve through the block stack, as depicted below:

Figure 2: Slot-routing flow at Block~06, showing the diversification of routing as depth increases.

Figure 3: Slot-routing flow at Block~11, corresponding to maximal routing diversity and mid-depth structural reorganization.

Empirical Study and Results

Training Setup and Baselines

The principal instantiation, GMT v7, utilizes 16 transformer blocks of 768 hidden units, each with a 128-slot graph memory cell. It is evaluated as a decoder-only LM with 82.2M parameters, compared to a dense GPT-style baseline of 103.0M parameters (identical in all respects except FFN replacement).

Both models are trained and validated on OpenWebText, using standard autoregressive objectives and the same data splits and tokenizer conventions. The GMT auxiliary training objective includes centroid-geometry, clustering, tracking, edge-entropy, and contrastive regularization terms to maintain slot utilization and transition diversity.

Predictive Performance

• Validation loss: The dense baseline achieves 3.2903 (PPL 26.85), while GMT v7 attains 3.5995 (PPL 36.58); the gap is expected given the parameter advantage maintained by the baseline.
• Zero-shot tasks: On ARC-Easy, PIQA, HellaSwag, and WinoGrande, the GMT v7 is within 1.7 percentage points or less of the dense baseline across most metrics, and outperforms the baseline by 1.0 percentage point on WinoGrande, a commonsense reasoning task.

Structural Analysis and Interpretability

GMT enables detailed introspection into learned memory utilization and routing:

• Slot routing evolution: Early blocks show concentrated intake through a dominant centroid, with destination routing fanning out over multiple slots. Deeper blocks exhibit more distributed utilization, with linguistic role (POS) and topical structure modulating path selection.

Figure 4: Topic-separated Block~11 routing flows, demonstrating shared regime but topic-conditioned centroid activation.

• Final output compression: The last block collapses routing into a dominant destination slot; however, source slot provenance remains topic-dependent, preserving distinctive upstream information.

  Figure 5: Slot-routing flow at Block~15, illustrating late-stage output compression with differentiated input sources.

• Edge structure visualization: Active edge matrices highlight topic-sensitive, slot-specific routing preferences and the emergence of selective, context-conditioned transitions.

  Figure 6: Active edge structure at Block~00; dark cells denote strong routing among centroids.

  

  Figure 7: Active edge structure at Block~06.

  

  Figure 8: Active edge structure at Block~11.

  

  Figure 9: Active edge structure at Block~15, corresponding to strong output address consolidation.

• Role switching in memory: Analysis of slot transitions (especially in Block~13 under political text) reveals centroids functioning alternately as source and destination, with distinctly different token classes participating in the moves.

  Figure 10: Slot-routing flow in block 13 for a political-text probe, illustrating dynamic source/destination switching within the memory graph.

Auxiliary Objectives and Stability

Stability of the memory mechanism is ensured by:

• Centroid orthogonality: Encouraging geometric diversity to prevent slot collapse.
• Usage entropy and clustering: Penalizing under-utilization and preserving effective slot count.
• Edge entropy and contrast: Maintaining transition diversity by discouraging premature or homogenized edge specialization.
• Online write-back and periodic maintenance: EMA-based centroid updates (hard assignment for write-back, soft for traversal) and reset/merge policies to recycle inactive or redundant centroids.

Empirical diagnostics indicate rapid reduction in centroid similarity during early training (from 0.587 to 0.155 mean cosine similarity) and persistent distributed slot usage (effective number typically above 60 for 128 slots).

Implications and Future Directions

Practical and Theoretical Impact

GMT demonstrates that block-level MLPs may be replaced with graph-based memory cells without substantial loss in standard autoregressive LM performance in the moderate-parameter regime. The architectural shift yields a model whose internal transformation variables—centroid usage, transition structure, source-target displacement—are natively interpretable. This considerably facilitates mechanistic attribution and analysis compared to post hoc interpretability of dense MLP weights.

The displacement-based readout, edge-augmented memory, and exposure of routing variables provide a promising substrate for deeper mechanistic interventions, facile memory editing, and potential integration with prototype-based or sparse-feature techniques.

Future Research

Key open questions and suggested pathways:

• Parameter-matched and multi-seed baselines: Deeper analysis of efficiency and robustness under strictly controlled conditions.
• Hybrid models: Investigation into hybrid dense/graph-memory computation and controlled residual scaling, merging benefits of both mechanisms.
• Multi-hop graph traversal: Extension of GMT with deeper memory navigation per block, testing compositional concept assembly.
• Maintenance strategies: Exploring alternative write-back policies (soft, confidence-gated, top-$k$), stabilizing memory during evaluation.
• Vision and anomaly domains: Adapting graph memory routing to patch-level or pixelwise memory in visual tasks, leveraging displacement and routing entropy for improved anomaly detection.
• Mechanistic interventions: Systematic ablations (graph structure, edge weights, displacements) to disentangle the specific contributions of each architectural element.

Conclusion

GMT (2604.23862) provides a structurally interpretable, memory-rich alternative to standard transformer FFNs. By embedding token transformation in a persistent, directed graph memory with explicit source-to-target displacement, GMT offers both operational viability and native introspectability. While some performance separation remains relative to dense baselines, the architectural innovation sets a concrete foundation for further research into interpretable, controllable, and compositionally expressive neural architectures.