Craft your own Attention layer in 6 lines — Story of how the code evolved

The essence of Attention across all its intoxicating flavours

An abundance of information may create a poverty of attention!

Sometimes when the Dalai Lama addresses large audiences, he gets caught up in the narrative and speaks in longer and longer chunks. This complicates the job of his interpreter, who has to wait for a pause to translate the speech to English ensuring that every essence of the message — every intent, pun, sarcasm, irony, humor — is passed on undiluted to the audience. This is an extraordinary job requiring extraordinary skills. There are occasions where the interpreter has to wait up to 20 mins to speak and this challenges even his razor-sharp memory honed by years of meditation. The tension in the air is palatable as the audience looks at the interpreter in sympathy, wishing him to somehow succeed. And succeed he does… much to the appreciation of the audience. If someone were to hack into the interpreter’s brain to find out how he does this, they might conclude that the interpreter would perhaps have made a mental note of all the important words to which to pay ‘attention’ during translation.

By now, the reader would have made the connection with the LSTMs of the machine learning world and how the problem of very long sequences of text could cause the ML model to falter, were it not for an innovation called ‘Attention’ that gained fame in 2014. Today, many state-of-the-art models use Attention in some way or another. Thankfully, there is a huge abundance of information on Attention that we can refer to. Unfortunately, it is this abundance of information that sometimes leads to a deficit of attention (no pun intended). Let me explain…

I remember sitting in the library of my first company in ‘99, going thru’ the voluminous mainframe manuals. They were fairly dry and not easy to comprehend as the material was crouched in technical jargon. It almost looked as if a focussed effort had been put to make the material difficult to understand (Jay Alammar would not have been accepted as an author in those times but S.Tharoor would probably have been given the red carpet). But I had one distinct advantage! There was only ONE source of information — take it or leave it. So I ploughed on, until solutions finally presented themselves.

Circa 2020. Me trying to contribute to a Kaggle competition which involved using AI/ML to identify the right mRNA structure which could serve as a stable Covid vaccine (link). I had in the recent past, dabbled with a few old-style ML programs (SVMs, Random forests etc) and was just beginning my first neural net program. Thanks to the basic pseudocode shared by some of the seniors, I was off to a decent start. We were convinced that the stability depended on the sequence. LSTMs were used liberally. I wanted to add an Attention mechanism on top of that because I felt LSTM was not capturing the complete sequence info. And then all my problems began!

Attention is one of those topics that is easy to understand intuitively. There are several great tutorials which bring out this ‘ intuitive’ explanation very well. Yet, when it comes to a custom practical implementation, it can be pretty confusing to understand. For starters — there are several variations and this lack of a ‘standard approach’ can sometimes be confusing. Even if we take a single standard approach, there could be several variations in the implementation. Even if we take a single standard approach and a single standard implementation, there could be minor variations in code based on the framework used and further variations within versions of the same framework itself. Moreover, some prefer using an off-the-shelf attention layer, others prefer custom layers while yet others use regular layers to achieve attention-like functionality. All of this could be confusing to the novice who is trying to peek under the hood.

I am sure you too will nod your head as I repeat the words of economist Herbert Simon who warned of an information overdose as early as 50 years back with his beautiful words —

“In an information-rich world, the wealth of information means a dearth of something else: a scarcity of whatever it is that information consumes. What information consumes is rather obvious: it consumes the attention of its recipients. Hence **a wealth of information creates a poverty of attention…**”

This article aims to introduce attention in its simplest essence possible. We then build upon this simple model to add various flavours in an easy-to-understand way. Attention can be used in a variety of use-cases and while its most common use is in sequence-sequence language translators which are reasonably well explained, they are several other use-cases as well. A de-linking of the ‘Attention’ concept from translators would help when we try to peek under the hood.

Prerequisites:

We will not spend much time on the intuition behind attention — this is covered well by Jay Allamar here — this is, in my view, one of the best introductions to ‘Attention’.

It is also assumed that you have a basic idea of RNN’s. Here is one of the best videos covering this topic — Thank you Brandon Rohrer!

This captivating 50-minute video by Jay Allamar gives you an instant intuitive overview of embeddings. Jay is indeed a master story-teller!

Last but not the least, if you are totally new to neural nets, here is a link to 3Blue1Brown’s excellent video. This introduces neural nets in the softest way possible. This 30 min lecture is equivalent to many weeks of classes. I include this link in nearly every article that I write, as a tribute to its creator Grant Sanderson. But there are other reasons why I admire Grant. Anybody who can get 3.2 Million viewers to subscribe for a set of ‘Maths’ courses must be an illusionist of the highest order.

Using ‘Attention’ in a classification problem:

Let us look at using ‘Attention’ in a classification problem instead of the regular translator example that we typically see. You want to take a tweet and tell whether the sentiment is positive or negative. This is a sequence classification exercise. Let us look at this hypothetical tweet: ‘Attention is so hard to understand. There are so many blogs and implementation and tend to confuse the novice’

Overall, the sentiment is somewhat negative. A typical classification exercise (without Attention) might involve feeding the above tweet to an LSTM which retains a memory of the sequence and then uses the final hidden state (the output after the last word) as an input to a feedforward dense neural net layer which learns and then classifies. It is assumed you already have a fair understanding of the below steps:

