A Quantum-Enhanced LSTM Layer. Using the PennyLane Quantum Machine... | by Riccardo Di Sipio | Towards Data Science

In recent years, Toronto-based startup company Xanadu has introduced a python framework called PennyLane which allows users to create hybrid Quantum Machine Learning (QML) models. While it’s still too early to claim that quantum computing has taken over, there are some areas where it can give an advantage as for example in drug discovery or finance. One field that so far has been poorly explored in QML is Natural Language Processing (NLP), the sub-field of Artificial Intelligence that gives computers the ability to read, write and to some extent comprehend written text.
As documents are usually presented as sequences of words, historically one of the most successful techniques to manipulate this kind of data has been the Recurrent Neural Network architecture, and in particular a variant called Long Short-Term Memory (LSTM). The “trick” that makes these networks so popular in the analysis of sequential data is a combination of “memory” and “statefulness” that help with identifying which components of the input are relevant to compute the output. While the mathematics is quite thick as we’ll see later on, suffices to say for now that LSTMs allowed machines to perform translations, classification and intent detection with state-of-the-art accuracy until the advent of Transformer networks. Still, it’s interesting at least from an educational point of view to dig into LSTMs to see what good quantum computing may bring to the field. For a more thorough discussion, please refer to “Quantum Long Short-Term Memory” by Chen, Yoo and Fang (arXiv:2009.01783) and “Recurrent Quantum Neural Networks” by J. Bausch (arXiv:2006.14619). On the same subject, see also my other posts Toward a Quantum Transformer, Classifying Scientific Documents with Quantum-Enhanced Transfer Learning and Spooky Computers at CERN.
To begin with, let’s review the inner workings of a LSTM. We assume that the input x is composed of a sequence of t time-steps (e.g. words), each represented by a N-dimensional feature vector (in practical applications N can be large, e.g. 512). Also, the network stores a hidden array of vectors h and a state vector c that are updated for each element of the input sequence. For example, if we are only interested in a summary of the sequence (e.g. in a sentiment analysis), the last element of the array h will be returned. Instead, if we are interested in a representation of each element (e.g. to assign to each word a part-of-speech tag such as noun, verb, etc...), we want to have access to every element of h.
The calculation can be summarized in the formulas below:
where v_t is a concatenation of the input element at step t and the hidden state at step t-1, i.e. v_t = [h_(t-1), x_t]. In the machine learning parlance, borrowed from electric circuit analysis, f_t is called forget gate, i_t is the input gate, C_t is the update gate and o_t is the output gate. The matrices W_f, W_i, W_C and W_o and the bias vectors b_f, b_i, b_C and b_o are the parameters that have be learned during the supervised training and implement the part of the calculation called linear dense layer that we want to replace with the quantum equivalent. As often the case, a non-linearity is introduced by the application of Sigmoid and Hyperbolic Tangent (tanh) functions to the output of these four dense layers, which effect is to determine whether a part of the input has to be considered (values close to 1) or ignored (values close to 0).
A key concept in Quantum Machine Learning is that information is stored in units of information called qubits that are a so-called superposition of the states 0 and 1. Physically, they may represent for example the orientation of the spin of a particle with respect to the axis of a magnetic field, or the circulation of electric current in a superconductor. For practical purposes, these details can safely be ignored. What’s important is to remember that the calculation is carried by first putting the input qubits into the initial state (e.g. a string such as 010010), then they are entangled between each other, rotated by an arbitrary angles, and finally observed (“measured”). The goal of the training is to find the rotation angles mentioned above that optimize some cost function such as the mean squared error (MSE) between the output and some pre-determined labels. This kind of quantum circuit is known as Variational Quantum Circuit. Using the PennyLane library, one can easily define a VQC as a series of AngleEmbedding and StronglyEntanglingLayers objects.
One problem that may not be obvious at the onset is that while bits are cheap, qubits are expensive: it’s hard to keep real qubits in superposition state, it’s also hard to simulate them. So forget about designing a circuit with 12 qubits. Good news is, thanks to their own nature, fewer qubits are probably needed to perform the calculation. To cut a long story short, what we’ll do here is to introduce a “dressed quantum circuit”, i.e. we sandwich the quantum layer between two classical linear layers to match the dimensionality. For example, if the feature dimension is 8 and the hidden dimension is 6, but we can only afford 4 qubits, what we need is a sequence like this:
x_in = torch.randn(8)x = nn.Linear(8, 4)(x_in)x_q = VQC(x)x_out = Linear(4, 6)(x_q)
We need one classical layer to “squeeze” the input to match the number of qubits, and one classical layer to “bloat” the output of f_t, i_t, C_t and o_t. to match the dimension of the hidden vectors. Thus, overall we still have a number of classical parameters that have to be learned during the training. Putting all together, the core of the calculation looks like this:
Does it work? To test if all the above makes sense, an intuitive but not trivial example is needed. In this case, I decided to adapt the part-of-speech (POS) tagging example from PyTorch to the case under consideration. The idea is pretty simple: we have a set of two sentences (“The dog ate the apple” and “Everybody read that book”) whose words have been annotated with POS tags. For example, the first sentence is thus [“DET”, “NN”, “V”, “DET”, “NN”]. The idea is to pass the two sequences through the LSTM, which will output the hidden array of vectors [h_0, h_1, h_2, h_3, h_4], one for each word. A dense layer “head” is attached to the LSTM’s outputs to calculate the probability that each word may be a determinant, noun or verb (for the curious, that is a softmax layer with dimension 3).
Following the example from the PyTorch website, we train the two networks (classical and quantum LSTM) for 300 epochs. This is just a toy example and has to be taken as such, but still, the results are encouraging. The loss function decreases as a function of the training epoch, and after 300 epochs both networks are able to tag correctly the first sentence. Due to the complexity of the simulation of the quantum circuit, it took approximatively 15 minutes to finish the training, to be compared to a mere 8 seconds for the classical case. Also, looking at the plots, it seems that the quantum LSTM needed to be trained for longer, but that’s enough for our purpose.
Sentence:  ['The', 'dog', 'ate', 'the', 'apple']Labels:    ['DET', 'NN', 'V', 'DET', 'NN']Predicted: ['DET', 'NN', 'V', 'DET', 'NN']Classical loss at 300 epochs: 0.029Quantum loss at 300 epochs:   0.250
To conclude, libraries such as PennyLane or TensorFlow Quantum offer a high-level of abstraction to define quantum/digital hybrid models for Quantum Machine Learning. This is a field that is rapidly expanding and more and more applications are being explored as one can happily see by attending topical conferences such as QTML 2020.
Give me the code! https://github.com/rdisipio/qlstm