RAT-SQL: Relation-Aware Schema Encoding and Linking for Text-to-SQL Parsers (1911.04942v5)

Abstract: When translating natural language questions into SQL queries to answer questions from a database, contemporary semantic parsing models struggle to generalize to unseen database schemas. The generalization challenge lies in (a) encoding the database relations in an accessible way for the semantic parser, and (b) modeling alignment between database columns and their mentions in a given query. We present a unified framework, based on the relation-aware self-attention mechanism, to address schema encoding, schema linking, and feature representation within a text-to-SQL encoder. On the challenging Spider dataset this framework boosts the exact match accuracy to 57.2%, surpassing its best counterparts by 8.7% absolute improvement. Further augmented with BERT, it achieves the new state-of-the-art performance of 65.6% on the Spider leaderboard. In addition, we observe qualitative improvements in the model's understanding of schema linking and alignment. Our implementation will be open-sourced at https://github.com/Microsoft/rat-sql.

This paper introduces RAT-SQL, a framework designed to improve the performance of text-to-SQL models, particularly their ability to generalize to unseen database schemas. The core challenge addressed is twofold: effectively encoding the database schema's structure and accurately linking mentions in the natural language question to the corresponding columns and tables in the schema.

Challenges Addressed:

• Schema Encoding: Representing the database schema (tables, columns, primary keys, foreign keys) in a way that's accessible to the model. Prior methods like Graph Neural Networks (GNNs) encoded schema structure but often didn't contextualize it with the specific question.
• Schema Linking: Aligning potentially ambiguous references in the question (e.g., "model", "cars") to the correct schema elements (e.g., car_names.model vs. model_list.model, or linking "cars" to multiple tables involved in a join). This requires considering both schema structure and question context.

Core Idea: Relation-Aware Self-Attention (RAT)

The foundation of RAT-SQL is the Relation-Aware self-Attention (RAT) mechanism, adapted from Shaw et al. (Shaw et al., 2018 ). Standard self-attention (Transformer) calculates attention weights based solely on learned query-key interactions. RAT enhances this by incorporating explicit information about predefined relationships between input elements.

In a standard Transformer layer:

$e_{ij}^{(h)} = \frac{(\vect{x_i} W_Q^{(h)}) (\vect{x_j} W_K^{(h)})^\top}{\sqrt{d_z / H}}$

$\vect{z}_i^{(h)} = \sum_{j=1}^n \alpha_{ij}^{(h)} (\vect{x_j} W_V^{(h)})$

In a RAT layer, relation embeddings ($\vect{r_{ij}}$) representing the known relationship between element $i$ and element $j$ are added:

$e_{ij}^{(h)} = \frac{(\vect{x_i} W_Q^{(h)}) (\vect{x_j} W_K^{(h)} + \color{blue}{\vect{r_{ij}^K}})^\top}{\sqrt{d_z / H}}$

$\vect{z}_i^{(h)} = \sum_{j=1}^n \alpha_{ij}^{(h)} (\vect{x_j} W_V^{(h)} + \color{blue}{\vect{r_{ij}^V}})$

Here, $\vect{r_{ij}^K}$ and $\vect{r_{ij}^V}$ are learned embeddings specific to the type of relation between $x_i$ and $x_j$. If no predefined relation exists, they can be zero vectors or a generic "no relation" embedding.

RAT-SQL Framework Implementation

1. Input Representation:
   • The input consists of the natural language question words ($Q$), schema columns ($C$), and schema tables ($T$).
   • Initial embeddings are obtained for each word in the question, column names (prepended with type like text), and table names using GloVe or BERT.
   • Multi-word names (e.g., "car names") are processed with a BiLSTM to get a single initial vector for each column/table. Question words are also processed with a separate BiLSTM.
   • The set of all initial embeddings forms the input sequence $X = (\vect{c}_{1}^\text{init}, \dots, \vect{t}_{1}^\text{init}, \dots, \vect{q}_{1}^\text{init}, \dots)$.