1. Input = set of tweets. To keep things simple, let us assume that each tweet contains exactly 19 words.
2. One layer of embedding to convert each word in each tweet into a 100 dimensional vector. So now you have 19 100-dimensional vectors for each tweet. 19 is going to be the sequence length of the RNN also called the number of time-steps.
3. Pass this through a bi-directional LSTM of 128 units. It generates one hidden state for each word. Since it is a bi-directional LSTM, so the number of dimensions of each hidden state (one for each word) becomes 2 128 = 256. By default, only the last word’s hidden state is returned as the output for each tweet. So the output is a single vector of 256 dimensions for each tweet — a sort of a ‘sentence embedding’ if you want to call it that… and you hope it captures the essence of the whole tweet well*
4. This is passed as input to a simple feedforward dense layer. The output is the sentiment of the tweet. One binary value for each tweet

The results seem good but now you want to bring in an ‘Attention’ mechanism to further improve the score. You feel that instead of using just the last hidden state of the LSTM in step 3 as an input to the feedforward layer in step 4, you need to pay attention to every single hidden state generated, decide which ones are important and then take a weighted sum and feed this as input to step 4. You decide to plunge headlong into the code. The LSTM layer at Step 3 has the option of returning a hidden state for every word with the ‘return_sequences=true’ option and you decide to make full use of it. So far so good.

Now, after step 3, you have 19 hidden states. Intuitively, you know that all that is needed now is a set of weights. You need 19 weights overall — one for each hidden state, i.e basically, one weight for each of the 19 words. You intuitively know that words like ‘hard’, ‘confuse’ etc should be given higher weight and if it were in your hands, you would probably give 49.5% weight to each of these two words and distribute the remaining 1% to each of the other 17 words. Unfortunately, that kind of hard-coding won't work and you need the machine to output these weights. Your intuition tells you that when you don't know something, you can use neural nets to guess those values. So you create an ‘ADDITIONAL’ simple feedforward layer between steps 3 & 4 to give you those 19 weights!

Before we get into the code, let us first try to understand a little more about this so-called ‘simple additional feedforward layer’. Maybe, because it is so simple, it is not typically explained in existing blogs.

• Firstly, let us see why we need an ‘Additional’ layer? It is quite possible to implement attention ‘inside’ the LSTM layer at step 3 or ‘inside’ the existing feedforward layer in step 4. However, it makes sense to bring in a clean new layer to segregate the attention code to understand it better.
• *This new layer can be a dense single layer *Multilayer Perceptron *(MLP) with a single unit*. There are many oxymorons here. Let us try to understand this in more detail. We often associate neural nets with hundreds of neurons and dozens of layers. So why would we want to use a “neural net” with only one neuron and a single layer? Well, for one you can replace this layer with any contraption of your choice as long as it returns the output needed (a set of attention weights — one per each word in the sequence) without overfitting. The fact is that this is a linear problem and can be solved with a simple structure. Adding complexity may not help in a big way. If it helps, think of a neural net layer as a function approximator. We feed it some inputs, outputs and expect it to train itself to find the definition of the function such that: o/p = function (i/p). Typically the function involves millions of variables involving ‘n’ degrees, but often the function could be as simple as: output y = 12x1 + 3x2 +5. In this case, 1 neuron and 1 layer more than suffice to calculate the function definition.
• Inputs to this ‘Dense single MLP — single unit layer’ will be all the hidden states of the LSTM sequence from the previous step. So we have 19 inputs for every tweet to this layer. Each input continues to have 256 dimensions. So the input shape should be: (?,19,256). The ‘?’ tends to confuse, but it is just a placeholder for the number of tweets that the model is going to process in each batch. Let us ignore that and look at the last 2 numbers in the shape which is what is relevant. This is (19, 256). So the input is basically a set of 19 vectors for each tweet, each vector of 256 dimensions.
• The goal of this layer is that it should give us the ‘attention weights’ — one for each word. Number of attention weights = sequence length of the RNN or the number of time-steps in the RNN or the number of words in the tweet which is just 19. So the shape of the output should be: (?,19,1). 19 1-dimensional probabilistic attention weights — a0, a1,….a18. These should add up to 1.
• This layer also needs to return the weighted sum. In fact, this is the actual output that goes to the next layer, not the weights. Let us call this output the ‘attention adjusted output state’. The shape of this is also (?, 1, 256). Basically, you use the attention weights discovered above to magnify or diminish each word (i.e. multiply each of the 256 dimensions of word 1 by a0, word 2 by a1 and so on for all 19 words in the tweet) and add them up to return one final ‘word’ which continues to be of 256 dimensions. This is the ‘attention adjusted output state’ and is to be fed to the final layer — at step 4. The attention weights themselves need not be fed to the next layer as their job is already done. However, if the user likes, they can be fed to a nice graph to display which words contribute to the sentiment in a big way
• Lastly, the purpose of every neural layer is to calculate the weights of the connections during the training process and arrive at a set of optimal weights and biases for that layer. So for a neural layer to work, we need to define the weights and bias needed for that layer. What should be the shape of the weight and bias for this ‘’Dense single MLP — single unit layer’ we are introducing? This is the last question we need to address before jumping into the code.

Intuitively, you may want to define the layer weights as (?,19,1). If you train this model to generate 19 weights, you feel this will give you a weight vector where the words number 4 (hard) and number 17 (confusing) have higher weights. Can you guess what could be the problem with this?

The problem is — You are training the layer to assign weights based on the position of the words instead of training the model to assign weights based on the meaning of the words. So if after the training (i.e., during production launch), another tweet were to come up for classification which read something like: ‘I understood attention after reading Jay Allamar's brilliantly written blog. It is too good and all misunderstanding was cleared’

