Multi dimensional input for LSTM in Keras
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
Keras LSTM layers expect sequence data, not a flat feature table. The key idea is that each training example is a sequence of timesteps, and each timestep contains one or more features, so the standard input shape is three-dimensional.
The Shape Keras Expects
For an LSTM in Keras or tf.keras, the input tensor shape is:
- '
samples' - '
timesteps' - '
features'
When you define the layer, you specify only timesteps and features, because the batch dimension is handled automatically.
That model expects each sample to contain 10 time steps and 5 numeric features at each step.
What "Multidimensional" Usually Means
In most practical questions, "multidimensional input" means one of two things.
The first meaning is simple multivariate sequence data, such as temperature, pressure, and humidity values observed at each time step. That fits directly into the normal LSTM shape of (samples, timesteps, features).
The second meaning is more complex structured data at each step, such as an image frame, a grid, or a vector of embeddings. In those cases, you often need an additional preprocessing layer before the LSTM or a different recurrent architecture.
A Runnable Multivariate Example
This example trains an LSTM on synthetic data with 100 samples, 12 time steps, and 3 features per step.
The important detail is not the toy target. It is the shape of a, which is (100, 12, 3).
Stacking LSTM Layers
If you want multiple LSTM layers, the first recurrent layer must return the full sequence, not just the final output. That is what return_sequences=True does.
Without return_sequences=True on the first layer, the second LSTM would receive a two-dimensional tensor instead of a sequence.
What If Each Timestep Is Itself High-Dimensional?
Sometimes each timestep is not just a feature vector. It may be an image, spectrogram slice, or other structured tensor.
For example, if each timestep is a 32 x 32 image with 3 channels, the raw input shape would conceptually be (samples, timesteps, height, width, channels). A plain LSTM does not consume that directly.
Common solutions are:
- extract per-frame features with
TimeDistributedand a smaller network - use
ConvLSTM2Dif the sequence is genuinely spatiotemporal - reshape only if the flattening still preserves the information you care about
A TimeDistributed example looks like this.
That converts each frame into a vector before the LSTM processes the temporal sequence.
Sequence Length and Padding
Real datasets often contain sequences of unequal length. Keras can still handle them if you pad shorter sequences to a common length and optionally use masking.
This is useful when some timesteps are just placeholders added to make the batch rectangular.
Common Pitfalls
The most common mistake is swapping the timesteps and features axes. If your model expects (12, 3) and you pass (3, 12), training may still run but the sequence meaning is wrong.
Another mistake is flattening the whole sample before the LSTM. That destroys the temporal structure the recurrent layer is supposed to learn.
A third issue is forgetting return_sequences=True when stacking recurrent layers.
Finally, if your per-timestep data is image-like or grid-like, a plain LSTM may be the wrong layer. Consider TimeDistributed, CNN feature extraction, or ConvLSTM2D instead.
Summary
- A standard Keras LSTM expects input shaped as
(samples, timesteps, features). - Multivariate sequence data fits naturally into that three-dimensional structure.
- Use
return_sequences=Truewhen feeding one LSTM into another. - Use masking or padding when sequence lengths vary.
- For image-like data per timestep, preprocess with
TimeDistributedor useConvLSTM2D. - Most shape errors come from mixing up axes or flattening away the time dimension.
Related reading
- Multi GPU architecture, gradient averaging - less accurate model?
- Multi GPU Training in Tensorflow Data Parallelism when Using feed_dict
- Multi label regression in Caffe
- Multilabel Text Classification using TensorFlow
- MultiHeadAttention attention_mask Keras, Tensorflow example
- MultiHeadAttention attention_mask Keras, Tensorflow example
- Multilayer Seq2Seq model with LSTM in Keras
- Multiple embedding layers in keras
.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.