Introduction

Putting attention on top of an LSTM or GRU is a common way to let the model focus on the most relevant time steps in a sequence. The crucial implementation detail is that the recurrent layer must return the full sequence, not just the last hidden state. Once you have sequence outputs, you can apply either Keras's built-in attention layers or a small custom scoring layer depending on the task.

Return the Full Sequence from the Recurrent Layer

If the LSTM or GRU returns only the last state, attention has nothing to attend over. That is why return_sequences=True is mandatory.

1import tensorflow as tf
2from tensorflow.keras import layers
3
4inputs = layers.Input(shape=(50, 32))
5sequence = layers.Bidirectional(
6    layers.LSTM(64, return_sequences=True)
7)(inputs)

At this point sequence has shape (batch, timesteps, features). That is exactly what an attention mechanism needs.

The same idea works with a GRU:

sequence = layers.GRU(64, return_sequences=True)(inputs)

The model choice does not change the attention requirement.

Use a Simple Custom Attention for Sequence Classification

For many classification problems, a lightweight attention layer is easier to reason about than the generic Attention layer.

1import tensorflow as tf
2from tensorflow.keras import layers, Model
3
4inputs = layers.Input(shape=(50, 32))
5x = layers.Bidirectional(layers.GRU(64, return_sequences=True))(inputs)
6
7score = layers.Dense(1, activation="tanh")(x)
8weights = layers.Softmax(axis=1)(score)
9context = tf.reduce_sum(x * weights, axis=1)
10
11outputs = layers.Dense(3, activation="softmax")(context)
12model = Model(inputs, outputs)
13model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
14model.summary()

Here the dense layer produces one score per time step, softmax turns those scores into attention weights, and the weighted sum produces a single context vector for classification.

This pattern is practical because it matches the question most people really mean: "how do I learn which time steps matter before the final classifier?"

Use Keras Attention When You Have Query-Key Structure

Keras also provides Attention and AdditiveAttention. These are more natural when you have a separate query and value sequence, as in encoder-decoder models or cross-attention-like setups.

1from tensorflow.keras import layers, Model
2
3encoder_inputs = layers.Input(shape=(50, 32))
4encoder_outputs = layers.LSTM(64, return_sequences=True)(encoder_inputs)
5
6query = layers.GlobalAveragePooling1D()(encoder_outputs)
7query = layers.Reshape((1, 64))(query)
8
9attended = layers.Attention()([query, encoder_outputs])
10attended = layers.Flatten()(attended)
11outputs = layers.Dense(1, activation="sigmoid")(attended)
12
13model = Model(encoder_inputs, outputs)
14model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

This works, but for simple sequence classification the custom score-and-weight pattern is often easier to interpret.

Masking Matters for Padded Sequences

If your sequences are padded, attention can accidentally place weight on padding tokens unless masking is wired correctly.

inputs = layers.Input(shape=(None,), dtype="int32")
x = layers.Embedding(input_dim=10000, output_dim=64, mask_zero=True)(inputs)
x = layers.GRU(64, return_sequences=True)(x)

Using mask_zero=True in the embedding layer helps downstream recurrent layers understand padding. With custom attention, you may also need to carry the mask logic explicitly if the built-in layers are not handling it for your case.

Padding bugs are subtle because the model still trains, but it may learn unstable or misleading attention patterns.

Choose Attention Based on the Output You Need

There are two common goals:

• produce one context vector for the whole sequence
• produce one attended output per decoder step

For the first case, a custom temporal attention over recurrent outputs is often enough. For the second, built-in attention layers or a fuller encoder-decoder design make more sense.

That distinction keeps the architecture honest. A lot of bad examples mix classification-style attention and seq2seq attention as if they were interchangeable.

Common Pitfalls

• Forgetting return_sequences=True on the LSTM or GRU layer.
• Applying attention to a single hidden state instead of a full sequence of states.
• Using the generic Attention layer when a simpler custom weighting layer would fit the task better.
• Ignoring padding masks and letting the model attend to padded positions.
• Assuming attention automatically improves results without checking whether the task actually benefits from it.

Summary

• Attention on top of LSTM or GRU starts with return_sequences=True.
• For sequence classification, a small custom score-and-weight layer is often the simplest solution.
• Keras Attention is more natural when you have distinct query and value tensors.
• Handle masking carefully if sequences are padded.
• Pick the attention pattern that matches the task instead of forcing one generic template onto every model.