Guess what happens now? The weights of the model you have trained are such that the 4th and the 17th word get importance. So ‘after’ and ‘misunderstanding’ gets all the attention and most likely the sentiment would be rated negative and Jay might feel sad.

So instead of training the model on the ‘position’ of the words, we need to train the model on the ‘meaning’ of the words. Since the ‘meaning’ of the word is encoded in its ‘embedding’, this basically means that the model needs to consider all 256 dimensions of each word during training and create weights for each one of those. So this means that the number of weights generated by the layer after training should be 256. So the weights in the additional layer we are defining should be (?, 256, 1) — one weighted connection each for the 256 inputs connected to the solitary neuron in this layer. This also makes sense mathematically speaking, during training, we matrix multiply each word embedding by this weight. So (?, 19, 256) matrix is multiplied by (?, 256, 1) and this yields us a set of weights (?, 19,1) which are our attention weights (well almost…we have to use an activation and add a softmax layer to ensure that all these weights add up to 1). The rest is a simple multiplication to get the weighted sum. This is the ‘attention adjusted output’ which is fed to the next layer.

With this, we have defined the basics of what our ‘Attention’ layer looks like. If you have understood the above — particularly the difference between the ‘attention weights’ which is of shape (19,1) and ‘layer weights’ which is of shape (256,1), then implementation is going to be a piece of cake irrespective of the framework you are using. There will also be 1 bias for every word, so the bias shape is going to be (19,1). We will do a custom implementation which takes less than half a dozen lines of code. Of course, this is the bare essence after peeling away all the layers of the onion skins. Once you have understood the bare essence, it is always better to add back the peels which help in keeping the code more structured/readable and maintainable.

By the way, which type of attention is this — you may ask — xyz or abc or def? To be frank, at this point in time, the answer is — I don't know. In fact, this reminds me of a funny story read in the Tinkle comic book more than 3 decades back — The owner of a house has a lot of fancy gadgets and likes to show off a bit. When a guest comes and requests to look at his latest acquisition, his assistant gets that object from the other room. After the guest leaves, the owner reprimands the assistant saying, he should always ask whether he needs to bring the American model or the Asian one or the European one etc… the idea being to show off his richness. The assistant, of course, gets miffed and when another guest comes the next day and the owner requests his assistant to call his father from the other room to make introductions, the assistant asks — Shall I get your American father or European father or your Asian father?

Well, humour never got me anywhere and may well get this article rejected from publications, nonetheless, I am going to keep this little joke. The fact is that it does not matter whether we are dealing with ‘this’ Attention or ‘that’ Attention or some ‘other’ Attention. We will take the intuitive approach for the problem at hand, write the code and then later reflect upon which attention model our code closely resembles. Most versions are minor tweaks of each other… I am sure it would resemble one of the many models present today and if it doesn't — hey, we got a new IEEE paper on our hands!

Writing your own Attention layer:

Adding a custom layer to any neural net is easy. In general, it follows the same pattern irrespective of whether you are using TF, Keras, Pytorch or any other framework. Let us use Keras. We start by sub-classing the base ‘layer’ class & create our own custom layer. Next, define the weights in the build() method. Put the logic in the call() method and we are pretty much done. Oh, we also need to define the init() layer. There are a couple of optional methods we can choose to define e.g. to save and load the layer later but we will stick to basic usage for now.

from keras import backend as K

class peel_the_layer(tf.keras.layers.Layer): 

    def __init__(self):    
        # Nothing special to be done here
        super(peel_the_layer, self).__init__()

    def build(self, input_shape):
        # Define the shape of the weights and bias in this layer
        # As we discussed the layer has just 1 lonely neuron
        # We discussed the shapes of the weights and bias earlier
        self.w=self.add_weight(shape=(256,1), initializer="normal")
        self.b=self.add_weight(shape=(19,1), initializer="zeros")
        super(peel_the_layer, self).build(input_shape)

    def call(self, x):
        # x is the input tensor of 256 dimensions

        # Below is the main processing done during training
        # K is the Keras Backend import
        e = K.tanh(K.dot(x,self.w)+self.b)
        a = K.softmax(e, axis=1)
        output = x*a

        # return the outputs. 'a' is the set of 19 attention weights
        # the second variable is the 'attention adjusted o/p state'
        return a, K.sum(output, axis=1)

So this is it! Exactly half a dozen lines of code and we have created an attention layer. We invoke this layer between steps 3 & 4 — let us call it step 3.5

# lstm_out is o/p of step 3 and is an i/p to step 3.5

a, attn_adjusted_op = peel_the_layer()(lstm_out)

# attn_adjusted_op is o/p of step 3.5 and is an i/p to step 4

We need to add some peels back though. For one thing — all the hardcoding business is not good. We have the input shape which comes to the layer class build() method as a param ‘input_shape’. Let us derive the layer weights and bias shapes from that. We discussed the input shape would be: (?,19,256). The layer-weight is currently hardcoded as (256, 1). The last index of the input_shape is nothing but the number of dimensions for each word. This should be the number of weights in the layer. The last-but-one index of the input_shape is the number of timesteps or words in the tweet. Note that I use reverse indexing instead of forward which eliminates hassles with presence or absence of batch size. Lastly, we have one solitary neuron in this layer. This explains the hard-coding of 1. Let us refine the build method slightly as:

