Source-linked AI summary
Graph Machine: Towards Better Pretraining via Edges
Lintai Hou
TL;DR
The Graph Machine addresses the challenge of accessing a large state with sparse, dynamic routing by maintaining graph state and differentiably updated edges. Hybrid Graph Language Machines replace most dense Transformer layers with sparse retrieval, preserving quality while reducing estimated compute at the evaluated scale and setup.
Problem
Dynamic access to a Θ(n)-sized state requires Θ(log n) directly usable addressing bits, motivating an architecture that combines large state with constant-sized retrieval.
Method
Graph Machine maintains Θ(n) node features, edge indices, and edge weights, using differentiable referral and sparse edge attention to retrieve and update selected positions.
Results
At the evaluated scale and setup, most Transformer layers could be replaced by sparse layers retrieving 2 or 4 of 4,096 positions per KV head without materially sacrificing quality, while reducing estimated compute.
Takeaways & Limitations
GM makes address passing and direct indexed retrieval architectural primitives, enabling trade-offs between dense computation and sparse memory operations.
Takeaways & Limitations
The experiments are small-scale, use only a pretraining setup, and evaluate test loss as an aggregate metric, leaving broader evaluations for future work.
Abstract
from arXiv · showhide
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.
1 Introduction
Existing sequence models trade off state size, access size, and dynamic addressing: Transformers retain a Θ(n)-sized state but incur Θ(n^2) attention, while sparse alternatives often use fixed access. Graph Machine instead combines Θ(n) state, Θ(1) access, and Θ(log n) dynamic addressing through differentiable sparse edges and referral.
- Recurrent models maintain a Θ(1)-sized state and must compress their history, whereas Transformers retain a Θ(n)-sized state in keys and values.
- Transformers’ unrestricted access to the full state produces Θ(n^2) complexity, while sliding-window attention reduces access to Θ(1) but remains static.
- Dynamic Θ(1)-sized access to a Θ(n)-sized state requires each token to provide Θ(log n) directly usable addressing bits.
- Graph Machine maintains Θ(n) state entries, accesses Θ(1) entries per step, and uses Θ(log n) dynamic-addressing bits, storing node features alongside integer edge targets.
- Referral constructs new neighborhoods by composing ℓ-hop edges, potentially yielding s^ℓ candidate paths before sparsification restores neighborhood size s.
- Graph Machine represents edges with integer indices and differentiable floating-point weights, which contribute to attention alongside query–key factors.
- The resulting Graph Language Machines hybridize Graph Machine and Transformer layers while maintaining and updating a graph with soft directed edges.
2 Architecture
Graph Machine layers maintain sparse, weighted edge states and update them through mixing, referral, and edge-based attention. These operations enable dynamic information routing while preserving sparse representations across hybrid GM–Transformer layers.
- States: GM maintains node features together with sparse edge indices and weights, representing generalized adjacency matrices across multiple edge copies.Each node has k edges, each with s1 member positions; indices identify targets and weights are nonnegative and normalized.
- Layers: Each sparse layer applies edge-referral submodules followed by edge-attention operations, alongside standard dense Transformer layers.The supplied architecture description specifies this hybrid layer structure.
- Operations: Sparsification coalesces duplicate edge indices, selects the highest-weight entries, and renormalizes them to the target sparsity.Only retained weights receive gradients, and sparsification is used before referral and attention and after referral.
- Sparse edge referral: Referral composes two sparse edge matrices into new edges and sparsifies the result, allowing multi-hop paths to form updated neighborhoods.The two-hop operation is expressed as A′ = Sparsify_s1(A1A2), with intermediate edge members reduced back to the stored sparsity.
- Sparse edge attention: Sparse edge attention retrieves positions from mixed stored edges and combines edge factors with query–key node factors as a product of experts.The resulting attention weights use both mixed edge weights and the usual scaled query–key scores, followed by standard value aggregation and output projection.
3 Results
The experiments evaluate GM configurations against Qwen3 under controlled pretraining and report sparse retrieval with lower estimated compute but hardware-sensitive runtime. Performance varies with edge count, referral depth, refresh, and sparsity, with the best Hyperion model slightly outperforming Qwen3.
- All models use the Qwen3-0.6B backbone, matched training hyperparameters, codebase, random seed, and a conventional untuned LLM recipe.
- GM uses a 3:1 sparse-to-dense layer ratio with 16–32 edges, 0–6 referral steps, and Theia or Hyperion sparsity budgets.The schedule is [S, S, D, S] × 7, with realignment enabled and optional dense refresh.
- 9%–23% more parameters accompany referral-equipped GLMs, while estimated training compute falls 5%–15% and referral-plus-attention compute falls 10%–30% relative to Qwen3.Most additional parameters belong to referral projections, which support relatively inexpensive operations.
- 2 or 4 retrieved positions per KV head correspond to 0.098% or 0.195% of dense causal KV access at sequence length 4,096.
- 0.003 lower final loss distinguishes Hyperion-K24-R3 from K16-R3, while full cross-edge mixing scales as k^2 and raises parameter and compute cost.
- 0.026 improvement at 100% distinguishes Theia-K24-R3 from the non-referral baseline, but Theia-K24-R4 is 0.001 worse than R3 under the sparser budget.
- 0.004 improvement at 100% distinguishes Hyperion-K16-R3-S from Hyperion-K16-R3, while the best Hyperion model reduces final loss by approximately 0.003.
- Hyperion-K16-R3-S achieves the best loss with 11% more parameters and 19% less referral-plus-attention compute than Qwen3.Within the tested configurations, parameter count and estimated compute are not strong performance predictors.
4 Related work
GM-2 extends the original GM work from a bespoke Sudoku benchmark to standard language-model pretraining and replaces its dense edge representation with sparse edge representations and operations. It is also positioned within efficient sequence modeling and sparse-attention research.
- GM-2 adopts standard language-model pretraining, making its results easier to contextualize than GM-1’s bespoke Sudoku benchmark.
- GM-1 studies edges with a dense representation whose cubic time and quadratic space complexity limited experiments to a few hundred nodes.
- GM-2 uses sparse edge representations and operations, with edge indices and weights forming a sparse coordinate representation of GM-1’s edge addresses.
- GM is related to efficient sequence models, hybrids, and the sparse-attention literature.
5 Limitations and conclusion
The work demonstrates that dense attention can be substantially trimmed at the evaluated scale and setup, while leaving important evaluation and architectural questions open. It frames sparse address-based computation as a design space requiring further exploration.
- Limitations: The work leaves substantial room for architectural and implementation improvements, including more efficient custom kernels.This is stated as an explicit direction for future improvement.
- Limitations: The experiments are small-scale, use only pretraining with test loss, and leave richer evaluations across scales, setups, and downstream capabilities for future work.The authors explicitly limit the evidence to this evaluation scope.
- Limitations: The updated architecture’s inductive bias remains unexamined in language modeling.The paper identifies this as an open question from prior work.
- Conclusion: At the authors’ scale, setup, and measure, global dense attention is highly trimmable, although some dense attention likely remains valuable.The conclusion preserves this scope qualification rather than claiming universal replacement.
- Conclusion: GM creates architectural trade-offs between dense computation and sparse memory operations, and between relational traversal and global search.The paper presents these as new design degrees of freedom whose optimal balance remains unresolved.