TimeDistributed(Dense) vs Dense in Keras - Same number of parameters

Question

I'm building a model that converts a string to another string using recurrent layers (GRUs). I have tried both a Dense and a TimeDistributed(Dense) layer as the last-but-one layer, but I don't understand the difference between the two when using return_sequences=True, especially as they seem to have the same number of parameters.

My simplified model is the following:

InputSize = 15
MaxLen = 64
HiddenSize = 16

inputs = keras.layers.Input(shape=(MaxLen, InputSize))
x = keras.layers.recurrent.GRU(HiddenSize, return_sequences=True)(inputs)
x = keras.layers.TimeDistributed(keras.layers.Dense(InputSize))(x)
predictions = keras.layers.Activation('softmax')(x)

The summary of the network is:

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         (None, 64, 15)            0         
_________________________________________________________________
gru_1 (GRU)                  (None, 64, 16)            1536      
_________________________________________________________________
time_distributed_1 (TimeDist (None, 64, 15)            255       
_________________________________________________________________
activation_1 (Activation)    (None, 64, 15)            0         
=================================================================

This makes sense to me as my understanding of TimeDistributed is that it applies the same layer at all timepoints, and so the Dense layer has 16*15+15=255 parameters (weights+biases).

However, if I switch to a simple Dense layer:

inputs = keras.layers.Input(shape=(MaxLen, InputSize))
x = keras.layers.recurrent.GRU(HiddenSize, return_sequences=True)(inputs)
x = keras.layers.Dense(InputSize)(x)
predictions = keras.layers.Activation('softmax')(x)

I still only have 255 parameters:

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         (None, 64, 15)            0         
_________________________________________________________________
gru_1 (GRU)                  (None, 64, 16)            1536      
_________________________________________________________________
dense_1 (Dense)              (None, 64, 15)            255       
_________________________________________________________________
activation_1 (Activation)    (None, 64, 15)            0         
=================================================================

I wonder if this is because Dense() will only use the last dimension in the shape, and effectively treat everything else as a batch-like dimension. But then I'm no longer sure what the difference is between Dense and TimeDistributed(Dense).

Update Looking at https://github.com/fchollet/keras/blob/master/keras/layers/core.py it does seem that Dense uses the last dimension only to size itself:

def build(self, input_shape):
    assert len(input_shape) >= 2
    input_dim = input_shape[-1]

    self.kernel = self.add_weight(shape=(input_dim, self.units),

It also uses keras.dot to apply the weights:

def call(self, inputs):
    output = K.dot(inputs, self.kernel)

The docs of keras.dot imply that it works fine on n-dimensional tensors. I wonder if its exact behavior means that Dense() will in effect be called at every time step. If so, the question still remains what TimeDistributed() achieves in this case.

Solution

TimeDistributedDense applies a same dense to every time step during GRU/LSTM Cell unrolling. So the error function will be between predicted label sequence and the actual label sequence. (Which is normally the requirement for sequence to sequence labeling problems).

However, with return_sequences=False, Dense layer is applied only once at the last cell. This is normally the case when RNNs are used for classification problem. If return_sequences=True then Dense layer is applied to every timestep just like TimeDistributedDense.

So for as per your models both are same, but if you change your second model to return_sequences=False, then Dense will be applied only at the last cell. Try changing it and the model will throw as error because then the Y will be of size [Batch_size, InputSize], it is no more a sequence to sequence but a full sequence to label problem.

from keras.models import Sequential
from keras.layers import Dense, Activation, TimeDistributed
from keras.layers.recurrent import GRU
import numpy as np

InputSize = 15
MaxLen = 64
HiddenSize = 16

OutputSize = 8
n_samples = 1000

model1 = Sequential()
model1.add(GRU(HiddenSize, return_sequences=True, input_shape=(MaxLen, InputSize)))
model1.add(TimeDistributed(Dense(OutputSize)))
model1.add(Activation('softmax'))
model1.compile(loss='categorical_crossentropy', optimizer='rmsprop')


model2 = Sequential()
model2.add(GRU(HiddenSize, return_sequences=True, input_shape=(MaxLen, InputSize)))
model2.add(Dense(OutputSize))
model2.add(Activation('softmax'))
model2.compile(loss='categorical_crossentropy', optimizer='rmsprop')

model3 = Sequential()
model3.add(GRU(HiddenSize, return_sequences=False, input_shape=(MaxLen, InputSize)))
model3.add(Dense(OutputSize))
model3.add(Activation('softmax'))
model3.compile(loss='categorical_crossentropy', optimizer='rmsprop')

X = np.random.random([n_samples,MaxLen,InputSize])
Y1 = np.random.random([n_samples,MaxLen,OutputSize])
Y2 = np.random.random([n_samples, OutputSize])

model1.fit(X, Y1, batch_size=128, nb_epoch=1)
model2.fit(X, Y1, batch_size=128, nb_epoch=1)
model3.fit(X, Y2, batch_size=128, nb_epoch=1)

print(model1.summary())
print(model2.summary())
print(model3.summary())

In the above example architecture of model1 and model2 are sample (sequence to sequence models) and model3 is a full sequence to label model.

Comments

[thon]

Thank you for the answer. I'm not sure I can follow though as I know that the output in both cases is a sequence. In both cases the recurrent layer has return_sequences=True, and the output shape in both cases is 3D and is exactly the same (batch_size, 64, 15). So it seems to me that the Dense layer is also applied at every time step.

[mujjiga]

I have updated my answer with better explanation, hope it helps you.

[thon]

Thank you. For avoidance of doubt, when you say "So for as per your models both are same, but if u change your second model to "return_sequences=True" then the Dense will be applied only at the last cell." - do you mean if I change return_sequences to False? Your answer seems to imply that if return_sequences is True, Dense() and TimeDistributed(Dense()) do exactly the same thing. Could you confirm this? This would make sense, but then why does Keras need TimeDistributed() at all?

[mujjiga]

Sorry for typo I have corrected it now, yes it should be "return_sequences=False". As per keras documentation "...This wrapper (TimeDistributed) allows to apply a layer to every temporal slice of an input." Using TimeDistributed you can apply the same layer but RNNs are by default unrolled in time. See this code snippet "model = Sequential() model.add(TimeDistributed(Conv2D(64, (3, 3)), input_shape=(10, 299, 299, 3)))". You cant achieve this using a simple dense layer without TimeDistributed wrapper"

[thon]

Thanks again. Yes, I agree that one needs TimeDistributed() for other layer types. It seems to me that a simple Dense() after a recurrent layer that returns sequences does work, but more by accident than by design. From old Keras examples I think there used to be a TimeDistributedDense() - it's still a mystery to me why it was needed if Dense() would have worked anyway.

[Gengiolo]

Hi @thon, I've come to the same conclusion. It's also very strange because Dense() should flatten the input dimensions if > 2 as cited in the doc: "Note: if the input to the layer has a rank greater than 2, then it is flattened prior to the initial dot product with kernel". Have you found the answer?

[thon]

Hi, thanks for the comment. No, unfortunately I haven't managed to find out anything more about TimeDistributed vs Dense.

[OverLordGoldDragon]

This is false; session graph results aren't equivalent to the full model graph results - latter involves gradients and weight updates. From documentation, "if the input to the layer has a rank greater than 2, then it is flattened prior to the initial dot product with kernel" - contrasting TimeDistributedDense, which doesn't flatten. Counterexample code