LSTM
Keras
neural networks
deep learning
machine learning

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.

Practice ML system design

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=True returns the hidden output for every timestep'
  • 'return_state=True returns only the final hidden state and final cell state'

A quick example makes the distinction clear:

python
1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(4, 3))
4outputs, final_h, final_c = tf.keras.layers.LSTM(
5    units=5,
6    return_sequences=True,
7    return_state=True,
8)(inputs)
9
10model = tf.keras.Model(inputs, [outputs, final_h, final_c])
11
12x = tf.random.normal((2, 4, 3))
13out_seq, h_last, c_last = model(x)
14print(out_seq.shape)
15print(h_last.shape)
16print(c_last.shape)

This prints shapes like these:

text
(2, 4, 5)
(2, 5)
(2, 5)

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.

python
1import tensorflow as tf
2
3class CellStateSequence(tf.keras.layers.Layer):
4    def __init__(self, units):
5        super().__init__()
6        self.cell = tf.keras.layers.LSTMCell(units)
7
8    def call(self, inputs):
9        batch_size = tf.shape(inputs)[0]
10        h = tf.zeros((batch_size, self.cell.units), dtype=inputs.dtype)
11        c = tf.zeros((batch_size, self.cell.units), dtype=inputs.dtype)
12
13        cell_states = []
14        for x_t in tf.unstack(inputs, axis=1):
15            _, [h, c] = self.cell(x_t, states=[h, c])
16            cell_states.append(c)
17
18        return tf.stack(cell_states, axis=1)
19
20x = tf.random.normal((2, 4, 3))
21layer = CellStateSequence(units=5)
22c_sequence = layer(x)
23print(c_sequence.shape)

Output:

text
(2, 4, 5)

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.

python
1inputs = tf.keras.Input(shape=(4, 3))
2cell_sequence = CellStateSequence(units=5)(inputs)
3inspect_model = tf.keras.Model(inputs, cell_sequence)
4
5sample = tf.random.normal((1, 4, 3))
6print(inspect_model(sample).shape)

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 LSTM does not directly return the full cell state sequence
  • 'return_sequences=True gives you the hidden outputs, not the cell memory'
  • 'return_state=True gives only the final hidden and final cell states'
  • Use LSTMCell and collect c_t manually when you need every timestep
  • Wrap that logic in a custom layer if you want it to fit cleanly into a Keras model

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
Basis Lab
No 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 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