Graph Machine: Towards Better Pretraining via Edges

AI summary

Overview

Research area: Efficient sequence modeling and large language model architecture, specifically sparse attention, graph-structured inductive biases, and language-model pretraining.

Technical level: Advanced. The paper assumes familiarity with Transformer internals (queries, keys, values, KV heads, causal attention), complexity notation (Θ and O), and concepts such as adjacency matrices, softmax temperature scaling, and product-of-experts factorisation.

Scope: The paper introduces the Graph Machine (GM), a sequence-model architecture that keeps a Θ(n)-sized state but accesses only Θ(1) entries per step using Θ(log n) bits of dynamic addressing, and validates it by replacing 75% of the dense Transformer layers in Qwen3-0.6B and pretraining from scratch on 15.7B tokens.

What This Paper Is About

Transformers retain a large state in their keys and values but pay Θ(n²) cost because every token can look at everything; sliding-window attention cuts that access to Θ(1) but picks which positions to look at in advance, independently of the current token's content. The paper asks whether a model can keep a large state while choosing, per step and based on content, a tiny number of positions to retrieve — and whether such a model can actually be pretrained as a language model. The answer offered is the Graph Machine, which stores integer edge indices alongside floating-point node features and edge weights, and updates those indices differentiably through a "referral" mechanism resembling pointer chasing.

Key Contributions

1. A four-category taxonomy of sequence models based on state size, per-step access size, and per-step dynamic address bits. GM occupies a fourth category — Θ(n) state, Θ(1) access, Θ(log n) dynamic address bits — that RNNs/SSMs, vanilla Transformers, and sliding-window Transformers do not fill.
2. Sparse edge states coupled with a differentiable referral mechanism. GM maintains node features plus edge indices and edge weights of shape n × k × s₁. Rather than learning a discrete selection policy (which the paper argues would create credit assignment difficulties), it softens the edges: it takes a soft sparse adjacency matrix, raises it to the ℓ-th power, and sparsifies the result to top-s.
3. Two new sparse submodules: sparse edge referral (SER), which composes two-hop paths via A′ = Sparsify_{s₁}(A₁A₂) and truncates the up-to-s₂s₃ candidates back to s₁; and sparse edge attention (SEA), which uses edge indices to select which positions supply keys and values and combines node factors (scaled query–key scores) with edge factors (log edge weights) as a product of experts.
4. A controlled pretraining study. Graph Language Machines (GLMs) hybridise GM and Transformer layers at a 3:1 sparse-to-dense ratio, scheduled as [S, S, D, S] × 7, and are compared against Qwen3-0.6B trained from scratch with the same backbone, hyperparameters, codebase, and random seed.

Main Findings

• Most dense attention is trimmable at this scale. Retrieving only 2 or 4 of 4,096 tokens per KV head in each sparse layer yields only slight loss degradation at 2, and the best model marginally improves loss at 4.
• Loss numbers. Qwen3 test loss is +2.869 at 25% of training, +2.744 at 50%, +2.660 at 75%, and +2.587 at 100%. Final test losses across models are around 2.60. The best Theia model increases final loss by approximately 0.014, while the best Hyperion model reduces it by approximately 0.003.
• Referral helps. Theia-K24-R3 improves over the non-referral Theia-K24 baseline by 0.026 at 100% of training. Hyperion-K16-R4 improves over Hyperion-K16-R3 by 0.003, but under the sparser Theia budget with more edges than heads, Theia-K24-R4 performs 0.001 worse than Theia-K24-R3, so more referral steps do not help in every setting.
• More edges help moderately. In the matched Hyperion comparison, K24-R3 outperforms K16-R3 by 0.003 at the end of training, at additional parameter and compute cost because full cross-edge mixing scales as k².
• Refresh and realignment help. Hyperion-K16-R3-S improves over Hyperion-K16-R3 by 0.004 at 100% of training; preliminary experiments also suggest benefits from realignment and input refresh.
• Compute and parameter trade-offs. Referral-equipped GLMs add 9%–23% more parameters than Qwen3 but reduce total estimated training compute by 5%–15% and referral-plus-attention compute by 10%–30%. Hyperion-K16-R3-S achieves the best loss despite having 11% more parameters and 19% less referral-plus-attention compute than Qwen3.
• KV access reduction. SEA KV access over a 4,096-token sequence is 0.098% of dense causal attention for the Theia models (sparsity budget (2,4,4,2)) and 0.195% for the Hyperion models (sparsity budget (4,4,4,4)). This follows from dense causal attention accessing an average of (n+1)/2 KV positions per query.
• Neither parameters nor estimated compute predicts performance well within the tested configurations.
• Throughput caveat. Under the prototype implementation, GLMs are generally several times slower than Qwen3 on a single H100 SXM, with runs taking 53–236 hours (Qwen3: 53 hours; most GLMs: around 150–160 hours). Preliminary RTX 4090 experiments show some configurations approaching Qwen3's training throughput.

Methodology in Plain English

The starting point is a simple accounting exercise: if you want to keep a large memory but read only a constant amount from it at each step, each new token must carry enough identifying information — about log n bits — to say exactly where to look. The Graph Machine takes this literally and stores two kinds of things: floating-point node features, which hold content, and integer edge indices with associated floating-point edge weights, which hold addresses.

