Training a Team of Language Models as Options to Build an SQL-Based Memory
Source abstract
Abstract
Despite the rapid progress in the capabilities of large language models, they still lack a reliable and efficient method of storing and retrieving new information conveyed over the course of their interaction with users upon deployment. In this paper, we use reinforcement learning methods to train a team of smaller language models, which we frame as options, on reward-respecting subtasks, to learn to use SQL commands to store and retrieve relevant information to and from an external SQL database. In particular, we train a storage language model on a subtask for distinguishing between user and assistant in the dialogue history, to learn to store any relevant facts that may be required to answer future user queries. We then train a retrieval language model on a subtask for querying a sufficient number of fields, to learn to retrieve information from the SQL database that could be useful in answering the current user query. We find that training our models on their respective subtasks results in much higher performance than training them to directly optimize the reward signal and that the resulting team of language models is able to achieve performance on memory tasks comparable to existing methods that rely on language models orders of magnitude larger in size. In particular, we were able to able to achieve a 36% gain in accuracy over a prompt engineering baseline and a 13% gain over a strong baseline that uses the much larger GPT-3.5 Turbo on the MSC-Self-Instruct dataset.
Source: crossref · Quality: baseline
Research context
OpenAlex topics: Topic Modeling · Natural Language Processing Techniques · Multimodal Machine Learning Applications
SciVX discovery tags are terminology-based. OpenAlex topics are source-provided classifications.
Affiliations
College of Informatics, Korea University, Seoul 02841, Republic of Korea
School of Electrical Engineering, Korea University, Seoul 02841, Republic of Korea
Department of Electrical and Computer Engineering, Catholic University of America, Washington, DC 20064, USA
References (44)
Show source-provided references
1. Radford, A., and Narasimhan, K. (2025, August 18). Improving Language Understanding by Generative Pre-Training. Available online: https://www.mikecaptain.com/resources/pdf/GPT-1.pdf.
2. Devlin, J., Chang, M.W., Lee, K., and Toutanova, K. (2019, January 2–7). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics, Minneapolis, MN, USA.
3. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. (2025, August 20). Language Models Are Unsupervised Multitask Learners. Available online: https://storage.prod.researchhub.com/uploads/papers/2020/06/01/language-models.pdf.
4. Brown, T.B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., and Askell, A. (2020). Language Models are Few-Shot Learners. arXiv.
5. Kaplan, J., McCandlish, S., Henighan, T.J., Brown, T.B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. (2020). Scaling Laws for Neural Language Models. arXiv.
6. Achiam, O.J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F.L., Almeida, D., Altenschmidt, J., Altman, S., and Anadkat, S. (2023). GPT-4 Technical Report. arXiv.
7. Lin, S.C., Hilton, J., and Evans, O. (2021, January 1–6). TruthfulQA: Measuring How Models Mimic Human Falsehoods. Proceedings of the Annual Meeting of the Association for Computational Linguistics, Online.10.18653/v1/2022.acl-long.229 ↗
8. Li, J., Cheng, X., Zhao, W.X., Nie, J., and rong Wen, J. (2023). HaluEval: A Large-Scale Hallucination Evaluation Benchmark for Large Language Models. arXiv.10.18653/v1/2023.emnlp-main.397 ↗
9. Min, S., Krishna, K., Lyu, X., Lewis, M., tau Yih, W., Koh, P.W., Iyyer, M., Zettlemoyer, L., and Hajishirzi, H. (2023). FActScore: Fine-grained Atomic Evaluation of Factual Precision in Long Form Text Generation. arXiv.10.18653/v1/2023.emnlp-main.741 ↗
10. Wei, J., Yang, C., Song, X., Lu, Y., Hu, N., Huang, J., Tran, D., Peng, D., Liu, R., and Huang, D. (2024). Long-form factuality in large language models. arXiv.
11. Lee, S., Kim, M., Cherif, L., Dobre, D., Lee, J., Hwang, S.J., Kawaguchi, K., Gidel, G., Bengio, Y., and Malkin, N. (2024). Learning diverse attacks on large language models for robust red-teaming and safety tuning. arXiv.
12. Hong, Z.W., Shenfeld, I., Wang, T.H., Chuang, Y.S., Pareja, A., Glass, J., Srivastava, A., and Agrawal, P. (2024). Curiosity-driven Red-teaming for Large Language Models. arXiv.
13. Inan, H., Upasani, K., Chi, J., Rungta, R., Iyer, K., Mao, Y., Tontchev, M., Hu, Q., Fuller, B., and Testuggine, D. (2023). Llama Guard: LLM-based Input-Output Safeguard for Human-AI Conversations. arXiv.
14. Adaptive Cooperative Fault-Tolerant Control for Output-Constrained Nonlinear Multi-Agent Systems Under Stochastic FDI Attacks10.1109/TCSI.2025.3534194 ↗
15. Adaptive Tracking Control for Uncertain Nonlinear Multi-Agent Systems With Partially Sensor Attack10.1109/TASE.2024.3440394 ↗
16. Lost in the Middle: How Language Models Use Long Contexts10.1162/tacl_a_00638 ↗
17. Du, Y., Tian, M., Ronanki, S., Rongali, S., Bodapati, S., Galstyan, A., Wells, A., Schwartz, R., Huerta, E.A., and Peng, H. (2025). Context Length Alone Hurts LLM Performance Despite Perfect Retrieval. arXiv.
18. Li, T., Zhang, G., Do, Q.D., Yue, X., and Chen, W. (2024). Long-context LLMs Struggle with Long In-context Learning. Trans. Mach. Learn. Res., 2025.
19. Tang, Z., Zhou, K., Li, J., Ji, B., Hou, J., and Zhang, M. (2024). L-CiteEval: Do Long-Context Models Truly Leverage Context for Responding?. arXiv.
20. Catastrophic Interference in Connectionist Networks: The Sequential Learning Problem10.1016/S0079-7421(08)60536-8 ↗
21. Catastrophic forgetting in connectionist networks10.1016/S1364-6613(99)01294-2 ↗
22. Robins, A.V. (1993, January 24–26). Catastrophic forgetting in neural networks: The role of rehearsal mechanisms. Proceedings of the First New Zealand International Two-Stream Conference on Artificial Neural Networks and Expert Systems, Dunedin, New Zealand.
23. Overcoming catastrophic forgetting in neural networks10.1073/pnas.1611835114 ↗
24. Zhong, W., Guo, L., Gao, Q.F., Ye, H., and Wang, Y. (2023). MemoryBank: Enhancing Large Language Models with Long-Term Memory. arXiv.10.1609/aaai.v38i17.29946 ↗
25. Learning without Forgetting10.1109/TPAMI.2017.2773081 ↗
26. Ramasesh, V.V., Lewkowycz, A., and Dyer, E. (2022, January 25–29). Effect of scale on catastrophic forgetting in neural networks. Proceedings of the International Conference on Learning Representations, Virtual.
27. Between MDPs and Semi-MDPs: A Framework for Temporal Abstraction in Reinforcement Learning10.1016/S0004-3702(99)00052-1 ↗
28. Reward-Respecting Subtasks for Model-Based Reinforcement Learning10.1016/j.artint.2023.104001 ↗
29. Mitchell, E., Lin, C., Bosselut, A., Finn, C., and Manning, C.D. (2021). Fast Model Editing at Scale. arXiv.
30. Meng, K., Bau, D., Andonian, A., and Belinkov, Y. (December, January 28). Locating and Editing Factual Associations in GPT. Proceedings of the Conference on Neural Information Processing Systems, New Orleans, LA, USA.
31. Meng, K., Sharma, A.S., Andonian, A., Belinkov, Y., and Bau, D. (2022). Mass-Editing Memory in a Transformer. arXiv.
32. Hartvigsen, T., Sankaranarayanan, S., Palangi, H., Kim, Y., and Ghassemi, M. (2022). Aging with GRACE: Lifelong Model Editing with Discrete Key-Value Adaptors. arXiv.
33. Gupta, A., Rao, A., and Anumanchipalli, G.K. (2024, January 11–16). Model Editing at Scale leads to Gradual and Catastrophic Forgetting. Proceedings of the Annual Meeting of the Association for Computational Linguistics, Bangkok, Thailand.10.18653/v1/2024.findings-acl.902 ↗
34. Lewis, P., Perez, E., Piktus, A., Petroni, F., Karpukhin, V., Goyal, N., Kuttler, H., Lewis, M., tau Yih, W., and Rocktäschel, T. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. arXiv.
35. Wang, W., Dong, L., Cheng, H., Liu, X., Yan, X., Gao, J., and Wei, F. (2023). Augmenting Language Models with Long-Term Memory. arXiv.
36. Packer, C., Fang, V., Patil, S.G., Lin, K., Wooders, S., and Gonzalez, J. (2023). MemGPT: Towards LLMs as Operating Systems. arXiv.
37. Xu, J., Szlam, A., and Weston, J. (2021). Beyond Goldfish Memory: Long-Term Open-Domain Conversation. arXiv.10.18653/v1/2022.acl-long.356 ↗
38. Simple statistical gradient-following algorithms for connectionist reinforcement learning10.1023/A:1022672621406 ↗
39. Guo, T., Chen, X., Wang, Y., Chang, R., Pei, S., Chawla, N., Wiest, O., and Zhang, X. (2024, January 3–9). Large Language Model based Multi-Agents: A Survey of Progress and Challenges. Proceedings of the International Joint Conference on Artificial Intelligence, Jeju Island, Republic of Korea.
40. Wang, Z., Moriyama, S., Wang, W.Y., Gangopadhyay, B., and Takamatsu, S. (2025). Talk Structurally, Act Hierarchically: A Collaborative Framework for LLM Multi-Agent Systems. arXiv.
41. Yuan, S., Song, K., Chen, J., Tan, X., Li, D., and Yang, D. (2024, January 16–21). EvoAgent: Towards Automatic Multi-Agent Generation via Evolutionary Algorithms. Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics, Mexico City, Mexico.
42. Yang, Y., Chai, H., Shao, S., Song, Y., Qi, S., Rui, R., and Zhang, W. (2025). AgentNet: Decentralized Evolutionary Coordination for LLM-based Multi-Agent Systems. arXiv.10.32388/WS0VIM ↗
43. Händler, T. (2023). Balancing Autonomy and Alignment: A Multi-Dimensional Taxonomy for Autonomous LLM-powered Multi-Agent Architectures. arXiv.10.5220/0012239100003598 ↗
44. Adaptive Event-Triggered Output Feedback Control for Nonlinear Multiagent Systems Using Output Information Only10.1109/TSMC.2025.3597016 ↗