def build(self, input_shape):

    # For convenience sake, add a couple of properties
    self.num_dim_perword = input_shape[-1] ##Last value of inp:256
    self.words_pertweet = input_shape[-2]  ##Last but 1 value:19
    num_units = 1

    self.W = self.add_weight((num_dim_perword,num_units),
                                 initializer='normal')

    self.b = self.add_weight((words_pertweet,num_units),
                                 initializer='zero')

Better, but we can do further improvements. There are different types of initializers, so you may want to soft-code that. Also, it makes more sense to define the num_units in the init() method as a method parameter instead of hardcoding it to 1 in build(). You can then reference it in the build() method by using self.units. Lastly, avoid specific variable names like num_dim_perword or words_pertweet. Rename it to inp_dimensions and seq_length etc which is more generic.

Let us turn our attention to the call() method now. Actually, it has only 3 lines and is as straightforward as can be:

# ‘w’ is the weight of the layer and ‘a’ is the attention weights

e = K.tanh(K.dot(x,self.w)+self.b)
a = K.softmax(e, axis=1)
output = x*a

return a, K.sum(output, axis=1)

We multiply the inputs ‘x’ of shape (19 256) by the layer weights ‘w’ of shape (256 1) and obtain a (19 1) matrix. We add the bias (19 1) and pass the output thru’ any activation layer. So we now have 19 1 values (I would not call them attention weights yet). We take a softmax of these values. Softmax squashes these into values in the range between 0, and 1 whose sum is 1. These are the 19 attention weights. We multiply each attention weight by the respective word and sum up and we are done. We now have the ‘attention adjusted output*’ state ready to be fed to the next layer.

It is a good idea to also look at the minor changes needed to get this code working across frameworks. For example, based on the framework you are using, you may need to explicitly ‘squeeze’ the (19 1) attention vector that you got above into a 1-D array of (19) before computing the softmax. Some frameworks do it automatically for you and some don't. After calculating the softmax, you need to ‘expand’ back the attention weights from (19) to (19 1)

 e = K.squeeze(K.tanh(K.dot(x,self.W)+self.b),axis=-1)
 a = K.softmax(e)
 a = K.expand_dims(a,axis=-1)
 output = x*a
 return a, K.sum(output, axis=1)

But the best approach I like is the functional one — something akin to the below. This helps us visualise the code as we go along and make it a little more intuitive and easy to understand.

from keras.layers import Flatten, Activation, RepeatVector, Permute, Multiply, Lambda

# No changes to init(), build() from before

def call(self, x):

    e = K.tanh(K.dot(x,self.w)+self.b)
    e = Flatten()(e)
    a = Activation('softmax')(e) 

    # Don't manipulate 'a'. It needs to be 'return'ed intact
    temp = RepeatVector(256)(a) 
    temp = Permute([2,1])(temp) 

    output = Multiply()([x,temp])
    output = Lambda(lambda values: K.sum(values, axis=1))(output)

    return a, output

The first 3 lines are easily understood. We now get ‘a’ which is a set of 19 weights — a0, a1…a19. Let us say it is 0.49999 for the 4th and the 17th word and a tiny 0.000005 or something for the rest of the 17 words.

We now need to multiply the first word (by word, I mean all of its 256 dimensions) with a0, 2nd word with a1 and so on for all 19 words. This is done by the below 3 lines.

temp = RepeatVector(256)(a)  # (?,19) becomes (?,256,19)
temp = Permute([2,1])(temp)  # change from (?,256,19) to (?,19,256)
output = Multiply()([x,temp])# Apply weight to each of the 256 dim

We take each weight-vector of 19 values and repeat it 256 times. We get 256 rows of 19 attention weights. What we actually need for the matrix multiplication is 19 rows of 256 (same valued) attention weights. So we swap the indices slightly (changing the stride). This gives us 19 vectors. Each vector has 256 dimensions. Each of the 256 dimensions has the same value which is nothing but the attention weight for the word in question. So in the 4th and the 17th row, we have a vector which has 256 values all of them [.49999, .499999, .49999,…] and for the rest of the words it is [.00005, .00005, .00005, .00005,…]. The rest of the logic is straightforward to understand. The Lambda layer wraps up any arbitrary function and gives it a ‘layer-like’ look.

All the code snippets above achieve the same thing. In fact, if you look carefully enough, we haven't done anything ‘grand’ in our new custom layer. Any layer of a feedforward network, by default, multiplies the input with the weights to arrive at the output. The only additional customization we have done is to do the softmax calculation and the weighted summation. So the next logical question is — Do we need a custom layer at all? Why don't we use the ‘regular layer’ command itself to create a layer and then do all this softmax business after we get the output from the layer? Why bother writing a custom layer? This is a very pertinent question and hence the below type of approach is also correct:

from keras import backend as K
from keras.layers import Flatten, Activation, RepeatVector, Permute, Multiply, Lambda, Dense, merge

# Define a regular layer instead of writing a custom layer
# This layer should have just one neuron - like before
# The weights and bias shapes are automatically calculated 
# by the Framework, based on the input
# This layer is defined at step 3.5 directly
e=Dense(1, activation='tanh')(lstm_out )

# Now do all the softmax business taking the above o/p
e=Flatten()(e)
a=Activation('softmax')(e)
temp=RepeatVector(256)(a)
temp=Permute([2, 1])(temp)