Each sparse layer does two jobs. First, referral builds new edges out of old ones: it mixes the stored edges into pairs, sparsifies each leg, then uses the first leg's indices to fetch the second leg's indices and weights, composing two-hop paths into one-hop paths. This is like neighbours telling you who their neighbours are. Because the number of candidate paths would blow up (up to s₂s₃ per edge), the result is cut back down to a fixed sparsity budget using hard top-s selection — only retained weights get gradients. Soft edge weights, not a learned discrete policy, carry the learning signal. Optional "refresh" edges pull in either the original input edges or the top-k positions from the most recent dense-attention layer, and a "realignment" step lets attention overwrite some stored edges.

Second, sparse edge attention uses the edge indices to decide which positions supply keys and values — just a handful of them. The final attention weight for a retrieved position combines the usual scaled query–key score with the logarithm of the mixed edge weight, which in probability space is a product of experts. RoPE is omitted from this submodule.

The whole thing is a hybrid: three sparse GM layers for every one standard dense Transformer layer. The authors call the resulting models Graph Language Machines, train them from scratch on a subset of FineWeb-Edu (documents packed into 4,096-token sequences, processed for less than one epoch, with a padding token at the start of each sequence to absorb underflowing edge indices), and compare them to Qwen3-0.6B under identical conditions. Models are named by sparsity budget, number of stored edges, and referral steps: Theia models retrieve 2 positions per KV head and Hyperion models retrieve 4.

Why This Matters

Impact on research. The paper reframes attention as one instance of a broader operation — passing and resolving addresses — and argues that if address resolution is part of attention's job, then logarithmic identification bits and indexed retrieval can be architectural primitives rather than emergent behaviour. It also opens an explicit design axis: relational traversal versus global search, and address passing versus content passing. The result that global dense attention is "highly trimmable" at this scale, while noting that some dense attention likely remains valuable, is a concrete data point for hybrid architectures.

Real-world applications (potential, not evaluated in this paper):

• Long-context inference, where KV access cost is the binding constraint and a model that reads 0.098%–0.195% of dense causal KV positions is attractive.
• Memory- or bandwidth-limited deployment, since the architecture trades dense computation for sparse memory operations.
• Retrieval-augmented systems, where the edge-index mechanism is structurally close to explicit pointer-based retrieval.
• Interpretability tooling, since the relational state is composed of readable indices and weights rather than opaque dense scores.

The paper does not test any downstream task, benchmark, or application; it evaluates only test loss in a pretraining setup.

Industry relevance. The compute figures are provocative: 5%–15% lower total estimated training compute and 10%–30% lower referral-plus-attention compute than Qwen3-0.6B, at 9%–23% more parameters. But the paper is candid that the prototype implementation is several times slower on an H100 SXM, so any practical payoff depends on custom kernels — the authors identify more efficient kernels as an important direction and note that RTX 4090 results already differ substantially from H100 results.

Future Directions

1. Efficient custom kernels. The authors state that under their prototype, GLMs are generally several times slower than Qwen3 on an H100 SXM, and that hardware and kernel implementation strongly shape relative performance. A custom Triton kernel is used only for sparsification in the current implementation.
2. Scale and richer evaluation. The experiments are small-scale in both model and training, and evaluate only a pretraining setup using test loss as an aggregate metric; richer evaluations across scales, setups, and downstream capabilities are left to future work.
3. Inductive bias and mechanistic interpretation. Although inductive bias was a central focus of GM-1, it remains to be examined for the updated architecture and in the language-modeling context. The authors suggest mechanistic interpretation could exploit GM's highly self-interpretable relational states and operations.
4. Finding the balance. On the efficiency side, dense computation can be traded against sparse memory operations; on the inductive-bias side, relational traversal can be traded against global search, and address passing against content passing. The authors state that finding the optimal balance will require further exploration.

Target Audience

Researchers and engineers working on efficient Transformer variants, sparse and approximate attention, long-context modeling, and alternatives to dense attention such as SSMs and their hybrids. It will also appeal to readers interested in graph-structured inductive biases in language models and in pointer-like or address-based computation. Because the paper adopts a standard language-model pretraining setting rather than a bespoke benchmark, its results are relatively easy to contextualise for practitioners — but the complexity notation, product-of-experts formulation, and sparse-coordinate implementation details make it most accessible to readers already comfortable with LLM architecture papers.

Authors’ abstract

We introduce the Graph Machine (GM), an architecture that maintains an $O(n)$-sized state and accesses it through sparse, dynamic routing. Unlike methods with fixed-size states or sparse but static routing, GM preserves $O(n)$ complexity in its sparse layers without restricting the potentially accessible state size to $O(1)$. Instead, GM uses edges - pointer-like objects updated differentiably by a referral mechanism resembling pointer chasing. We replace 75% of the dense Transformer layers in Qwen3-0.6B with GM sparse layers and pretrain from scratch on 15.7B tokens. With only 2 of 4,096 tokens retrieved per KV head in each sparse layer, loss degrades only slightly; with 4, the best model marginally improves loss.