Beyond Flat Walks:
Compositional Abstraction for Autoregressive Graph Generation

A flat walk hides the structure a graph is built from

To generate graphs with a language model, you first have to write a graph down as a sentence. The standard sentence is a flat walk: start at a node, step from neighbor to neighbor, and whenever the walk returns to a node it has already written, record a bracket of back-edges. This is lossless and efficient, and on small molecules it works well. But nothing in the walk knows what a ring, a functional group, or a community is. That structure is left implicit in the token stream, and the model has to rediscover it on its own.

We ask whether the sentence itself can carry the structure. Everything else is held fixed: the same causal transformer, the same training recipe, and matched training time. Only the way the graph is written down changes. The figure below plays the two serializations of one molecule side by side.

The intuition is coarse before fine. Every hierarchical tokenizer we study emits the higher-level community structure before the atom-level detail inside it, so when the model has to place an atom, it has already written which community that atom belongs to. A fine-to-coarse order would defeat the purpose: the model would have to commit to local structure before seeing the organization meant to guide it. The cost is length. Writing the hierarchy down takes tokens, and HDT sequences are roughly twice as long as SENT’s, which is why every comparison in the paper matches wall-clock training time rather than steps or epochs.

The pipeline: coarsen, serialize, predict, decode

The framework splits generation into stages that can be studied separately. Coarsening builds a partition tree $\mathcal{H}$ from the input graph G. Tokenization serializes $\mathcal{H}$ into a flat sequence in coarse-before-fine order. A standard causal transformer learns that sequence by next-token prediction, and decoding inverts the tokenization exactly, so any syntactically valid sample becomes a graph. All of our contributions concern the serialization; the model is an ordinary GPT-2.

Decoding is exact for the encoder–decoder pair: decode ∘ tokenize ∘ coarsen(G) ≅ G, because the partition covers every node and each edge is written exactly once. That is a property of the tokenizer, not a promise about samples: a generated sequence may be malformed, and then it is counted as invalid rather than repaired. Validity is measured for every method, including the flat baseline.

Three design axes

Separating the stages exposes three independent choices: how the graph is coarsened, how the hierarchy is written into tokens, and how the traversal order is randomized during training. The figures and table below run each axis on real molecules.

How the graph is coarsened

We use one domain-agnostic strategy everywhere and one chemistry-aware strategy on molecules. HAC, hierarchical agglomerative clustering with an adjacency constraint, merges clusters bottom-up only when they share an edge, yielding tight communities of controlled size for any graph. Motif Community (MC) detects ring systems with SMARTS queries, merges overlapping rings, and absorbs leftover atoms where it can. Recursion continues until communities fall below a small threshold (τ = 4), so $\mathcal{H}$ is shallow on drug-like molecules (mean depth 1.13 on MOSES) and deeper on large ones (1.86 on Macrocycle).

How the hierarchy is written into tokens

Given $\mathcal{H}$, the question is how explicit the hierarchy should be in the token stream. HSENT declares it: a catalogue of community blocks followed by a separate bipartite section that wires them together. HDTC is a typed variant for molecules that labels each block and appends a super-graph of inter-community bonds. HDT, the tokenizer we use throughout, traverses it: one depth-first pass in which cross-community bonds are ordinary back-edges, syntactically identical to bonds inside a community.

Declaring the hierarchy did not pay off. On GuacaMol and MOSES, both HSENT configurations have lower validity and V.U.N. and higher MoleculePGD than the corresponding HDT configurations, and HDTC trails HDT on the same metrics, while HSENT sequences run about 2.9× SENT’s length against 1.9× for HDT. One plausible reading: a declarative grammar poses a second prediction problem. A sample can declare a set of communities and then wire them in a way that contradicts the declaration. In HDT there is no separate wiring record to contradict; community membership lives entirely in the [ENTER] / [EXIT] nesting.

How the order is randomized

A graph admits many valid serializations, and resampling the traversal at every training presentation is free data augmentation. A flat random walk gets this for free: SENT draws a new walk each epoch. A deterministic traversal shows the model one string per graph. We compare two policies that keep the graph and $\mathcal{H}$ fixed. Within-community randomization (-R) keeps the stored community order and randomizes the traversal inside each terminal community. Hierarchy-wide randomization (-FR) also samples an adjacency-constrained DFS over communities at every level, entering each community through an atom adjacent to the one it was reached from.

Hierarchy improves distributional fidelity on every molecular benchmark

We train the same GPT-2 XS backbone (≈ 11M parameters) on five molecular benchmarks (GuacaMol, MOSES, COCONUT 2.0, a disjoint split of GEOM-Drugs, and Macrocycle) with one training recipe and matched training time, and compare the flat baseline with the two hierarchical configurations fixed in advance, HDT-HAC-R and HDT-HAC-FR.