# multiply weight with lstm layer o/p
output = merge.Multiply()([lstm_out, temp])

# Get the attention adjusted output state
output = Lambda(lambda values: K.sum(values, axis=1))(output)
# Pass output to step 4 and 'a' to any nice display

Take good note of this approach for this is the style that will serve us best going ahead. Note that we don't need to define any custom layer here and the above code directly comes after the LSTM layer at step 3. Also, notice how we don't need to define layer weights and just need to give the number of units. The framework does all the shape-calculations by itself.

e=Dense(num_of_units, activation=’tanh’)( lstm_out )

In the above case, the framework sees that the input ‘lstm_out’ has the shape (?,19,256) and the number of units we have given to this layer is 1 and therefore the framework itself defines the layer weight shape as (?,256,1) and the output of this layer has a shape (?,19,1). We will use this style when creating more serious implementations of Attention.

There are other interesting approaches as well. For e.g. in one implementation, I found attention added in the LSTM layer — step 3 itself by overriding the necessary methods in the appropriate classes used by the LSTM layer (without harming the LSTM flow).

We discussed about half-a-dozen different possible implementations of a single flavour of Attention. While we haven't coded each one in entirety, we have examined the bare essence of each. Now you can safely navigate any ‘attention’ code and clearly know what is happening at any point in time. Now let us start tweaking this code to bring in the various flavours of attention. Before we do that, it may help to summarise intuitively what we have done so far. I find it easy to do so by putting in plain English what we expect the layer to do. In the above case, we would make a plain English request as follows — Layer, Layer in the model, Given a tweet, could you highlight all the words carrying sentiment in them?

or to be more specific — Could you give me a set of 256 weights that transform any input word such that the words involving sentiment generate a higher output signal. Oh! by the way, since this is a straightforward linear problem, I have given you only 1 neuron to help you with the calculations.

Which paper is the above simple model based on? The closest match I could find was the work done by Raffel et al. here. How about a working code? There are many implementations of this type of simple Attention for NLP in Kaggle, most inspired directly by GrandMaster qianqian solving the Toxic comment classification challenge.

Now let us turn our attention to different ‘flavours’ of adding attention. A new variable here or a minor code tweak there and hey presto — you get a new approach. For e.g. one paper defines a new variable — the ‘context’, and this ‘context’ is supposed to best summarize the sentiment of the sentence or the tweet in one word. Who determines the value of this ‘context’? Well, the authors say — let the additional feedforward layer determine the value of this context along with the weights and bias during training. So after the training, you have the weights, bias AND a ‘context’ vector that best summarizes the sentiment of the tweet. Now, we calculate the attention weights based on the similarity of each of the 19 words with the ‘context’ vector as is done here. Other than an additional line or 2 of code (one line in build method to define the context ‘u’, another in call method to do the dot-product of the output with the self.u), the rest of the processing is the same as before.

Note that the product of 2 vectors gives a measure of their similarity. Softmaxing it returns a set of 19 probabilities adding up to 1. Each probability indicates how close the word is to the context vector. The rest of the processing is the same and we finally end up with the ‘attention adjusted output’.

Intuitively, this ‘Context’ approach may help the model return a higher score as compared to the plain-vanilla ‘no-context’ attention. Instead of blindly shooting in the dark, this ‘context’ provides a reference point for the layer — it asks the layer to select weights such that after transformation, the words nearer to this ‘context’ should have a better say in determining the label.

But this Attention implementation is just a side effect of the actual intent of the authors which is to show how attention could help in document classification. Texts tend to have a hierarchical structure and the importance of words and sentences are highly context-dependent. To include this insight, Yang et al. proposed the hierarchical model here, which uses two levels of attention — one at the word level and one at the sentence level. This makes it an extremely powerful model to analyse text corpora and hence the importance of this paper cannot be emphasised enough. There is another pleasant side effect of this approach. The attention mechanism here can also function as a pure ‘reduction’ operation, which could be used in place of any pooling step. This is because the ‘context’ that is derived, is 1 word and it best summarises the sentiment of the 19-word tweet — a classic ‘reduction’.

There are several other variations of attention. For instance, Parikh et al. (2016) introduce an attention mechanism that takes two sentences and outputs a single vector. Another take on this is Attention-over-Attention — Cui et al, 2016. The idea here is to place another attention layer over the primary attentions, to indicate the relative importance of each attention. We can go on and on…

Understanding ‘Attention’ in a language translator:

In fact, this brings us back to a point where we can couple back the ‘language translator’ aspect to the above ‘attention’ code. Let us say we want to convert the tweets from English to German. This is a sequence to sequence translation. Some call it an NMT — a neural machine translator. The plain vanilla translator without Attention is straightforward to design. There is a decoder at the fourth step and this is another LSTM layer. By default, this takes the value of the last hidden state from the encoder and starts spitting out the translation for the first word (there is a Dense final layer on top which actually does this). To translate the second word, it uses the hidden state generated by the decoder in the previous time-step and also the just-translated word# 1 and generates the second word and a second hidden state. Similarly, the third word is generated with two inputs 1) hidden state generated from the decoder after time-step two & 2) the second word just translated. Note: I am giving a quick intuitive overview here and not covering specifics like end of sentence placeholders etc. Before proceeding further, you may want to recheck the reference links above to get a clear and intuitive understanding of how a plain decoder (without attention) works. The discussion henceforth focusses on the part where we add ‘Attention’ to a translator.

