Keras
LSTM
attention mechanism
neural networks
machine learning

Keras - Add attention mechanism to an LSTM model

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

To add attention to an LSTM model in Keras, the most important structural step is to keep the full sequence of LSTM outputs by setting return_sequences=True. Attention needs access to all time-step representations, not just the final hidden state, so the layer arrangement matters more than the word "attention" itself.

Why attention changes an LSTM model

A plain LSTM used for sequence classification often compresses the entire input sequence into one final vector. That works, but it forces the model to squeeze all relevant information into a single hidden state.

Attention changes that by letting the model compute a weighted view over the sequence outputs. Instead of trusting only the final time step, the model can focus more on the positions that matter for the prediction.

Keep the LSTM sequence output

This is the key requirement:

python
1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4
5inputs = keras.Input(shape=(20, 16))
6x = layers.LSTM(64, return_sequences=True)(inputs)

If return_sequences=False, the LSTM emits only one vector, which removes the time dimension that the attention layer needs.

Add Keras attention over the sequence

A simple self-attention style pattern with built-in Keras layers looks like this:

python
1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4
5inputs = keras.Input(shape=(20, 16))
6x = layers.LSTM(64, return_sequences=True)(inputs)
7attended = layers.Attention()([x, x])
8pooled = layers.GlobalAveragePooling1D()(attended)
9outputs = layers.Dense(1, activation="sigmoid")(pooled)
10
11model = keras.Model(inputs, outputs)
12model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
13model.summary()

Here the same sequence x is used as both query and value input, which gives a simple attention mechanism over the LSTM outputs.

Understand what the attention layer is doing

The Attention layer computes relevance scores between sequence elements and produces a weighted combination of the sequence representations. The result still has sequence structure, which is why pooling or a downstream decoder step is often added afterward.

For a classification task, common next steps include:

  • global pooling
  • selecting a summary time step
  • feeding the attended sequence into another recurrent or dense block

The right follow-up depends on whether the task is classification, tagging, or sequence generation.

Encoder-decoder attention is a different pattern

If you are building a sequence-to-sequence model, attention usually connects decoder queries to encoder outputs, not the sequence to itself.

That looks more like this conceptually:

  • encoder LSTM produces sequence outputs
  • decoder LSTM produces query states
  • attention combines decoder queries with encoder values

So when people say "add attention to an LSTM," make sure you know whether they mean self-attention over one sequence or decoder-to-encoder attention in a seq2seq model.

Masking still matters for padded sequences

If your input sequences are padded, attention should respect the mask. Keras masking support helps, but only if the earlier layers propagate it correctly. Otherwise the model may attend to padding tokens, which weakens the whole point of attention.

That means sequence preprocessing and masking are part of the attention design, not an afterthought.

Start simple before writing a custom attention layer

Many older tutorials jump straight into subclassing Layer and hand-writing alignment logic. That is useful when you need custom behavior, but it is often unnecessary for a first working model.

Use the built-in attention layers first. Write a custom layer only when you know what built-in behavior is insufficient.

Common Pitfalls

  • Forgetting return_sequences=True on the LSTM before the attention layer.
  • Expecting attention to help when the task and data do not actually benefit from long-range focus.
  • Ignoring sequence masking so the model attends to padding.
  • Using a custom attention implementation before understanding the simpler built-in Keras layers.
  • Confusing self-attention on one sequence with encoder-decoder attention in seq2seq models.

Summary

  • Attention needs access to the full LSTM sequence output, so return_sequences=True is crucial.
  • A simple Keras pattern is LSTM -> Attention -> Pooling -> Dense.
  • Attention lets the model weight important time steps instead of relying only on the final hidden state.
  • Masking matters when sequences are padded.
  • Start with built-in Keras attention layers before implementing custom attention logic.

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.