Shubhashis Roy Dipta

<p dir="ltr">Text2SQL systems promise to democratize access to enterprise data but often fail to handle the complexity of real-world database queries, even if they perform well on test datasets. We found that Reinforcement Learning from Human Feedback (RLHF) is a viable approach for active learning from incorrect production queries to improve Text2SQL accuracy.</p>
<p dir="ltr">Using Mistral&rsquo;s <a href="https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1">mixtral-8x7b-instruct</a> model and the <a href="https://bird-bench.github.io/">BIRD</a> dataset, we demonstrate that RLHF with a hybrid reward model&mdash;which combines database execution and LLM evaluation&mdash;can improve SQL generation accuracy. Our approach achieves better performance than traditional supervised fine-tuning across most complexity levels by using 75% of training data for fine-tuning and 25% for RLHF with augmentation.</p>
<p dir="ltr">In this blog post, we share our methodology, the effectiveness of&nbsp;<a href="https://arxiv.org/abs/2402.01306">Online Kahneman-Tversky Optimization</a> (OKTO), and practical insights for leveraging RLHF for Text2SQL tasks.</p>
<h2 dir="ltr">Text2SQL systems face two key challenges that impact accuracy: data scarcity and data distribution shift</h2>
<h3>Data scarcity</h3>
<p dir="ltr">Whereas current LLM models are trained over a trillion tokens (15T for <a href="https://ai.meta.com/blog/meta-llama-3/">Llama3</a>), the largest Text2SQL dataset (<a href="https://bird-bench.github.io/">BIRD</a>) has 80k query-SQL pairs (~2M tokens). Developing a large Text2SQL dataset is expensive and time-consuming, unlike a text-generation dataset, which can be scraped from the internet. Even though the positive text-SQL pair is hard to annotate, the negative data from incorrectly generated SQL is practically unlimited. RLHF can help us to use that negative data to align the model with the positive pairs.</p>
<h3 dir="ltr">Data distribution shift&nbsp;</h3>
<p dir="ltr">The second issue is that SQL from real-life databases are much different from academic datasets. For example, consider the largest and most well-known SQL dataset: BIRD.&nbsp;<br><br>Below is one of the &ldquo;challenging&rdquo; SQL queries from the BIRD dataset:</p>
<pre class="language-markup"><code>SELECT T2.admfname1,
       T2.admlname1,
       T2.admfname2,
       T2.admlname2,
       T2.admfname3,
       T2.admlname3
FROM satscores AS T1
INNER JOIN schools AS T2
    ON T1.cds = T2.cdscode
ORDER BY T1.numge1500 DESC
LIMIT 1;</code></pre>
<p dir="ltr">And here is a randomly chosen, sanitized SQL query from a real-life database:</p>
<pre class="language-markup"><code>SELECT   u.name,
         Count(DISTINCT (t.reference_number))
FROM     table_a                    AS p,
         table_b                    AS prod,
         table_c                    AS pa,
         table_d                    AS aff,
         table_e                    AS u,
         table_f                    AS t
WHERE    u.user_id = aff.rvp_id
AND      prod.id = p.product_id
AND      p.ref_number = pa.ref_number
AND      pa.assignment_id = aff.affiliation_id
AND      t.transaction_id = aff.trx_code
AND      NOT prod.family_code ilike '%TYPE1%'
AND      NOT prod.family_code ilike '%TYPE2%'
AND      (
                  t.trx_type ilike '%Payment%'
         OR       t.trx_type ilike '%Cancellation%' )
AND      t.year = date_part('year', CURRENT_DATE) - 1
AND      date_part('quarter', t.date) = 3
AND      (
                  NOT p.parent_org ilike '%Company1%'
         OR       NOT p.parent_org ilike '%Company2%'
         OR       NOT p.parent_org ilike '%Company3%' )
AND      u.region IN ('Region1',
                     'Region2',
                     'Region3')
