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.
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:
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
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.
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:
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
yshape 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=Trueis required when you want one prediction per time step.' - Check
model.output_shapefirst instead of guessing label dimensions.
Related reading
- YOLO object detection how does the algorithm predict bounding boxes larger than a grid cell?
- 4D input in LSTM layer in Keras
- About tensorflow Metadata and RunOptions
- About tensorflow.initialize_all_variables
- Yolo object detection include images that do not contain classes to be predicted?
- Yolo v1 bounding boxes during training step
- Accessing PyTorch GPU matrix from TensorFlow directly
- Accuracy issue in caffe
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.