SNACK: A Sequential Notation Framework for Probabilistic Graph Generation

AI summary

Overview

• Research area: Machine learning for generative modeling of graphs, spanning string-based graph representations, autoregressive sequence models, and Generative Flow Networks (GFlowNets), with molecular design as the main application domain.
• Technical level: Intermediate. The core encoding idea is easy to grasp, but the GFlowNet training section and the automorphism-based treatment of symmetry assume familiarity with probabilistic modeling.
• Scope: The paper introduces a graph-to-string encoding scheme, proves it is bijective with adjacency matrices, and evaluates it on generic graph generation, molecular distribution learning, goal-directed molecular optimization, and probing of learned representations.

What This Paper Is About

Generating graphs with discrete node and edge features can be done either with graph-native models (which are expensive and need specialized architectures) or by turning graphs into strings and using standard sequence models (which are efficient but typically cannot guarantee valid outputs or handle graph symmetries). SNACK proposes a string format that linearizes the nonzero entries of a graph's lower-triangular adjacency matrix, so that each valid sequence maps to exactly one ordered adjacency matrix. The goal is to combine the efficiency of sequence models with the structural guarantees normally associated with graph-based generative models.

Key Contributions

1. A bijective string representation. SNACK encodes only the nonzero entries of the lower-triangular adjacency matrix, including the diagonal, using exactly N + 2M tokens for an attributed graph with N nodes and M edges. A proposition (with proof in the appendix) establishes that the SNACK representation is bijective with respect to the adjacency matrix — a property the authors state they are the first to explicitly formalize and realize in a string-based representation.
2. String-based GFlowNet training with symmetry correction. Because multiple node orderings can yield the same adjacency matrix, the paper derives q(x|G) by summing an ordering policy over the automorphism group Aut(A), giving |Aut(A)| · q^ord(π|G). This enables Trajectory Balance training on strings, plus a conditioned TB loss for generating around fixed substructures.
3. Invalid logit masking for guaranteed validity. Prefix-decodability lets the model inspect the partial graph at each step and mask tokens that violate grammar constraints, chemical valence, or BFS ordering. These masks enforce structural constraints before token selection rather than filtering outputs afterward.
4. Empirical validation and probing. SNACK is evaluated on generic graph generation, molecular distribution learning, conditioned generation, and reward-driven optimization, and bilinear probing is used to test whether the trained model's hidden states encode adjacency and shortest-path-distance structure.

Main Findings

• Generic graph generation: On eight graph datasets, SNACK matches training-distribution statistics using one shared hyperparameter set with no per-dataset tuning. On Planar it records MMD scores of 0.001 (degree), 0.005 (clustering), and 0.003 (orbit); on Enzymes 0.003, 0.013, and 0.001; on SBM 0.018, 0.013, and 0.039; and on Grid 0.000, 0.000, and 0.000.
• Molecular distribution learning: With grammar, valence, and BFS masking (V+B), SNACK reaches 1.000 validity on both QM9 and MOSES. On QM9 it records FCD 0.037, SNN 0.627, Frag 0.996, and Scaf 0.977; on MOSES, FCD 0.103, SNN 0.518, Frag 1.000, and Scaf 0.895. Without any masking it still reaches 0.998 validity on QM9 and 0.996 on MOSES.
• Substructure-conditioned generation: Using 129 of the smallest MOSES test molecules and 100 random input sequences each, SNACK achieves a success rate of 0.390, compared with 0.181 for SELFIES, 0.152 for SMILES, 0.057 for SAFE, 0.003 for STGG, and 0.000 for both GEEL and G2PT.
• Goal-directed optimization: SNACK (V+B) is competitive with a graph-based GFlowNet on the HOMO–LUMO gap task (max N = 9) and substantially better on the two similarity tasks (max N = 16). On albuterol it reaches an average reward of 0.529, Top 100 reward of 0.736, and diversity of 0.930, versus 0.116, 0.171, and 0.960 for the graph baseline. On aripiprazole it reaches 0.292 average reward and 0.469 Top 100 reward. SNACK attains a unique fraction of 1.000 on all three tasks, while the graph-based GFlowNet does not.
• Conditioned scaffold optimization: On optimizing around one of 800 molecules with low penalized logP, SNACK reports an improvement of 6.45 ± 1.6 and similarity of 0.46 ± .15 with 100% success, compared with GraphDF (5.93 ± 0.8, 0.30 ± .1, 100%), GCPN (4.20 ± 1.3, 0.32 ± .1, 100%), and JT-VAE (1.91 ± 2.0, 0.28 ± .1, 97.5%).
• Scalability: Graph-based generation requires N + M forward passes to score a trajectory, whereas SNACK computes trajectory-level log-probabilities in a single forward pass. With dense attention in both cases, this reduces computational and memory complexity from cubic to quadratic in the number of nodes, and SNACK-GFlowNet shows higher training throughput as molecular size grows.
• BFS ordering beats alternatives: Among uniform, CM, and BFS node orderings, BFS gives the best FCD on both datasets — 0.037 on QM9 and 0.103 on MOSES, versus 0.044 and 0.113 for CM, and 0.203 and 6.513 for uniform ordering. The authors note this contrasts with prior work favoring CM ordering.
• Edge-ordering design matters: Moving connectivity tokens before the node token (as in GEEL and many graph-based methods) consistently increased validation loss on QM9, supporting SNACK's current block design that lists distances and edge types before the node token.
• Probing reveals internalized structure: A bilinear probe on a SNACK (V+B) model trained on MOSES recovers adjacency matrices with AUROC 0.9998 and accuracy 0.9976, and shortest-path distances with MAE 0.9250 and MSE 1.6659. A randomly initialized SNACK model reaches AUROC 0.9362 and MAE 1.2897; an SMILES model reaches AUROC 0.9970 and MAE 0.8148, indicating SMILES better preserves long-range distances while SNACK better captures local connectivity.
• Automorphism cost is practical: Although computing |Aut(A)| is asymptotically expensive, the paper states that the nauty package computes it in a fraction of a millisecond for the graph sizes used in existing generative models, and notes this cost is also inherent to graph-based GFlowNets.

