Variable-sized K-ary Tree Modeling via Discrete Flow Matching

Tree growth — sampling trajectory — Sampling trajectory — a maze (spanning tree) materializes step by step from the root-only reference under the edit-flow CTMC.

Motivation

Tree-structured data is everywhere — code ASTs, CSG solid-modeling trees, parse trees — yet existing generative models do not handle them well. Graph-diffusion methods such as DeFoG achieve strong results on general graphs but pay O(n²) complexity for the adjacency matrix and cannot natively express tree-specific constraints (acyclicity, single-parent rule). Tree-diffusion methods, on the other hand, rely on hand-crafted CFG rules and are typically bounded to a fixed depth or domain. Variable-length sequence flow matching (Edit Flows, NeurIPS 2025) shows that edit operations can model length variability — but only in 1D.

Goal

Construct a generative framework that produces variable-sized rooted K-ary trees while keeping implicit structural rules (acyclic, single root, valid parent–child) satisfied throughout the entire generation trajectory.

Method

Discrete Flow Matching on trees. The interpolant between a root-only reference tree and a target tree is projected onto the manifold of valid trees at every t. The projection cascades deletions of unreachable subtrees, keeping state-space complexity at O(kN) instead of k^d.

Edit-flow CTMC over tree slots. Three operations act on the tree's empty / non-empty slots:

Ins(x, i, r, c) — insert a node of type c into an empty slot
Del(x, i, r) — delete a node and cascade-remove its descendants
Sub(x, i, r, c) — substitute a node's type

The network outputs per-slot OP rates (λ_ins, λ_del, λ_sub) and per-operation type logits (π_ins, π_sub), which together define the generating velocity in a Continuous-Time Markov Chain that is well-defined only on valid trees — a semi-autoregressive scheme.

Training loss. A Bregman-KL objective on the rates with three branches (ins / del / sub) keeps the velocity field consistent with the projected interpolant.

Sampling. Tau-leaping sampler — at each step, the number of events per slot is drawn from a Poisson process with rate λΔt, giving event probability 1 − e^{−λΔt}.

Experiment — Maze Generation

The framework is evaluated on maze generation: a maze is a spanning tree of a grid, so the task probes whether the model genuinely learns the tree manifold rather than just plausible-looking graphs. Comparison is against DeFoG (graph-level discrete flow matching) and an Insertion-only ablation that disables the deletion and substitution operations.

Qualitative comparison of maze samples — Qualitative comparison. Full-Editing (ours) produces dense, well-formed mazes; the Insertion-only ablation lacks the corrective capacity of deletion/substitution; DeFoG samples often fail to be acyclic at all.

Quantitative comparison table — Quantitative comparison over 100 generations. Our method produces 5× more perfect samples than DeFoG and is 100% acyclic — DeFoG is acyclic only 44/100 of the time, illustrating the value of operating directly on the tree manifold.

The clearest signal is in acyclicity: 100/100 valid trees from our method versus 44/100 for DeFoG — direct evidence that the projection step plus edit-based CTMC gives the structural guarantees that adjacency-matrix diffusion does not. Perfect samples are also markedly higher (10 vs 6 vs 2).

Takeaways

Designed a non-trivial extension of discrete flow matching to a structured (non-1D, non-fixed-size) state space
Engineered the projection step that makes the CTMC well-defined on the tree manifold
Derived and implemented the rate-based Bregman-KL training loss and the tau-leaping sampler