Add masking to GRU and LSTM

Author: wxs

keras/layers/recurrent.py

@@ -198,7 +198,7 @@ class GRU(Layer):
     def __init__(self, input_dim, output_dim=128, 
         init='glorot_uniform', inner_init='orthogonal',
         activation='sigmoid', inner_activation='hard_sigmoid',
-        weights=None, truncate_gradient=-1, return_sequences=False):
+        weights=None, truncate_gradient=-1, return_sequences=False, mask_val=default_mask_val):
 
         super(GRU,self).__init__()
         self.input_dim = input_dim
@@ -211,6 +211,7 @@ def __init__(self, input_dim, output_dim=128,
         self.activation = activations.get(activation)
         self.inner_activation = activations.get(inner_activation)
         self.input = T.tensor3()
+        self.mask_val = shared_scalar(default_mask_val)
 
         self.W_z = self.init((self.input_dim, self.output_dim))
         self.U_z = self.inner_init((self.output_dim, self.output_dim))
@@ -234,29 +235,35 @@ def __init__(self, input_dim, output_dim=128,
             self.set_weights(weights)
 
     def _step(self, 
-        xz_t, xr_t, xh_t, 
+        xz_t, xr_t, xh_t, mask_tm1,
         h_tm1, 
         u_z, u_r, u_h):
-        z = self.inner_activation(xz_t + T.dot(h_tm1, u_z))
-        r = self.inner_activation(xr_t + T.dot(h_tm1, u_r))
-        hh_t = self.activation(xh_t + T.dot(r * h_tm1, u_h))
-        h_t = z * h_tm1 + (1 - z) * hh_t
+        h_mask_tm1 = mask_tm1 * h_tm1
+        z = self.inner_activation(xz_t + T.dot(h_mask_tm1, u_z))
+        r = self.inner_activation(xr_t + T.dot(h_mask_tm1, u_r))
+        hh_t = self.activation(xh_t + T.dot(r * h_mask_tm1, u_h))
+        h_t = z * h_mask_tm1 + (1 - z) * hh_t
+        #return theano.printing.Print("h_t")(h_t)
         return h_t
 
     def get_output(self, train):
         X = self.get_input(train) 
         X = X.dimshuffle((1,0,2)) 
+        mask, padded_mask = get_mask(X, self.mask_val, steps_back=1)
 
         x_z = T.dot(X, self.W_z) + self.b_z
         x_r = T.dot(X, self.W_r) + self.b_r
         x_h = T.dot(X, self.W_h) + self.b_h
         outputs, updates = theano.scan(
             self._step, 
-            sequences=[x_z, x_r, x_h], 
+            sequences=[x_z, x_r, x_h, padded_mask], 
             outputs_info=T.unbroadcast(alloc_zeros_matrix(X.shape[1], self.output_dim), 1),
             non_sequences=[self.U_z, self.U_r, self.U_h],
             truncate_gradient=self.truncate_gradient
         )
+
+        outputs = mask * outputs + (1 - mask) * self.mask_val
+
         if self.return_sequences:
             return outputs.dimshuffle((1,0,2))
         return outputs[-1]
@@ -302,13 +309,14 @@ class LSTM(Layer):
     def __init__(self, input_dim, output_dim=128, 
         init='glorot_uniform', inner_init='orthogonal', 
         activation='tanh', inner_activation='hard_sigmoid',
-        weights=None, truncate_gradient=-1, return_sequences=False):
+        weights=None, truncate_gradient=-1, return_sequences=False, mask_val=default_mask_val):
 
         super(LSTM,self).__init__()
         self.input_dim = input_dim
         self.output_dim = output_dim
         self.truncate_gradient = truncate_gradient
         self.return_sequences = return_sequences
+        self.mask_val = shared_scalar(mask_val)
 
         self.init = initializations.get(init)
         self.inner_init = initializations.get(inner_init)
@@ -343,19 +351,23 @@ def __init__(self, input_dim, output_dim=128,
             self.set_weights(weights)
 
     def _step(self, 
-        xi_t, xf_t, xo_t, xc_t, 
+        xi_t, xf_t, xo_t, xc_t, mask_tm1,
         h_tm1, c_tm1, 
         u_i, u_f, u_o, u_c): 
-        i_t = self.inner_activation(xi_t + T.dot(h_tm1, u_i))
-        f_t = self.inner_activation(xf_t + T.dot(h_tm1, u_f))
-        c_t = f_t * c_tm1 + i_t * self.activation(xc_t + T.dot(h_tm1, u_c))
-        o_t = self.inner_activation(xo_t + T.dot(h_tm1, u_o))
+        h_mask_tm1 = mask_tm1 * h_tm1
+        c_mask_tm1 = mask_tm1 * c_tm1
+
+        i_t = self.inner_activation(xi_t + T.dot(h_mask_tm1, u_i))
+        f_t = self.inner_activation(xf_t + T.dot(h_mask_tm1, u_f))
+        c_t = f_t * c_mask_tm1 + i_t * self.activation(xc_t + T.dot(h_mask_tm1, u_c))
+        o_t = self.inner_activation(xo_t + T.dot(h_mask_tm1, u_o))
         h_t = o_t * self.activation(c_t)
         return h_t, c_t
 
     def get_output(self, train):
         X = self.get_input(train) 
         X = X.dimshuffle((1,0,2))
+        mask, padded_mask = get_mask(X, self.mask_val, steps_back=1)
 
         xi = T.dot(X, self.W_i) + self.b_i
         xf = T.dot(X, self.W_f) + self.b_f
@@ -364,14 +376,16 @@ def get_output(self, train):
 
         [outputs, memories], updates = theano.scan(
             self._step, 
-            sequences=[xi, xf, xo, xc],
+            sequences=[xi, xf, xo, xc, padded_mask],
             outputs_info=[
                 T.unbroadcast(alloc_zeros_matrix(X.shape[1], self.output_dim), 1),
                 T.unbroadcast(alloc_zeros_matrix(X.shape[1], self.output_dim), 1)
             ], 
             non_sequences=[self.U_i, self.U_f, self.U_o, self.U_c], 
             truncate_gradient=self.truncate_gradient 
         )
+
+        outputs = mask * outputs + (1 - mask) * self.mask_val
         if self.return_sequences:
             return outputs.dimshuffle((1,0,2))
         return outputs[-1]

tests/manual/check_masked_recurrent.py

@@ -6,7 +6,7 @@
 from keras.models import Sequential
 from keras.layers.core import Dense, Activation, Merge, Dropout, TimeDistributedDense
 from keras.layers.embeddings import Embedding
-from keras.layers.recurrent import SimpleRNN, SimpleDeepRNN
+from keras.layers.recurrent import SimpleRNN, SimpleDeepRNN, LSTM, GRU
 from keras.layers.core import default_mask_val
 import theano
 
@@ -15,6 +15,7 @@
 # (nb_samples, timesteps, dimensions)
 X = np.random.random_integers(1, 4, size=(500000,15))
 
+print("About to compile the first model")
 model = Sequential()
 model.add(Embedding(5, 4, zero_is_mask=True))
 model.add(TimeDistributedDense(4,4)) # obviously this is redundant. Just testing.
@@ -108,3 +109,26 @@
 if score < uniform_score*0.9:
     raise Exception('Somehow learned to copy timestep 0 despite masking 1, score %f' % score)
 
+
+# Another testing approach, just initialize models and make sure that prepending zeros doesn't affect
+# their output
+print("About to compile the second model")
+model2 = Sequential()
+model2.add(Embedding(5, 4, zero_is_mask=True))
+model2.add(TimeDistributedDense(4,4)) # obviously this is redundant. Just testing.
+model2.add(LSTM(4,4, return_sequences=True))
+model2.add(GRU(4,4, activation='softmax', return_sequences=True))
+model2.add(SimpleDeepRNN(4,4, depth=2, activation='relu', return_sequences=True)) 
+model2.add(SimpleRNN(4,4, activation='relu', return_sequences=False))
+model2.compile(loss='categorical_crossentropy',
+        optimizer='rmsprop', theano_mode=theano.compile.mode.FAST_RUN)
+print("Compiled model2")
+
+X2 = np.random.random_integers(1, 4, size=(1,5))
+ref = model2.predict(X2)
+for pre_zeros in range(1,10):
+    padded = np.concatenate((np.zeros((1, pre_zeros)), X2), axis=1)
+    pred = model2.predict(padded)
+    if not np.allclose(ref, pred):
+        raise Exception("Different result after left-padding %d zeros. Ref: %s, Pred: %s" % (pre_zeros, ref, pred))
+