Implementing Attention in Tensorflow

Question

I would like to check if my attention implementation is correct in TensorFlow.

Basically, I am using the attention mentioned in https://arxiv.org/pdf/1509.06664v1.pdf. (Just the baseline attention, not the word-for-word attention). So far I implemented it without using the last hidden state h_N.

def attention(hidden_states):
    '''
    hidden states (inputs) are seq_len x batch_size x dim 
    returns r, the weighted representation of the hidden states by attention vector a

    Note:I do not use the h_N vector and also I skip the last projection layer. 
    '''
    shape = hidden_states.get_shape().as_list()

    with tf.variable_scope('attention_{}'.format(name), reuse=reuse) as f:
        initializer = tf.random_uniform_initializer()

        # Initialize Parameters
        weights_Y = tf.get_variable("weights_Y", [self.in_size, self.in_size], initializer=initializer)
        weights_w = tf.get_variable("weights_w", [self.in_size, 1], initializer=initializer)

        # Equation => M  = tanh(W^{Y}Y)
        tmp = tf.reshape(hidden_states, [-1, shape[2]])
        Y = tf.matmul(tmp, weights_Y)
        Y = tf.reshape(Y, [shape[0], -1, shape[2]])
        Y = tf.tanh(Y, name='M_matrix')

        # Equation => a = softmax(Y w^T)
        Y = tf.reshape(Y, [-1, shape[2]])
        a = tf.matmul(Y, weights_w)
        a = tf.reshape(a, [-1, shape[0]])
        a = tf.nn.softmax(a, name='attention_vector')

        # Equation => r = Ya^T
        # This is the part I weight all hidden states by the attention vector
        a = tf.expand_dims(a, 2)
        H = tf.transpose(hidden_states, [1,2,0])  
        r = tf.matmul(H, a, name='R_vector')
        r = tf.reshape(r, [-1, shape[2]])

        # I skip the last projection layer  since I do not use h_N
        return r

This graph compiles, runs and trains properly. (loss is decreasing etc) but the performance is lower than I would expect. I would appreciate if I could get a check if I'm doing it right.

Generally,

1) For multiplications that are [?, seq_len, dim] matrix multiplied by [dim, dim]. Is it correct to use tf.reshape from [?, seq_len, dim] to [-1,dim] and apply matmul of shapes [-1, dim] with [dim, dim] and then reshape back to [?, seq_len, dim] after matmul?

2) I notice I obtain (?, seq_len) for a, the attention vector. So I am required to do (?, seq_len) x (?, dim, seq_len).

Is it correct to cast and expand_dims from (?, seq_len) to (?, seq_len, 1) and then do a matmul (I think this was what batch_matmul does in previous versions).

Thanks in advance!

Answer 1

Not sure tf.einsum in TF1.0 is implemented efficiently, but it would make the computation quite elegant.

import tensorflow as tf
import numpy as np

batch_size = 3
seq_len = 5
dim = 2
# [batch_size x seq_len x dim]  -- hidden states
Y = tf.constant(np.random.randn(batch_size, seq_len, dim), tf.float32)
# [batch_size x dim]            -- h_N
h = tf.constant(np.random.randn(batch_size, dim), tf.float32)

initializer = tf.random_uniform_initializer()
W = tf.get_variable("weights_Y", [dim, dim], initializer=initializer)
w = tf.get_variable("weights_w", [dim], initializer=initializer)

# [batch_size x seq_len x dim]  -- tanh(W^{Y}Y)
M = tf.tanh(tf.einsum("aij,jk->aik", Y, W))
# [batch_size x seq_len]        -- softmax(Y w^T)
a = tf.nn.softmax(tf.einsum("aij,j->ai", M, w))
# [batch_size x dim]            -- Ya^T
r = tf.einsum("aij,ai->aj", Y, a)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    a_val, r_val = sess.run([a, r])
    print("a:", a_val, "\nr:", r_val)