Artificial Intelligence and Data Analytics Research Group

Project Leviathan 🐙

[DESCRIPTION WIP]

Project Leviathan is president Jason Banuelos's personal research project; the goal of this project is to train a language model that has been pre-trained for program synthesis to use a virtual terminal to solve programming problems.

Language models trained for program synthesis such as Codex, CodeGen, and AlphaCode have become remarkably adept at a generating short, useful blocks of code. However, they are not nearly as proficient at generating long, complete programs, and are certainly not capable of creating multi-file programs due to the inherent nature of their functionality. Unlike human programmers, since they are autoregressive, they are only capable of adding to what they have already written, and are thus incapable of editing and testing their code. But since the input and output of language models entirely consist of text, it is possible that a language model could control a terminal by retrieving the state of the terminal as Unicode text and outputting text to control it. If a model were hooked up to a terminal, it could potentially write, edit and test code and write long, complicated programs for its user.

Unfortunately, we have not been able to find a library or environment that has the functionality needed to give an agent complete control over a terminal like a human user. Because of this, the first goal of this project is to develop an environment that can get the state of a virtual terminal as a string of text and control the terminal only by inputting text into it. The virtual terminal should be able to do anything that a typical terminal used by humans can do. 

To test a model's ability to solve programming problems, we would like to use a datasets such as APPS consisting of programming problems with respective test cases and code solutions to assign the model programming problems to solve and check if its code successfully solves the problem. The model's programming problem will be given to it as a part of each environment observation, and the model's performance will be measured by the percentage of programming problems it is able to solve.

The second goal of this project is to develop increasingly powerful models for solving problems in this environment. We do not anticipate that a standard pre-trained model, regardless of its size, will perform very well in this environment. This is because models have a limited context window, meaning they will only be able to remember so much of the code they're working on at a time. As such, we anticipate they will only perform well on simple programming problems in which their code solution can fit within their context windows. We will attempt to train one of these models via reinforcement learning by rewarding the agent when its code successfully solves its given problem and giving it no reward otherwise. However, we anticipate that these models, on their own, will only have limited success at performing these tasks due to their inherent limitations.

To allow these models to create longer, more complex programs, we plan to augment our model with retrieval capabilities similar to those of the Retrieval-Enhanced Transformer, and save previous environment states to the model's retrieval database so relevant information from the previous states can be recalled when needed. After implementing the environment and testing a default model, this is what we are currently planning on pursuing next.

To further improve the model's performance, we intend to develop a planning algorithm similar to Stochastic MuZero that uses transformer-based language models instead of convolutional neural networks. Rather than selecting from an existing set of potential actions as Stochastic MuZero does, our model will generate its own actions by generating a set of potential strings to send to the terminal environment, and will evaluate these strings the same way that Stochastic MuZero evaluates potential actions. In my experience, language models typically output satisfactory completions to prompts after a few tries, so essentially our planning algorithm will allow the agent to come up with a few potential actions and select the best among them.

With the general goal of teaching an AI to use a terminal to solve problems, we plan to continue exploring other developments in natural language processing and general-purpose agents so we can push the capabilities of our model as far as we can.

The project has been split into several smaller sub-projects to make its goals more manageable and to allow us to publish multiple papers about the project as it progresses. 

The Environment

The first thing we must implement, of course, is the environment in which the agent will act and be given its problems.

• Virtual Terminal Module

The implementation of the virtual terminal is what Jason is currently focusing on. He is currently working on implementing a Python module for this virtual terminal for his Capstone project.

• Programming Problem Assignment Manager

automatically assign programming problems and check generated solutions

Training a Base Model via Reinforcement Learning

sounds simple, but I anticipate it will be harder than it looks

The New Model

To create the most proficient agent we can, we have several research areas we plan to explore.

• RETROfitting CodeGen

We anticipate that giving the model a form of long-term memory will be the most important step towards making it proficient at solving problems in this environment. As such, 

• Transformer-based Stochastic MuZero

Giving the agent planning capabilities. might need to check out PALMER first before going all-in on this? i think there was another thing that introduced a thing that could plan across arbitrary time steps but idk I'll have to look back over all the tabs on my phone

The "minimum viable paper" we can publish for this project is to hook up an existing pre-trained language model to the environment and see how well it solves programming problems without any fine-tuning. As stated earlier, we do not anticipate a standard model, even a powerful one like CodeGen, to perform very well, but this first paper will allow us to introduce the environment and provide us with a base benchmark with which to evaluate the performance of future models.

To get up to speed on the general research topics related to the project, here are some select resources

• Language Models

https://lena-voita.github.io/nlp_course/language_modeling.html

• Reinforcement Learning

https://towardsdatascience.com/reinforcement-learning-101-e24b50e1d292

• Transformers

https://huggingface.co/course/chapter1/4?fw=pt

https://youtu.be/kCc8FmEb1nY

https://huggingface.co/course/chapter1/1