2. Joint Encoding with RAT:
   • A stack of RAT layers processes the combined sequence $X$. Crucially, this allows all elements (question words, columns, tables) to attend to each other, creating jointly contextualized representations.
   • The $\vect{r_{ij}}$ embeddings encode various relationships:
     • Schema Structure Relations: Based on the database schema definition (see Table 1). Examples:
       • Foreign-Key-Col-F: Column $i$ is a foreign key referencing column $j$.
       • Primary-Key-F: Column $i$ is the primary key of table $j$.
       • Belongs-To-F: Column $i$ belongs to table $j$.
       • Same-Table: Columns $i$ and $j$ are in the same table.
       • Relations between tables based on foreign keys.
     • Schema Linking Relations (Name-Based): Capture textual overlap between question n-grams (length 1-5) and column/table names.
       • Question-Column-ExactMatch: An n-gram in the question exactly matches column $j$'s name.
       • Question-Table-PartialMatch: An n-gram is a subsequence of table $j$'s name. (Also includes NoMatch and corresponding Column/Table-Question relations).
     • Schema Linking Relations (Value-Based): Links a question word $q_i$ to a column $c_j$ if $q_i$ appears as a (part of a) value within column $c_j$ in the database.
       • Implementation: This requires querying the database: SELECT 1 FROM table WHERE column LIKE '%word%' LIMIT 1. This is done efficiently using DB indices during preprocessing. The model itself doesn't see the DB content, only the existence of the match. Relation type: Column-Value.
     • Auxiliary Relations: (Appendix A) Include relative positional embeddings between question words (Question-Dist-d), identity relations (Column-Identity), and generic type-pair relations (Column-Table, Table-Table, etc.).
3. Decoder:
   • Uses a standard grammar-based LSTM decoder (following Yin et al. (Bousso et al., 2017 )) that generates the SQL query's Abstract Syntax Tree (AST) depth-first.
   • Actions include ApplyRule (expand AST node based on grammar) and SelectColumn/SelectTable (choose a schema item for a leaf node).
   • The decoder state attends over the final contextualized embeddings ($\vect{c}_{i}, \vect{t}_{i}, \vect{q}_{i}$) from the RAT encoder to generate context vectors.
   • For SelectColumn/SelectTable, it uses an attention mechanism that incorporates an alignment matrix computed using relation-aware attention between the decoder state and the final column/table embeddings. This helps ground the selection in the relevant input elements.

Key Results and Findings:

• Performance: Achieved state-of-the-art results on the challenging Spider benchmark, significantly outperforming previous non-BERT methods (57.2% vs 48.5%) and BERT-based methods (65.6% with BERT vs 60.6%) at the time of publication.
• Ablation Studies: Showed that both the schema structure relations and the schema linking relations (name/value-based) provide significant performance gains over a baseline without them. Value-based linking provided a large boost.
• Alignment Visualization: The learned attention weights qualitatively demonstrate the model's ability to align question phrases (like "largest horsepower") to relevant schema columns (Horsepower) and implicitly resolve table references ("cars" aligning to both cars_data and car_names).
• Error Analysis: Revealed common errors include selecting wrong columns (schema linking challenge), missing WHERE clauses (semantic understanding), and producing semantically equivalent but syntactically different SQL. Oracle experiments showed errors are roughly split between incorrect structure and incorrect column/table selection.
• Consistency: Adding BERT improved the model's consistency in producing the same output (or same correctness) for paraphrased questions.

Practical Implications:

• Provides a unified and effective way to inject structured knowledge (schema relations, name matches, value matches) into a Transformer-based encoder for text-to-SQL.
• The relation-aware mechanism is general and could be applied to other tasks involving structured inputs alongside text.
• Value-based linking, implemented via efficient database lookups, offers a practical way to leverage database content without exposing the full database to the model or incurring high computational costs during inference.
• The code was open-sourced, facilitating replication and extension.

Implementation Considerations:

• Computational Cost: Stacking multiple self-attention layers can be computationally intensive, especially with long input sequences (question + all columns + all tables). BERT further adds to this.
• Preprocessing: Requires schema analysis to extract relations and potentially database queries for value-based linking.
• Hyperparameters: Requires tuning of standard Transformer parameters plus RAT-specific aspects like the number of layers and potentially the dimensionality of relation embeddings.
• Value Decoding: While effective on Spider (which doesn't require predicting specific values in the SQL query), the paper notes the value decoding mechanism used for WikiSQL experiments was basic and needed improvement for tasks requiring value prediction.