Graph Memory Transformer（GMT）：用可学习记忆图替代 Transformer 的 FFN 子层

View PDF HTML (experimental)

Abstract:We investigate whether an explicit learned memory graph can replace the Feed-Forward Network (FFN) sublayer in a decoder-only transformer while preserving the surrounding autoregressive architecture. In the proposed graph Memory Transformer (GMT), a memory cell routes token representations over a learned bank of centroids connected by a learned directed transition matrix, while causal self-attention remains unchanged. In the base v7 instantiation, each of 16 transformer blocks contains 128 centroids and a 128 x 128 edge matrix. Each cell uses gravitational source routing and token-conditioned target selection to return a gated displacement from the estimated source memory state toward the target memory state, rather than a retrieved value. The resulting decoder-only language model has 82.2M trainable parameters and no dense FFN sublayers; the evaluated dense GPT-style baseline has 103.0M. It trains stably and exposes centroid usage, transition structure, and source-to-target movement for direct inspection in the forward computation. It trails the larger 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 support the viability and structural interpretability of replacing dense within-token transformation with graph-mediated memory navigation, without claiming state-of-the-art performance. Broader scaling, optimized kernels, and more extensive benchmark evaluation remain for future work.

Comments:	48 pages, 10 figures, 5 numbered tables. Revised manuscript with selected generation ablations and editorial revisions. Code available at this https URL
Subjects:	Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL)
ACM classes:	I.2.6; I.2.7
Cite as:	arXiv:2604.23862 [cs.LG]
	(or arXiv:2604.23862v4 [cs.LG] for this version)
	https://doi.org/10.48550/arXiv.2604.23862 arXiv-issued DOI via DataCite

Submission history

From: Niccolò Ferrari [view email]
[v1] Sun, 26 Apr 2026 20:09:25 UTC (6,149 KB)
[v2] Wed, 27 May 2026 15:28:10 UTC (6,150 KB)
[v3] Thu, 28 May 2026 09:21:15 UTC (6,150 KB)
[v4] Thu, 22 Oct 2026 14:23:49 UTC (6,173 KB)