TimeDistributed LSTM
Y label shape
machine learning
neural networks
deep learning

Y label shape for time_distributed lstm

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

With a time-distributed LSTM model, the correct target shape depends on what the model outputs at each time step. If the network returns one prediction for every step, your y labels must also have one target for every step. If the network collapses the sequence to one final output, then y should not have a time dimension at all.

Start From The Model Output Shape

The easiest rule is: make y match the model's output shape, excluding only the batch dimension.

A common sequence-to-sequence setup looks like this:

python
1import tensorflow as tf
2from tensorflow.keras import layers
3
4model = tf.keras.Sequential([
5    layers.LSTM(32, return_sequences=True, input_shape=(10, 8)),
6    layers.TimeDistributed(layers.Dense(4))
7])
8
9model.summary()

This model takes input with shape (batch, 10, 8) and produces output with shape (batch, 10, 4).

So the target labels must also have shape (batch, 10, 4).

Sequence-To-Sequence Labels

If you want one target vector per time step, y needs three dimensions:

  • batch size
  • time steps
  • output features or classes
python
1import numpy as np
2
3x = np.random.randn(32, 10, 8).astype("float32")
4y = np.random.randn(32, 10, 4).astype("float32")
5
6model.compile(optimizer="adam", loss="mse")
7model.fit(x, y, epochs=2)

This is the correct pattern for time-distributed regression where every time step has its own numeric target vector.

For time-distributed classification with one-hot targets, the last dimension would be the number of classes instead of regression outputs.

Sequence-To-One Is Different

If the LSTM does not return sequences, the output no longer includes the time axis.

python
1import tensorflow as tf
2from tensorflow.keras import layers
3
4model = tf.keras.Sequential([
5    layers.LSTM(32, return_sequences=False, input_shape=(10, 8)),
6    layers.Dense(4)
7])
8
9model.summary()

Now the output shape is (batch, 4), so y should be shaped (batch, 4).

That is a common source of confusion: people keep the old 3D label tensor even though the model produces only one output per sequence.

Sparse Classification Targets

If your last time-distributed layer predicts class probabilities with softmax, your label format depends on the loss.

With one-hot labels:

python
1model = tf.keras.Sequential([
2    layers.LSTM(32, return_sequences=True, input_shape=(10, 8)),
3    layers.TimeDistributed(layers.Dense(5, activation="softmax"))
4])
5
6x = np.random.randn(16, 10, 8).astype("float32")
7y = tf.one_hot(np.random.randint(0, 5, size=(16, 10)), depth=5)
8
9model.compile(optimizer="adam", loss="categorical_crossentropy")
10model.fit(x, y, epochs=1)

With sparse integer class labels, the target shape can be (batch, time_steps) when using sparse_categorical_crossentropy.

The key point is still the same: labels must match the output semantics of the model.

Why return_sequences=True Matters

TimeDistributed on top of an LSTM only makes sense for per-time-step outputs if the LSTM is returning the full sequence.

If return_sequences=False, the LSTM produces only the last hidden state, so there is no time axis left for TimeDistributed to operate over in the way you probably intended.

That means the shape problem is often caused earlier in the model than in the labels themselves.

Common Pitfalls

The most common mistake is giving a 2D label array to a model whose output is 3D, or the reverse. Always inspect model.output_shape before shaping y.

Another issue is forgetting that return_sequences=True changes the LSTM output from (batch, units) to (batch, time_steps, units).

It is also easy to mix one-hot and sparse labels with the wrong loss function. If the label rank looks right but training still fails, check whether the loss expects integers or full class vectors.

Finally, do not guess the label shape from the input shape alone. The output architecture determines the target format.

Summary

  • The target y shape should match the model output shape, excluding the batch dimension.
  • For time-distributed outputs, labels are usually 3D: (batch, time_steps, output_dim).
  • For sequence-to-one outputs, labels are usually 2D: (batch, output_dim).
  • 'return_sequences=True is required when you want one prediction per time step.'
  • Check model.output_shape first instead of guessing label dimensions.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.