How to extract cell state from a LSTM at each timestep 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
In Keras, an LSTM layer can return the output sequence and the final recurrent states, but it does not expose the internal cell state for every timestep through a single built-in flag. If you need the full c_t sequence for visualization, debugging, or research, you have to run the recurrence at the cell level.
What the Standard LSTM Layer Returns
The public tf.keras.layers.LSTM API gives you two options that are easy to confuse:
- '
return_sequences=Truereturns the hidden output for every timestep' - '
return_state=Truereturns only the final hidden state and final cell state'
A quick example makes the distinction clear:
This prints shapes like these:
So Keras gives you the output at every step and the final cell state, but not the full sequence of cell states.
Why the Cell State Is Different
An LSTM carries two recurrent tensors:
- the hidden state
h_t, which is the exposed output - the cell state
c_t, which is the internal memory
Most downstream layers consume h_t, so the high-level Keras layer is designed around that output. The internal memory is still there, but it is treated as implementation detail except for the final state when return_state=True is enabled.
That design is fine for ordinary sequence modeling. It becomes limiting only when you want to inspect how memory evolves across time.
Extract the Cell State Sequence with LSTMCell
The clean approach is to use tf.keras.layers.LSTMCell and explicitly collect the state after each step. That gives you direct access to both h_t and c_t.
Output:
That tensor is exactly the full cell state sequence.
Integrating It into a Keras Model
If you want this behavior inside a normal functional model, wrap the custom logic in a reusable layer and expose it as a model output.
This pattern works well when you want to log intermediate state dynamics during experiments.
When You Do Not Need This
Many people ask for the cell state sequence when they actually only need the output sequence. If you are feeding the LSTM into another layer, computing attention over timesteps, or plotting predictions over time, return_sequences=True is usually enough.
Reach for a custom cell-level loop only when the internal memory itself matters.
Common Pitfalls
The most common mistake is assuming return_sequences=True exposes the cell state sequence. It does not. It exposes the hidden output sequence.
Another mistake is expecting tf.keras.layers.RNN(tf.keras.layers.LSTMCell(...)) to automatically publish the full c_t history. It still follows the standard RNN return contract unless you write custom collection logic.
A third issue is building a Python loop and then not checking tensor shapes. Always test with a small random tensor first so you can confirm the returned shape is (batch, timesteps, units).
Finally, remember that the cell state is an internal recurrent feature. It can be useful for analysis, but it is not automatically the best representation for downstream modeling.
Summary
- A standard Keras
LSTMdoes not directly return the full cell state sequence - '
return_sequences=Truegives you the hidden outputs, not the cell memory' - '
return_state=Truegives only the final hidden and final cell states' - Use
LSTMCelland collectc_tmanually when you need every timestep - Wrap that logic in a custom layer if you want it to fit cleanly into a Keras model
Related reading
- How to feed back `RNN` output to input in tensorflow
- How to feed into LSTM with 4 dimensional input?
- How to feed into LSTM with 4 dimensional input?
- How to find dynamically the depth of a network in Convolutional Neural Network
- How to extract classes from prefetched dataset in Tensorflow for confusion matrix
- How to extract classes from prefetched dataset in Tensorflow for confusion matrix
- How to extract data/labels back from TensorFlow dataset
- How to extract feature importances from an Sklearn pipeline
.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 courseNo course covers this one yet
Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.
Build my 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.