AST-Driven Sub-SQL Augmentation

• AST-driven Sub-SQL augmentation is a method that uses SQL Abstract Syntax Trees to systematically decompose queries and generate negative examples for robust model training.
• LearNAT employs Monte Carlo Tree Search to decompose natural language queries into sub-SQL fragments, while HeroSQL applies AST mutations to derive semantically altered negatives.
• Empirical results demonstrate significant improvements in NL2SQL performance, with gains up to 16 pp in AUPRC and enhanced semantic validation across models.

AST-driven Sub-SQL augmentation encompasses a class of techniques at the intersection of program synthesis, robust semantic parsing, and diagnostic validation for Text-to-SQL systems. The central idea is to exploit the compositional and hierarchical structure of SQL queries—formalized as Abstract Syntax Trees (ASTs)—to either (a) guide the systematic decomposition and merging of sub-query fragments or (b) generate fine-grained, syntactically valid but semantically incorrect negative examples for model training and evaluation. This approach leverages tree-theoretic representations to both enhance decompositional reasoning in generation (as in NL2SQL decomposition via LearNAT) and amplify discriminative power in semantic validation settings (as in HeroSQL).

1. Formalization of ASTs and Sub-SQL Fragments

In this paradigm, the AST of an SQL query, denoted $AT(Y) = (N, E)$, encodes SQL as a directed acyclic graph with nodes $N$ and edges $E$. Nodes are partitioned into clause nodes ($N_c$, e.g., SELECT, WHERE), operator nodes ($N_o$, e.g., AND, OR, $>$), and operand nodes ($N_v$, e.g., table/column names and constants). Parent-to-child edges capture grammatical relationships, with the root usually being SELECT (Liao et al., 3 Apr 2025). Sub-SQLs correspond to subtrees of $AT(Y)$; formal subtree relations $$isSubtree(AT_1, AT_2) = 1 \Leftrightarrow N_1 \subseteq N_2, E_1 \subseteq E_2$$ underpin validity checks in decomposition and augmentation procedures.

This structure enables precise targeting of logical fragments—for example, isolating the WHERE clause or a specific aggregation for mutation or extraction. In HeroSQL, each AST node can be explicitly aligned with a Logical Plan (LP) node, producing a strong mapping between sub-SQL semantics and tree structure (Qiu et al., 28 Dec 2025).

2. Methodologies: Generation, Augmentation, and Validation

LearNAT: AST-Guided Task Decomposition for Sub-SQL Generation

The LearNAT framework leverages ASTs to decompose natural language queries into sequences of sub-SQL generation tasks. Subtask selection, expansion, and pruning are controlled during Monte Carlo Tree Search (MCTS), where each search state contains the cumulative merged AST and a reward based on AST similarity to the gold SQL's AST. Actions are classified as progressive, redundant, or invalid via subtree relations, enabling aggressive pruning (Liao et al., 3 Apr 2025).

Reward is computed as a convex combination of node-level overlap and tree edit distance:

• Node-level similarity:

$$sim_{node}(AT_1, AT_2) = \sum_{t \in \{c, o, v\}} w_t \cdot \frac{|N_t(AT_1) \cap N_t(AT_2)|}{|N_t(AT_1) \cup N_t(AT_2)|}$$

• Structural similarity:

$N$0

Margins derived from these rewards inform fine-grained Reinforcement Learning via DPO objectives.

HeroSQL: AST-Driven Sub-SQL Augmentation for Negative Mining

HeroSQL employs AST-level transformations to automatically generate negative SQL examples at sub-query granularity (Qiu et al., 28 Dec 2025). The process involves:

• Parsing gold SQLs into AST and LP structures.
• Sampling AST nodes and applying transformations such as operator inversion, identifier substitution, constant replacement, and aggregation mutation.
• Filtering mutated queries by requiring different execution results from the original.
• Labeling each negative example with the specific LP node (sub-SQL fragment) affected.

This enables syntactically valid but semantically incorrect negatives, supporting robust training and evaluations of both query-level and sub-query-level semantic validators. The final augmented dataset balances positive and negative ratios (typically 1:1) for optimal discriminative learning.

3. Algorithmic Procedures and Implementation Logic

