Cache the attention key-value pairs and encoder outputs
\nWhen generating sequences, you may feed tokens by selecting from the calculated next-token distributions through various decoding strategies (e.g. greedy, beam-search, top-k, top-p, etc.). The important thing is that the distribution of next-token can only be changed by current token (because the previous tokens are already chosen and the attention vectors of them are always same at every timesteps). Let's see the below figure.
\n![\"attention](\"https://i.ibb.co/p3dkjvv/a.png\")
\nLet N be a length of current predicted sequence, then as you can see, the computational cost of attention layer is O(N^2). Since the previous tokens have same key-value vectors at every timesteps, we can determine that the above attention is really inefficient.
\n![\"attention](\"https://i.ibb.co/ZKZsstk/b.png\")
\nHowever, when we cache the previously-calculated key-value pairs and reuse them to only the current query vector, we can reduce the complexity to O(N). Many transformer-related libraries support caching attention. huggingface transformers, representatively, use past to reuse the attention vectors in sentence generation.
\nOf course, it can be applied to the encoder-decoder attention as well. Once you get encoder outputs, do not recalculate and just caching them. It also can improve the inference time.
Use torchscript, if you are using PyTorch
\ntorch.jit.script provides significant performance improvement. Using this function, your model will be compiled to torchscript. Typically, the torchscript-compiled models are much faster than the models which are performed through Python. That's because one of the features of torchscript is to reduce the bottleneck (or latency) between the operations. All you have to do is just to wrap your model with torch.jit.script.