Tree Variational Autoencoder for Code

Autoencoder models of source code are an emerging alternative to autoregressive large language models with important benefits for genetic improvement of software. We hypothesize that encoder-decoder architectures are suboptimal for source code because they ignore the grammatical structure that can be derived with an Abstract Syntax Tree parser. We propose a structured Variational Auto-Encoder based on TreeLSTM that operates directly on the AST. We train it along with a baseline sequence VAE on a dataset of competitive programming submissions We find the structured model to perform better in most tests, with some notable exceptions. These findings suggest structured autoencoder models could enable more effective generation and manipulation of source code for tasks like automated bug fixing and generative programming.