Both generation and augmentation protocols are formalizable with precise pseudocode.

LearNAT Decomposition (Generation):

$N$4 (Liao et al., 3 Apr 2025)

HeroSQL AST Augmentation:

$N$5 (Qiu et al., 28 Dec 2025)

Perturbations used include operator inversion, identifier substitution, constant replacement, and aggregation mutation, with constraints to maintain syntactic correctness and meaningful semantic drift.

4. Integration into Representation Learning and Training

In systems such as HeroSQL, the augmented dataset with sub-SQL-labeled negatives is fed into a hierarchical encoding and detection pipeline:

• Each SQL (correct or perturbed) is parsed into an LP graph, then each LP node is parsed into an AST.
• Node embeddings are produced using LLM-based schema-aware models.
• A nested message passing neural network (NMPNN) propagates information: a lower-level MPNN aggregates ASTs per LP node, and a higher-level MPNN aggregates over the LP structure.
• Outputs are pooled, fused with question embeddings, and scored for semantic correctness.

The loss is a query-level binary cross-entropy objective, with potential for extension to sub-query node findings. Negative-to-positive sampling ratios and network depth hyperparameters are empirically optimized, with best performance at a 1:1 negative-to-positive ratio and two-layer MPNN depth (Qiu et al., 28 Dec 2025).

5. Empirical Effects and Impact on NL2SQL Systems

Empirical studies demonstrate that AST-driven sub-SQL augmentation substantially improves the robustness and granularity of Text-to-SQL models:

• In HeroSQL, removing the AST-driven negative augmentation caused a performance drop of 4.7 pp in AUPRC and 2.3 pp in AUROC on BIRD with Qwen3-0.6B (Qiu et al., 28 Dec 2025).
• Consistent gains (4–16 pp AUPRC in ablations) are attributed solely to this augmentation, confirming its role as a principal driver of semantic validation performance.
• For generation, the LearNAT framework enables open-source 7B-parameter LLMs to match GPT-4 on complex NL2SQL datasets, rationalized by the systematic coverage of the gold AST throughout multi-step decomposition and recomposition (Liao et al., 3 Apr 2025).

6. Demonstrative Workflow and Practical Considerations

A stepwise illustration, using LearNAT, is as follows (Liao et al., 3 Apr 2025):

• The NL input, e.g., “Find the average rating of movies released after 2010 and count how many distinct directors they have,” is parsed for gold SQL and AST.
• The first subtask (extract average rating) yields a SQL and AST subtree; its reward reflects partial overlap (e.g., $N$1).
• The next subtask (distinct director count) merges additional AST components; accumulated coverage of the AST drives $N$2, signifying full coverage and a complete, correct query.
• Each decomposed and recomposed step is automatically validated, with errorful branches pruned or stored as negative (contrastive) examples for reinforcement learning.

In HeroSQL's augmentation, filtered mutations ensure only execution-changing negatives are retained, each precisely attributed to the responsible sub-SQL fragment. This enables fine-grained error localization and systematic improvement of semantic validators.

7. Adaptive Demonstration and Self-Improving Pools

LearNAT employs adaptive demonstration selection by embedding training questions and, at inference, retrieving the top-$N$3 nearest neighbors for in-context demonstration. Simultaneously, after every training synthesis round, a “demo pool” is updated with successful decomposition traces, where tasks with high AST similarity to current queries are prioritized in future rounds (Liao et al., 3 Apr 2025). This mechanism is self-improving: successful decompositions propagate as new demonstrations, providing continual bootstrapping of sub-SQL augmentation capabilities.

A similar principle applies in HeroSQL: automatic labeling of sub-query errors allows refined feedback loops during model evolution.

---

In summary, AST-driven sub-SQL augmentation—via systematic decomposition (LearNAT) or strategic negative mining (HeroSQL)—exploits the fine structure of SQL’s AST to both guide LLMs through compositional reasoning and provide densely annotated, syntactically correct negative examples for discriminative modeling. The result is a measurable, architecture-agnostic boost in both generative and validation performance on complex Text-to-SQL tasks, underpinned by interpretable tree- and fragment-level representations (Liao et al., 3 Apr 2025, Qiu et al., 28 Dec 2025).