The gap is largest, and most consistent, on the metrics that compare distributions rather than count successes. Against SENT, HDT-HAC-R lowers FCD by 51.6% to 64.6% and MoleculePGD by 34.3% to 86.6% on all five benchmarks, and improves SNN, fragment similarity and scaffold similarity on all five as well. The margins are large relative to the seed spread: on GuacaMol, FCD falls from 1.799 ± 0.045 to 0.871 ± 0.019, and on Macrocycle, MoleculePGD falls from 0.639 ± 0.003 to 0.085 ± 0.002. All four hierarchical configurations have lower FCD and MoleculePGD than SENT on every benchmark, and the improvements persist when metrics are recomputed only on novel molecules, so they are not explained by reproducing training molecules. SENT keeps a slightly higher internal diversity on all five benchmarks, by less than 0.02.

Randomizing across the hierarchy closes most of the V.U.N. gap

Fidelity is not the whole story. Switch the chart above to V.U.N. (the fraction of samples that are valid, unique and novel). With within-community randomization, HDT-HAC-R trails SENT by 3.6% on GuacaMol, 4.0% on MOSES, 17.0% on COCONUT 2.0 and 37.2% on Macrocycle, where novelty falls from SENT’s 0.999 to 0.658: the model reproduces training molecules more often.

Randomizing the traversal across the whole hierarchy largely removes this cost. Across the two fully randomized configurations the largest V.U.N. deficit is 10.85%, and for HDT-MC-FR it is 2.98%. HDT-MC-FR still improves FCD, MoleculePGD, SNN and fragment similarity over SENT on all five benchmarks, and has slightly higher mean V.U.N. on GuacaMol. A residual deficit remains on Macrocycle (V.U.N. 0.762 for HDT-HAC-FR against SENT’s 0.854), carried by validity rather than novelty. How the hierarchy is presented during training matters as much as whether it is present.

Beyond molecules: hierarchy helps when it matches the data

The framework only needs a coarsening, so it applies to any graph. On four non-molecular benchmarks, with the same domain-agnostic HAC, the picture is no longer uniform, and that is informative. The figure lets you feel why: turn up the community gravity and see whether the coarsening carves the graph into clean islands or into arbitrary pieces.

On Proteins, HDT-HAC-R improves V.U.N. by 62.8% over SENT and reaches the lowest MMD Ratio among the methods in the paper’s literature table, so the benefit does not depend on a chemical catalogue. On Sparse SBM it improves V.U.N. by 59.3% and cuts Degree MMD, Orbit MMD and Ratio by more than 90%. On Planar, where there is no community structure to find, both hierarchical configurations are worse than the flat walk on every graph MMD.

When the hierarchy misleads

The standard SBM is the instructive failure. Default HAC recovers the planted blocks well at the top level, but keeps splitting: communities averaging about 31 nodes are recursively cut into leaves averaging 2.25 nodes. The resulting model reaches validity 0.926, yet its uniqueness and novelty fall to about 0.41 and its V.U.N. is only 0.276, against SENT’s 0.669. High validity alone does not establish successful generation.

Two fixes confirm the diagnosis. Raising the stopping threshold to τ = 44 raises V.U.N. to 0.709 and 0.694 for -R and -FR; the planted hierarchy reaches 0.887 and 0.848. On connected Recursive SBM, HAC with τ = 10 reaches V.U.N.† 0.854 under -R, close to the planted hierarchy’s 0.858, whereas -FR stays far below the planted -FR (0.083 against 0.744). The benefit comes from alignment between the hierarchy and the data, and how much of it survives depends on the randomization policy.

In context

For orientation, here are the controlled runs next to published numbers. Rows marked ∗ are transcribed as published; budgets and evaluation lineages differ, so this is placement rather than a superiority claim. The controlled comparison is the one above.

Generated samples

Numbers summarize distributions; the samples show what they look like. Below are the unconditional-generation samples from the paper for every benchmark: held-out references, the flat SENT baseline, and our two hierarchical configurations, with each model’s headline numbers beside its row.

Takeaways

Limitations

Hierarchy has a computational cost. HDT’s lossless coarse-before-fine encoding lengthens sequences by roughly 2× over SENT on small molecules and more on large graphs, so the same backbone spends roughly 4× the attention memory and compute; at matched training time the hierarchical variants take between a third and two-thirds of SENT’s optimizer updates. The benefit also depends on the coarsening: a coarsening has to be assessed for each domain, and hierarchies do not transfer across datasets. Finally, a residual V.U.N. deficit remains on Macrocycle even with full randomization.

Citation

@inproceedings{bian2026beyond,
  title     = {{Beyond Flat Walks: Compositional Abstraction for Autoregressive Graph Generation}},
  author    = {Bian, Kaiwen and Yang, Andrew H. and Parviz, Ali and Mishne, Gal and Wang, Yusu},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}