AND      u.name &lt;&gt; ''
GROUP BY 1
ORDER BY 2</code></pre>
<p dir="ltr">As we can see, enterprise SQL is much more complex, longer, and hard to generate. One of the use cases of RLHF is to make stylistic changes to a previously trained model&rsquo;s output. We hypothesized that we could force our model to generate more complex SQL through stylistic alignment that better represents real-life database SQL queries.</p>
<p dir="ltr">Motivated by the issues above, our goal was to investigate if the negative data could push the alignment in the right direction and generate complex SQL to answer complex questions.</p>
<h2 dir="ltr">We explored using RLHF to help with SQL alignment</h2>
<p dir="ltr">Previous papers have already shown that RLHF helps align text generation to follow instructions (<a href="https://arxiv.org/abs/2203.02155">Ouyang et al.</a>), to generate helpful and harmless texts (<a href="https://arxiv.org/abs/2204.05862">Bai et al.</a>), and to train general-purpose LLMs (<a href="https://arxiv.org/abs/2307.09288">Lllama2</a>,<a href="https://ai.meta.com/blog/meta-llama-3/"> Llama3</a>,<a href="https://openai.com/index/chatgpt/"> ChatGPT</a>). Our central hypothesis was that we could use RLHF to make stylistic changes in SQL generation and force the model to generate more valid semantically correct SQL. Additionally, since we have a limited number of positive text-SQL pairs, can we use negative data with RLHF to get an LLM to generate more valid and complex SQL queries?</p>
<h2 dir="ltr">Experimental Design</h2>
<h3 dir="ltr">The setup</h3>
<p dir="ltr">We used the <a href="https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1">mixtral-8x7b-instruct</a> models as the base model. To start, we fine-tuned this model using 75% of the <a href="https://bird-bench.github.io/">BIRD</a> training data. Next, we took the remaining 25% of the dataset and synthetically augmented it. This was done by randomly reordering the in-context table schemas. We then used this augmented 25% of the data for the RLHF algorithm.</p>
<p dir="ltr">We used the original dev and test data from the BIRD dataset during validation and testing. For the reward model, we used a combination of the DB_executor function and a general-purpose LLM. We implemented a dual-component reward model that combines a DB_executor function with a general-purpose Large Language Model (LLM). The DB_executor function processes SQL queries through the appropriate SQL engine and generates a binary reward based on successful query execution.&nbsp;</p>
<p dir="ltr">To obtain a more nuanced evaluation, we then pass four elements to a judge LLM (specifically, mixtral-8x22b-instruct): the generated SQL query, the ground truth SQL query, and the database query results from both queries. This judge LLM provides a comprehensive reward assessment beyond simple binary execution success.</p>
<h3 dir="ltr">Reward model</h3>
<p dir="ltr">One of the crucial parts of any RLHF algorithm is the reward model. Our first choice of the reward model for a Text2SQL is the DB_executor function. The issue is that the DB executor can only say if the generation is right or wrong, which gives very sparse rewards. The RLHF algorithm needs a partial reward to get the signal that the generation is going in the right direction; otherwise, there is not enough signal to improve the generation.&nbsp;</p>
<p dir="ltr">Conventionally, a good reward model has to follow two rules:&nbsp;</p>
<ol>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">It should correctly say if the SQL is good or bad</p>
</li>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">If a SQL is partially correct or on the right path to correct, it should return a partial reward (partially correct SQL). Training such a reward model for SQL is tricky because, to our knowledge, there is no preference dataset for SQLs.&nbsp;</p>
</li>
</ol>
<p dir="ltr">As a result, we decided to use a hybrid reward model where the generated SQL first goes to the DB_executor function. Then, the output/error from the DB execution goes to a general-purpose LLM to work as the final reward model. For the reward scoring, we used a custom prompt that mimics the professor-student role: the general purpose LLM acts as Professor and grades the generated SQL out of 50. The whole reward model pipeline is shown in Figure 1.</p>
<p dir="ltr"><img src="https://img.plasmic.app/img-optimizer/v1/img?src=7108b46697f300b03a8b96671c267ede.png&amp;f=webp&amp;q=75&amp;w=3840" alt="" width="1512" height="509"><br>Figure 1: Reward Model Pipeline</p>
<h3 dir="ltr">Main pipeline</h3>
<p dir="ltr">At Scale, we performed various experiments to find the optimal split or configuration to align the model with positive SQL data. We fine-tuned the base model on 75% of the train data and kept the other 25% for RLHF train data. We also synthetically augmented 25% of the RLHF train data. The reasoning behind both of these choices is described in the next section. Figure 2 shows the system diagram of the whole pipeline.</p>
<p dir="ltr"><img src="https://img.plasmic.app/img-optimizer/v1/img?src=6ef0ee475acd4c3d9094eb0e2e08504b.png&amp;f=webp&amp;q=75&amp;w=3840" alt="" width="1512" height="723">Figure 2: High-Level System Pipeline&nbsp;</p>
<h2 dir="ltr">Key Findings</h2>
<p>In this section, we will briefly go through our main findings:&nbsp;</p>
<ol>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">Online KTO is the best RL algorithm for Text2SQL&nbsp;</p>
</li>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">The holdout dataset for RLHF is necessary for providing enough signal for the policy model to learn&nbsp;</p>
</li>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">Augmentation of the holdout dataset helps the RLHF algorithms</p>
</li>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">Neither LLM nor DB execution is enough as a reward model on its own, but combined, they provide a good signal</p>
</li>
</ol>
<h3 dir="ltr">Online KTO works better than all other RLHF algorithms:</h3>
<p dir="ltr">We have tried out different RLHF algorithms - <a href="https://arxiv.org/abs/1707.06347">PPO</a>, <a href="https://arxiv.org/abs/2305.18290">DPO</a>, <a href="https://arxiv.org/abs/2402.01306">KTO</a>, and both their online (<a href="https://arxiv.org/abs/2405.08448v1">ref</a>) and offline versions. In the following figure, we can see the results for OKTO (Online KTO), ODPO (Online DPO), PPO (Offline), and Reinforce (Offline) RLHF algorithms.</p>
<p dir="ltr"><img src="https://img.plasmic.app/img-optimizer/v1/img?src=f126f2fb80a259248528ade966424231.png&amp;f=webp&amp;q=75" alt="" width="1512" height="870">Figure 3: Results for OKTO (Online KTO), ODPO (Online DPO), PPO (Offline), and Reinforce (Offline) RLHF algorithms</p>
<p dir="ltr">The results show that Online KTO performs better than all others. The hypothesis is three-fold:&nbsp;</p>
<ol>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">As stated in the <a href="https://arxiv.org/abs/2402.01306">KTO</a> paper, KTO works better even with imbalanced data (negative data is 10x the positive data)</p>
</li>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">As stated in the <a href="https://arxiv.org/abs/2407.21783">Llama3</a> paper, DPO falls short when there are too many special tokens (SELECT, FROM, &hellip; in SQL)</p>
</li>
<li dir="ltr" aria-level="1">
<p dir="ltr" role="presentation">The offline RLHF algorithm limits learning to a fixed set of negative examples, while online RLHF enables dynamic learning through real-time SQL generation and reward feedback at each timestep t</p>
</li>
</ol>
<h3 dir="ltr">Holdout RLHF data is necessary</h3>
<p dir="ltr">The first question that comes to mind is why keep a holdout RLHF dataset? Why not use the whole dataset for fine-tuning and then the same dataset for the RLHF? We found out that doing so restricts any learning from the RLHF itself (shown in Figure 4). The main reasoning is that by fine-tuning, the model has already learned all of the query-SQL in the train set, so, during RLHF, it doesn&rsquo;t make any mistakes, hence gets perfect reward, hence again very sparse reward scores, which doesn&rsquo;t help in alignment.</p>
<p dir="ltr"><img src="https://img.plasmic.app/img-optimizer/v1/img?src=024839fbb4a4e43f0858b4a7634adb18.png&amp;f=webp&amp;q=75" alt="" width="1512" height="870">Figure 4: Test scores comparing SFT vs Online KTO, both on the full dataset</p>
<h3 dir="ltr">Augmenting the dataset helps</h3>
<p dir="ltr">As mentioned in the above system pipeline diagram, the RLHF was done on 25% augmented train data. In this experiment, we explored the importance of augmentation. We augmented the data by shuffling the database schema on the context.&nbsp;</p>
<p dir="ltr">So if for one question, Q1, the context sequence is C1, C2, C3, then on another instance, even if the question is Q1, the database schema can be C3, C1, C2. Our initial dev experiments revealed that 4x augmentation yields the most gain, so our initial 25% holdout dataset (~2k) has been augmented to ~8k query-SQL pairs.</p>
<p dir="ltr"><img src="https://img.plasmic.app/img-optimizer/v1/img?src=1d1ef1d2b4f92d1c7652dca7c78dadcf.png&amp;f=webp&amp;q=75" alt="" width="1512" height="870">Figure 5: Test scores comparing Online KTO vs Online KTO with augmentation</p>
<p dir="ltr">As we can see in the results, the augmented RLHF train data gave around a 2% boost in almost all the domains. The hypothesis is that shuffling the database schemas has two benefits: (1) More data/signal to train on, and (2) The model becomes more robust to the initial order of the database schemas and learns to retrieve the relevant DB schema.</p>
<h3 dir="ltr">Using only LLM is not enough as a reward model</h3>
<p dir="ltr">In this experiment, we wanted to explore whether we can remove the database execution overhead and use only the general-purpose LLM as a reward model. The main motivation behind this was that DB execution is slow and hard to parallelize during multi-node training. In this experiment, the reward model (LLM) only gets the ground truth and predicted SQL to evaluate.</p>
<p dir="ltr"><img src="https://img.plasmic.app/img-optimizer/v1/img?src=1ac163e81b65a4dde8761a0cc80f9cfc.png&amp;f=webp&amp;q=75" alt="" width="1512" height="870">Figure 6: Test scores comparing Online KTO (LLM + DB as reward function) vs Online KTO (Only mixtral 8x22b as reward function) vs Online KTO (Only llama3.1-70b as reward function)</p>
<p dir="ltr">As expected, the results show the need for a database executor function. We hypothesize that this is because the LLM has no prior knowledge about the database schema or what output it will give when running on an actual database; hence, it can&rsquo;t provide a reliable reward.&nbsp;</p>
<h2 dir="ltr">Final Result</h2>
<p dir="ltr">Finally, we wanted to see what we achieved. Until now, we have used an SFT model trained on 75% of the train data and the other 25% for the RLHF. But what have we achieved compared to a full SFT model trained on complete data?</p>
<p dir="ltr"><img src="https://img.plasmic.app/img-optimizer/v1/img?src=371f09c7da5da33b6ade3e404bc98088.png&amp;f=webp&amp;q=75" alt="" width="1512" height="870">Figure 7: Test scores comparing SFT (full train dataset) vs SFT (75% train data) vs OKTO (SFT 75% + RLHF 25%)</p>
<p dir="ltr">As we can see from the results, even if we train the data using SFT on complete train data, our RLHF model (75% SFT + 25% RLHF) beats the full SFT model in all cases other than the Challenging domain. The decrease in performance in the challenging domain is due to the major gap in performance based on SFT on 75% vs. 100% of the data (potentially due to the random data split). The 8%+ gap was too large for the OKTO step to overcome.&nbsp;</p>
<h2 dir="ltr">What&rsquo;s Next</h2>
<p dir="ltr">From our experiments, we found that dataset size or split is a significant factor in Text2SQL RLHF. Future work will focus on developing synthetic or human-annotated Text2SQL datasets more aligned with real-life databases. Another direction to explore is creating preference datasets for SQL pairs, which will help to make a good reward model. Another plausible direction for training the RLHF Text2SQL model is to explore the mixture of different open-source datasets and how different split strategies affect the final performance.</p>
<h2 dir="ltr">Conclusion</h2>
<p dir="ltr">Our research demonstrates that RLHF creates an effective production flywheel for Text2SQL systems. By learning from incorrect SQL queries in real-world use, models can continuously improve without requiring extensive manual annotation of correct query pairs. This transforms production errors from mere logs into valuable training signals, enabling Text2SQL systems to automatically adapt to enterprise-specific SQL patterns over time.</p>
<p dir="ltr">Key findings show that reserving 25% of data for RLHF training and 4x augmentation through schema shuffling produces optimal results. The hybrid reward model combining DB execution and LLM evaluation proved superior to using either component alone.&nbsp;</p>
<p dir="ltr">Online KTO emerged as the most effective RLHF algorithm for handling imbalanced datasets and SQL's abundant special tokens. Notably, the RLHF model trained on 75% SFT + 25% RLHF outperformed full SFT models in most domains, though it struggled with challenging queries due to the significant initial performance gap.&nbsp;</p>
<p dir="ltr">This work provides valuable insights for improving Text2SQL systems using RLHF, particularly in scenarios with limited high-quality training data and complex target SQL patterns. If you&rsquo;re interested in using our Text2SQL model for your enterprise use case or want to learn more about our&nbsp;<a href="https://scale.com/enterprise/generative-ai-solutions">GenAI solutions</a> offering for Enterprise, please <a href="https://scale.com/demo">book a demo here</a>.</p>

We found that RLHF is a viable approach to improve Text2SQL accuracy

Shubhashis Roy Dipta

When RLHF Meets Text2SQL