TimeDistributed(Dense()) vs Dense() after lstm

Viewed 125
input_word = Input(shape=(max_len,))
model = Embedding(input_dim=num_words, output_dim=50, input_length=max_len)(input_word)
model = SpatialDropout1D(0.1)(model)
model = Bidirectional(LSTM(units=100, return_sequences=True, recurrent_dropout=0.1))(model)
out = TimeDistributed(Dense(num_tags, activation="softmax"))(model)
#out = Dense(num_tags, activation="softmax")(model)
model = Model(input_word, out)
model.summary()

I get the same result when I use just Dense layer or with TimeDistributed. In which case should I use TimeDistributed?

2 Answers

TimeDistributed is only necessary for certain layers that cannot handle additional dimensions in their implementation. E.g. MaxPool2D only works on 2D tensors (shape batch x width x height x channels) and will crash if you, say, add a time dimension:

tfkl = tf.keras.layers
a = tf.random.normal((16, 32, 32, 3))
tfkl.MaxPool2D()(a)  # this works

a = tf.random.normal((16, 5, 32, 32, 3))  # added a 5th dimension
tfkl.MaxPool2D()(a)  # this will crash

Here, adding TimeDistributed will fix it:

tfkl.TimeDistributed(tfkl.MaxPool2D())(a)  # works with a being 5d!

However, many layers already support arbitrary input shapes and will automatically distribute the computations over those dimensions. One of these is Dense -- it is always applied to the last axis in your input and distributed over all others, so TimeDistributed isn't necessary. In fact, as you noted, it changes nothing about the output.

Still, it may change how exactly the computation is done. I'm not sure about this, but I would wager that not using TimeDistributed and relying on the Dense implementation itself may be more efficient.

According to the book Zero to Deep Learning by Francesco Mosconi in chapter 7:

If we want the model return an output sequence to be compared with the sequence of values in the labels, we will use the TimeDistributed layer wrapper around our output Dense layer. This method of training is called Teacher Forcing. If we didn’t create output sequences we wouldn't need Teacher Forcing(i.e. wouldn't need TimeDistributed wrapper).

Related