title: How can I speed up this Keras Attention computation?
body: <p>I have written a custom keras layer for an <code>AttentiveLSTMCell</code> and <code>AttentiveLSTM(RNN)</code> in line with keras' <em>new</em> approach to RNNs. This attention mechanism is described by <a href="https://arxiv.org/abs/1409.0473" rel="noreferrer">Bahdanau</a> where, in an encoder/decoder model a "context" vector is created from all the ouputs of the encoder and the decoder's current hidden state. I then append the context vector, at every timestep, to the input.</p>
  
  <p>The model is being used in to make a Dialog Agent, but is very similar to NMT models in architecture (similar tasks).</p>
  
  <p>However, in adding this attention mechanism, I have slowed down the training of my network 5 fold, and I would really like to know how I could write the part of the code that is slowing it down so much in a more efficient way. </p>
  
  <p>The brunt of the computation is done here:</p>
  
  <pre><code>h_tm1 = states[0]  # previous memory state
  c_tm1 = states[1]  # previous carry state
  
  # attention mechanism
  
  # repeat the hidden state to the length of the sequence
  _stm = K.repeat(h_tm1, self.annotation_timesteps)
  
  # multiplty the weight matrix with the repeated (current) hidden state
  _Wxstm = K.dot(_stm, self.kernel_w)
  
  # calculate the attention probabilities
  # self._uh is of shape (batch, timestep, self.units)
  et = K.dot(activations.tanh(_Wxstm + self._uh), K.expand_dims(self.kernel_v))
  
  at = K.exp(et)
  at_sum = K.sum(at, axis=1)
  at_sum_repeated = K.repeat(at_sum, self.annotation_timesteps)
  at /= at_sum_repeated  # vector of size (batchsize, timesteps, 1)
  
  # calculate the context vector
  context = K.squeeze(K.batch_dot(at, self.annotations, axes=1), axis=1)
  
  # append the context vector to the inputs
  inputs = K.concatenate([inputs, context])
  </code></pre>
  
  <p>in the <code>call</code> method of the <code>AttentiveLSTMCell</code> (one timestep).</p>
  
  <p>The full code can be found <a href="https://gist.github.com/DylanModesitt/c54af5b5531c0d4cc57a8b5ab0733437" rel="noreferrer">here</a>. If it is necessary that I provide some data and ways to interact with the model, then I can do that. </p>
  
  <p>Any ideas? I am, of course, training on a GPU if there is something clever here.</p>