We now want to add ‘attention’ to this process. So instead of taking just the last hidden state as input, the decoder needs to take all hidden states generated by the encoder and pay attention to the ones that matter most. This is similar to the approach we have seen earlier in Yang’s paper — with one exception — instead of determining a ‘context’ through training, the ‘context’ (or more specifically ‘an indicator to the context’ which I call the pseudo-context) in this case is nothing but the last hidden state. The layer has to match all the hidden states generated by the encoder to this ‘pseudo-context’ to generate the attention weights. The ‘final context’ is the weighted average of all the hidden states. Thus, every step of the decoder now has 3 inputs instead of 2. The semantics can get a bit confusing here. So let us spend a couple of minutes to get a crystal-clear idea on the role of each player. We will use these semantics for the rest of the article.

• Player 1 — This is who I will call the pseudo-context vector or an ‘indicator of the context’. At the first step of the decoder, this is the person who holds the essence of the tweet. As the decoder proceeds to translate word by word, the decoder also outputs an ‘updated’ hidden state at each step and this ‘updated’ hidden state becomes the pseudo-context vector for the next step
• Player 2 — The alignment scores — These set of 19 scores are an indication of how similar or aligned each of the 19 hidden states of the encoder is to Player 1. Similarity of vectors are nothing but the dot product of the vectors. Higher the dot product, higher the similarity
• Player 3 — The attention weights — These are obtained from the alignment scores which are softmaxed to give the 19 attention weights
• Player 4 — This is the real context vector. The attention weights above are multiplied with the encoder hidden states and added to give us the real context or the ‘attention-adjusted’ output state. Going ahead, whenever we talk about context vector, this will mean Player 4.

Player 4, Player 1 are concatenated and fed to the next time-step along with the translated word of the previous step as input. The output is the next translated word and a new hidden state (which becomes player 1 for the next step). This process repeats 19 times. This is a raw essence of Attention in sequence to sequence models and both Bahdanau et al. and Luong et al. have variations over this essence.

We can see how a trail is built from Raffel (generate a set of attention weights, use it to calculate the ‘attention-adjusted’ state and use this to predict) to a version of Yang (generate a context, using that context generate a set of attention weights and use this to calculate the ‘attention-adjusted’ state and predict) to the above model (use the prev hidden state as a pseudo-context, using that pseudo-context, generate a set of attention weights and use this to calculate the ‘attention-adjusted’ state. Now concatenate the prev hidden state, the attention-adjusted state and the translated word and use this to predict). Surprisingly it transpires that the papers came in exactly the reverse timeline order!!!

We will spend a short time on Bahdanau & Luong variations as these are important papers and I didn't find a satisfying explanation of these approaches anywhere on the web. The best option is to go thru’ the papers themselves. Bahdanau et al. introduced ‘Attention’ for the first time here. They also tried to increase non-linearity in the model. The final output in Bahdanau continues to be a function of 3 variables listed below with one minor change:

• last translated word
• context vector
• decoder hidden state at time-step ‘t’ (surprise — it is not t-1)

The next surprise is in the logic for calculation of decoder hidden state at time-step ‘t’. They calculate this value somewhat indirectly. They use another simple function which takes 3 inputs:

• last translated word
• context vector
• decoder hidden state at time-step t-1

So basically Bahdanau has an additional step of calculating ‘decoder hidden state at time-step t’. Why do we need a separate function for this? I am not very sure, but when I tried going thru’ this paper by Pascanu et al., it seems that this kind of approach makes the simple RNN (which is sort of shallow, let us admit) deeper. So possibly, it may perform better.

There are a couple of other minor changes to the way the context vector is calculated. Like before, we have 2 key inputs to the function determining the context: the decoder hidden state at t-1 (s_tminus1) as the pseudo context and all the encoder hidden state outputs (let us call it hj instead of lstm_out to be more generic). The first change is that they recommend s_tminus1 value to be uni-directional, not bi-directional. hj values continue to be bi-directional. The second change is that there is no dot product between them involved here. They train and find 3 separate weights u,v,w such that the attention weights are the softmax of:

*v (w s_tminus1 + uhj)**

They continue to use single layer MLP’s to generate 3 weights — a weight w for the prev hidden state S_tminus1, a weight u for all the encoder hidden states(hj) and another mystery weight ‘v’. Giving a separate set of weights w, u sounds like a good idea, but I felt that there should be a dot product somewhere between the prev hidden state S_tminus1 and hj while calculating attention. There isn't one here. This is one aspect that I couldn't yet understand intuitively. It is just an addition of weighted values multiplied with yet another derived weight…hence the name additive attention. Maybe it has some other benefit that I am overlooking. Anyway, let us see how u,v,w are derived.

We discussed earlier that the problem of getting attention weights was a simple linear one and needs just 1 unit single MLP in the new layer. However, in Bahdanau we have a choice to use more than one unit to determine w and u — the weights that are applied individually on the decoder hidden state at t-1 and the encoder hidden states. Having done that, we need to massage the tensor shape back to (?,19,1) and hence there is a need for another weight v. Determining v is a simple linear transformation and needs just 1 unit. So you may see Bahdanau implementation in the following way to determine context:

class Bahdanau(tf.keras.layers.Layer):
    def __init__(self, n):
        super(Bahdanau, self).__init__()
        self.w = tf.keras.layers.Dense(n)
        self.u = tf.keras.layers.Dense(n)
        self.v = tf.keras.layers.Dense(1)

    def call(self, stminus1, h):
        stminus1 = tf.expand_dims(stminus1, 1)
        e = self.v(tf.nn.tanh(self.w(stminus1) + self.u(h)))

        a = tf.nn.softmax(e, axis=1)
        c = a * h
        c = tf.reduce_sum(c, axis=1)

        return a,c

# Say we want 10 units in the single layer MLP determining w,u
attentionlayer = Bahdanau(10)

# Call with i/p: decoderstate @ t-1 and all encoder hidden states
a, c = attentionlayer(stminus1, h)

So we have layers inside layers. Keras provides this kind of nice functionality. w, u are determined using a single layer MLP of ‘n’ units. ‘v’ needs a single unit-single layer MLP as discussed earlier. All the layer weights for these 3 feed forward layers are calculated based on the respective input shapes. All three feedforward layers do the default layer processing — it multiples the layer weights by the input and adds a bias if any. All of that need not be explicitly coded as it is the default behaviour in a Dense layer. Once u,v,w are determined, the context is calculated in call() and the rest of the code above is more or less straightforward.

Assume our old 19-word tweet conversion problem where we have an encoder which is bidirectional LSTM of 128 units and a decoder LSTM of 128 units. Assume n=10 units for the alignment layer to determine w,u. Then: the shapes for stminus1 and hj would be (?,128) and (?,19,256). Note that stminus1 is the single decoder hidden state at t-1 and hj are the 19 hidden states of the bi-directional encoder. We have to expand stminus1 to (?,1,128) for the addition that follows later along the time axis. The layer weights for w,u,v will be automatically determined by the framework as (?,128,10), (?,256,10) and (?,10,1) respectively. Notice how self.w(stminus1) works out to (?,1,10). This is added to each of the self.u(hj) to give a shape of (?,19,10). The result is fed to self.v and the output is (?,19,1) which is the shape we want. Softmaxing this gives the attention weights.

Multiplying this attention weight with each encoder hidden state and summing up returns the weighted context. This is the essence of Bahdanau. Of course, as discussed, we now need to use this freshly minted context, the last translated word and the decoder hidden state at t-1 to generate the decoder hidden state at t. Then we use this hidden state at t, along with context and the last translated word to finally generate our translated word. Phew!!