Methodology in Plain English

The researchers start from the observation that a graph's adjacency matrix, restricted to its lower triangle, fully describes the graph. They read that triangle row by row, writing down only the non-zero entries: each edge becomes a pair of tokens (a distance to the earlier neighbor, then the edge type), and each node's own label appears as its diagonal entry. The example in the paper is [C]1=[C]2-[C]2-[N]2=1-[C]1-[O].

Two design decisions follow. First, node ordering is chosen to keep distances small — a bandwidth restriction — using random breadth-first search, which shrinks the distance vocabulary. Second, since any prefix decodes into a partial graph, the model can check at every step which next tokens would be illegal and set their probabilities to zero. Three mask types are used: grammar (syntax and termination), valence (chemical validity), and BFS (ordering consistency). Masking matters because training on BFS-ordered strings does not by itself make the model generate in BFS order.

For distribution learning, the model is a standard autoregressive Transformer trained with next-token negative log-likelihood, which is a variational lower bound on the graph log-likelihood under the ordering distribution q(x|G). Because SNACK is one-to-one with adjacency matrices, q(x|G) can be written exactly as |Aut(A)| times the probability the ordering policy assigns to a single canonical ordering. This exact quantity is what makes GFlowNet training possible: the Trajectory Balance objective compares forward trajectory probability against q(x|G) times the reward, so reward is spread correctly across all valid strings for a graph. The authors let the model generate arbitrary fragments rather than only molecules, which means the same machinery handles conditioned optimization. Finally, they train bilinear probes on hidden states at node-token positions to check whether structure is present inside the model without any explicit graph supervision.

Why This Matters

The paper argues that string-based graph generation has been held back by three things at once: no guarantee of valid output, difficulty handling graph symmetry in probabilistic objectives, and poor support for goal-directed search such as GFlowNets. SNACK addresses all three with one representation, and the authors show this is not merely a convenience — it changes what is computationally feasible as graphs grow, since scoring a full trajectory becomes a single forward pass instead of N + M.

Real-world applications:

• Drug discovery and lead optimization: The paper explicitly frames scaffold-constrained generation as critical for lead optimization, and shows SNACK can optimize properties while strictly preserving a required fragment.
• Molecular property optimization: Generating molecules with target HOMO–LUMO gaps or high similarity to reference molecules such as albuterol and aripiprazole.
• Graph-structured scientific data: The generic benchmarks include Planar, Enzymes, SBM, and Grid graphs, pointing to modeling of structured networks beyond chemistry where no per-dataset tuning is available.
• Generative model and pipeline engineering: Because SNACK runs on standard LLaMA-style Transformer blocks and a single forward pass, it lets practitioners reuse sequence-model infrastructure instead of building graph-native architectures.

Industry relevance: the immediate beneficiary is computational chemistry and biotech, where the combination of guaranteed chemical validity and faster GFlowNet training directly affects how many candidate molecules can be explored. The secondary impact is on ML infrastructure — replacing specialized graph neural networks and per-state forward passes with ordinary sequence stacks lowers engineering and memory costs.

Future Directions

• Learning the node ordering policy jointly with the model, since the current ordering policy is fixed and manually designed — the authors name this as a limitation.
• Extending SNACK to represent stereochemistry and chirality, which the current encoding does not cover.
• Adapting the representation for fragment-level rather than node/edge-level generation, to capture higher-level structural patterns.
• Revisiting ordering comparisons more broadly, since the paper's finding that BFS beats CM ordering contradicts prior studies and the authors suggest the earlier conclusions may be dataset-dependent or based on metrics that do not fully capture generative quality.

Target Audience

Researchers and practitioners working on generative models for graphs and molecules, especially those already using or considering GFlowNets and looking for a sequence-based alternative to graph-native architectures. It is also relevant to anyone interested in string linearizations of graphs (SMILES, SELFIES, SAFE, STGG, GEEL, G2PT), in representation probing, or in the general question of whether autoregressive sequence training can internalize structured, non-sequential data.

Authors’ abstract

We introduce SNACK, a string representation for probabilistic graph generation that enables sequence models to generate and model graphs with discrete node and edge features. SNACK linearizes graphs by encoding the nonzero entries of the lower-triangular adjacency matrix, establishing a bijective correspondence between valid sequences and ordered adjacency matrices. This design enables invalid logit masking to enforce structural constraints--such as chemical valence and ordering rules--guaranteeing 100% valid generated sequences. SNACK also enables principled string-based GFlowNet training by explicitly accounting for node-ordering symmetries. Across standard graph-generation tasks and molecular distribution-learning benchmarks, SNACK achieves competitive sample quality, while substantially improving the throughput of GFlowNet training relative to graph-based generation. Probing analyses further show that models trained on SNACK encode meaningful local and global graph structure despite receiving only sequential supervision. These results position SNACK as a practical framework for probabilistic graph generation with deep sequence models.