Data2Vis: Automatic Generation of Data Visualizations Using Sequence-to-Sequence Recurrent Neural Networks
By: Victor Dibia
Tags: Machine Learning, Recurrent Neural Network, Towards Data Science, Data Visualization, Generative Model

What we did

If we then think of visualizations as text specifications, the problem is simplified to generating a bunch of these specifications, given some input data. And so our next task was to explore models that learn the space of a visualization specification format and perhaps can “hallucinate” some new visualizations? We began our initial experiments following notes from Karpathy’s blog — we generated visualization specifications (data + spec), concatenated them and attempted to train an RNN that would generate something plausible.

From this first experimentations, we learned quite a bit about how well character RNNs (LSTMs) perform with respect to learning the structure of data contained within nested structures like a Vega-Lite specification. We found that while character based RNNs work well for the problem (compared to word-based RNNs), their limitations in modeling long sequences and memory requirement for long term sequences made them challenging to train (unintelligible tokens after a few characters, limited learning progress even after thousands of steps, millions of parameters without preprocessing). Also, character RNNs did not provide a concrete mapping between our input and output pairs for training. These challenges motivated us to think of other ways to formulate the problem —approaches that learn the structure of our output space, and generates new data given some input. Essentially translating from source data specification to target visualization specification.

Sequence to Sequence Models

Sequence to sequence models [1,2,3] are capable of doing just that — they take in some input (e.g. text in one language) and generate an output (e.g. same text in a different language). They have been applied with wide success to problems such as language translation, text summarization, image captioning. Variations of sequence to sequence models that utilize encoder decoder architectures have also been explored for translating between formal programming languages, learning domain specific programs and general program synthesis.

We train a sequence to sequence model based on the work of Britz et al 2017. Training data consists of a source token (a single row from a data set) and a target token (a valid Vegalite visualization specification) pair.

Data and Training

First we assembled a training dataset (source and target pairs) based on examples of valid Vega-Lite visualizations. A total of 215k pair samples based on 11 distinct datasets and 4300 Vega-Lite examples were used. Please see the paper for more details on how the data is sampled and pre-processed.

We train a sequence to sequence model using the architecture and sample code provided by Britz et al 2017. The model is an encoder-decoder architecture with attention mechanism (please see the Britz paper for more details on the architecture). We also perform some simple normalization on the dataset to simplify training. We replace string and numeric field names using a short notation — “str” and “num” in the source sequence (dataset). Next, a similar backward transformation (post processing) is replicated in the target sequence to maintain consistency in field names. These transformations help scaffold the learning process by reducing the vocabulary size, and prevents the LSTM from learning field names (learning field names in the training set is not useful to our generation problem). In turn we are able to reduce the overall source and target sequence length, reduce training time and reduce the number of hidden layers which the model needs to converge.

Beam Search Decoding

To explore a variety of generated visualizations, we use beam search decoding algorithm. As opposed to outputting the most likely (highest probability) translation of an input sequence, beam search expands all possible next steps during generation and keeps the k most likely, where k is a user specified parameter known as the beam width. Unlike conventional language translation systems where beam search is applied mainly to improve translation quality by maximizing conditional probabilities of generated sequences, we also explore beam search as a way to generate a diverse set of candidate visualizations by outputting all parallel beam results. With beam search, we observe the model generates diverse plots, exploring combinations of chart types and the use of multiple variables.

Early Results