Luong et al (https://arxiv.org/pdf/1508.04025.pdf) introduce local attention. The idea of global attention (or soft attention) is to use all the hidden states of the encoder when computing each context vector which is what we have done in all the models explained so far. The downside of a global attention model is that it has to attend to all words on the source side for each target word, which is computationally costly. Local attention is an interesting mix of hard and soft attention. It first chooses a position in the source sentence. This position will determine a window of words that the model attends to. Calculating Local attention during training is slightly more complicated and requires techniques such as reinforcement learning to train. So there are some tradeoffs.

Luong also gives not one but 4 different ways to calculate context (hope you are keeping track of the number of flavours of attention) All four are minor variations of each other and similar to the approaches we have discussed earlier. In fact, I was able to relate to these approaches better. Most involve taking a dot product of hj (all encoder hidden states) with the decoder hidden state (s), hence these are usually called multiplicative attention. There are several other simplifications as well. They don’t have separate weights for encoder hidden states and decoder hidden state. Luong just adds them up and uses one weight to calculate the context. Luong also uses the hidden state of current time-step directly from the decoder and does not calculate it using a separate function based on <context, decoder hidden state t-1 and last translated word> like how Bahdanau does. So there is one function less. They also don't see why encoder needs to be uni-directional. They simply take the bi-directional RNN output from both encoder and decoder. They also suggest using only the topmost LSTM layer outputs and not worry about the others. These simplifications by Luong show that it is the concept of ‘Attention’ itself that matters and not the other adornments. But they made one other change which proved to generate excellent results and must be used when we write any Attention model.

The context generated by them is concatenated with the decoder hidden state at time t and then a new entity called attentional vector is calculated and this is percolated to the next step…. the idea is to inform the model about past alignment decisions. Intuitively, this strikes me as a very good idea irrespective of the Attention approach you are using and this is what is the crux. Based on the problem context, it becomes necessary to experiment with hybrid models of Attention inspired by various papers and this is the reason why we need to understand Attention in all its flavours!

Of course, you could always choose to implement the pre-written versions of Attention layers that come with the latest version of most frameworks… it then boils down to 1 line of code (well at least this is the way one would expect it to work):

# Input 1 = the last decoder hidden state: stminus1
# Input 2 = All hidden states of the encoder: lstm_out
# Apply Bahdanau additive attention and give me the 
# output = context

context = tf.keras.layers.AdditiveAttention()([stminus1, lstm_out])
# Well this should work but it may not. Read on for why

But there is a catch here. I was wondering where to specify the ’n’ units needed for the alignment layer in Bahdanau and tried digging thru’ the Keras source code for the above class. The first line I could see was: This class is suitable for Dense or CNN networks, and not for RNN networks. You would have imagined this type of disclaimer would have been given here but unfortunately, that is not the case. So use it cautiously. This code also uses the post 2018-Attention terminology and can be confusing to someone trying to map the Bahdanau or Luong papers directly to the above code.

In general, there are many hybrid versions of Attention on the web, claiming to implement Bahdanau or Luong. Most of these work because at the end of the day, there are only a couple of things that matter — the direct connection between the target and source (basic Attention in any form) and passing the Attention vector along to the next time step at the decoder (Luong). Pay attention to the scoring function as well. Surprisingly, the remaining nuts and bolts of the implementation don’t matter at all. One reliable guide to Attention that I found I could blindly trust is here.

Well, we are pretty much done now with all the key essences, but there is one last thing…

Attention as Memory?

If you have arrived at this page from Google, then you may have typed a query in the search box, the Google search engine would have mapped your query against a set of keys indexing all the available literature on the Internet, then present you the best-matched values (web-pages) and hopefully, this page-link would have been somewhere at the top. This is the query-key-value semantic that you may associate with easily.

Now, let us turn our attention back to ‘Attention’. There is another interesting way of looking at this whole ‘Attention’ business. Attention is nothing but simply giving the model access to a memory. In our case, memory is nothing but all the hidden states of the encoder. The model chooses what to retrieve from this memory. All this is not very different from traditional memory systems… the databases you may have used in your software programs. A major difference is that the memory access mechanism used here is ‘soft’. This basically means that the network retrieves a weighted combination of all memory locations when you query something, not a value from a single discrete location unlike your traditional databases. We can say that there is no single key to open the lock (to retrieve information). There are many keys and the lock is opened only when the relevant bunch of keys are used correctly. Looking back at the Attention models we have studied, we can deduce that query was the last hidden state of the decoder, the keys are the various hidden states of the encoder and the ‘attention adjusted output’ is the final response. In order to respond to a query, the model must grab a key from the memory. However as discussed, there is no perfect key here. Nonetheless, some keys help unlock the output better than the others. The model knows that the product between ‘query’ and ‘keys’ give the ‘match’ probability (remember when similar vectors are multiplied, we get higher scores). The sum of ‘value’ vectors weighted by this probability gives the final response.

*End-to-End Memory Networks* by Sukhbaatar et al. is an important paper in which the model uses such kind of a soft memory, making multiple hops and is trained end to end, over multiple sentences. Multiple hops mean that the network reads the input sequence multiple times before producing an output, and updates the memory contents at each step.

How else can ‘Attention’ be leveraged? Is there a benefit in ‘attending’ to the input sentence itself to create a more efficient memory database? This is an interesting question because we are not even talking of the decoder or the output sentence at this juncture. What we are doing here is to carefully pay attention to all parts of the input sentence itself, trying to capture how the words in the sentence best relate to one another. Maybe ‘Attention’ helps capture different possible nuances, meanings and relationships of each word in the input sentence with every other word in the input sentence in a better way and gives a better representation for each word as compared to LSTMs? This sort of self-introspection benefits humans and models alike and is called self-attention and if this step precedes all the rest of the decoder business, immense benefits can be seen. Cheng et al probably came out with the first version of self-attention saying “In our model, memory and attention are added within a sequence encoder allowing the network to uncover lexical relations between tokens” here. Interestingly enough they use the term intra-attention, not self-attention.

By the way, all this key-quey-value stuff is the same ol’ Attention we have been discussing all along. But this different view-point and new terminology will serve us better here-onwards, as we move on from our pitstop at Attention towards our next destination — BERT — the queen of all Embeddings.

When text-processing was reaping such immense benefits from ‘Attention’, would image processing remain behind? Definitely not and so we had a series of papers showing the relevance of ‘Attention’ in image processing as well. For e.g. this paper — Show, Attend, Tell by Xu et al. inspired by Bahdanau’s original paper, applies ‘attention’ to image captioning. Instead of using the whole image embedding(equivalent to the final LSTM hidden state in the NLP world) to come out with a caption using an RNN, they use the features obtained from the last but one layer of the CNN that has processed the image. These features represent various objects in the image (equivalent to all the encoder hidden states in NLP world). Now we apply attention using the last hidden state of the RNN generating the caption and ‘attend’ to all the above features, find out which one needs to be given more weightage and then sum them up to generate a context. This context now contains data on which all features (objects) in the image to pay attention to when generating the next word. This context along with the last generated word and the prev RNN hidden state is used to generate the next word of the caption.

Attention is important even if it occasionally doesn't produce better results because it always has the advantage of making the AI model more Explainable. Having said that, Attention does work well most of the time. In fact, it works so darn well that the RNN piece is now decoupled and the ‘attention’ piece (which was a decorator when originally conceived) becomes the centre-piece. Bahdanu’s paper in 2014, set in motion a chain of events that was to culminate in 2018 with a landmark moment for NLP… the likes of which probably come once in a decade. The stage is set in 2018 for the announcement of a revolutionary paper by a young team at Google. The name of the paper is ‘Attention is all you need’ by Vaswani et al. and its bare essence is as simple as its name. Attention made it possible for the rise of the transformers and it is now possible for a simple device in your pocket to translate the Dalai Lama’s live speech into any language that you want without the need of an interpreter!

We will peel the layers of these interesting innovations in the follow-up article.

PS — A well-meaning friend of mine, who has seen me sit multiple late nights and weekends to write these articles, advised me to keep it short. ‘People like 10-12 minute stories at max. For every minute you go beyond that, you lose a quarter of your potential readers’, he said. Well, I have gone on for a good 24 additional minutes and assuming 1000 readers to start out with, I have now brought down my readership to a grand total of one. I will await the weary Internet traveller who may someday land on this page, brimming with more questions than what I had when I started out, and hopefully, this article retains her/his ‘attention’ till the very end.

I would be more than satisfied with that one view.

Please feel free to connect with me on LinkedIn.