Keras Attention Mechanism
Simple attention mechanism implemented in Keras for the following layers:
- Dense (attention 2D block)
- LSTM, GRU (attention 3D block)
inputs = Input(shape=(input_dims,)) attention_probs = Dense(input_dims, activation='softmax', name='attention_probs')(inputs) attention_mul = merge([inputs, attention_probs], output_shape=input_dims, name='attention_mul', mode='mul')
Let's consider this Hello World example:
- A vector v of 32 values as input to the model (simple feedforward neural network).
- v = target.
- Target is binary (either 0 or 1).
- All the other values of the vector v (v and v[2:32]) are purely random and do not contribute to the target.
We expect the attention to be focused on v only, or at least strongly. We recap the setup with this drawing:
The first two are samples taken randomly from the training set. The last plot is the attention vector that we expect. A high peak indexed by 1, and close to zero on the rest.
Let's train this model and visualize the attention vector applied to the inputs:
We can clearly see that the network figures this out for the inference.
Behind the scenes
The attention mechanism can be implemented in three lines with Keras:
inputs = Input(shape=(input_dims,)) attention_probs = Dense(input_dims, activation='softmax', name='attention_probs')(inputs) attention_mul = merge([inputs, attention_probs], output_shape=32, name='attention_mul', mode='mul')
We apply a
Dense - Softmax layer with the same number of output parameters than the
Input layer. The attention matrix has a shape of
input_dims x input_dims here.
Then we merge the
Inputs layer with the attention layer by multiplying element-wise.
Finally, the activation vector (probability distribution) can be derived with:
attention_vector = get_activations(m, testing_inputs_1, print_shape_only=True).flatten()
1 is the index of definition of the attention layer in the model definition (
Inputs is indexed by
Recurrent Layers (LSTM, GRU...)
Application of attention at input level
We consider the same example as the one used for the Dense layers. The attention index is now on the 10th value. We therefore expect an attention spike around this value. There are two main ways to apply attention to recurrent layers:
- Directly on the inputs (same as the Dense example above):
APPLY_ATTENTION_BEFORE_LSTM = True
Application of attention on the LSTM's output
- After the LSTM layer:
APPLY_ATTENTION_BEFORE_LSTM = False
Both have their own advantages and disadvantages. One obvious advantage of applying the attention directly at the inputs is that we clearly understand this space. The high dimensional space spanned by the LSTM might be a bit trickier to interpret, although they share the time steps in common with the inputs (
return_sequences=True is used here).
Attention of multi dimensional time series
Also, sometimes, the time series can be N-dimensional. It could be interesting to have one attention vector per dimension. Let's say we have a 2-D time series on 20 steps. Setting
SINGLE_ATTENTION_VECTOR = False gives an attention vector of shape
(20, 2). If
SINGLE_ATTENTION_VECTOR is set to
True, it means that the attention vector will be of shape
(20,) and shared across the input dimensions.
SINGLE_ATTENTION_VECTOR = False
SINGLE_ATTENTION_VECTOR = True