Introduction

A Bi-LSTM already captures context from both directions, but it still compresses sequence information in a fixed way unless you add another mechanism on top. Attention improves that by letting the model learn which timesteps matter most for the final prediction instead of treating every hidden state equally.

The Core Idea

A Bi-LSTM produces one hidden representation per timestep when return_sequences=True. Attention then assigns a learned weight to each timestep and builds a weighted summary vector.

So the workflow is:

1. encode the sequence with a Bi-LSTM
2. compute attention scores over timesteps
3. normalize those scores into weights
4. take the weighted sum of hidden states
5. send the resulting context vector to the output layer

That is the simplest mental model for sentence classification, sequence labeling, and many other sequence tasks.

A Minimal Keras Example

Here is a compact TensorFlow and Keras model that adds a custom attention layer on top of a Bi-LSTM.

1import tensorflow as tf
2
3class AttentionPooling(tf.keras.layers.Layer):
4    def __init__(self, **kwargs):
5        super().__init__(**kwargs)
6        self.score_dense = tf.keras.layers.Dense(1)
7
8    def call(self, inputs, mask=None):
9        scores = self.score_dense(inputs)
10        scores = tf.squeeze(scores, axis=-1)
11
12        if mask is not None:
13            mask = tf.cast(mask, tf.float32)
14            scores += (1.0 - mask) * -1e9
15
16        weights = tf.nn.softmax(scores, axis=1)
17        weights = tf.expand_dims(weights, axis=-1)
18        context = tf.reduce_sum(inputs * weights, axis=1)
19        return context
20
21
22inputs = tf.keras.Input(shape=(50,), dtype="int32")
23x = tf.keras.layers.Embedding(input_dim=5000, output_dim=64, mask_zero=True)(inputs)
24x = tf.keras.layers.Bidirectional(
25    tf.keras.layers.LSTM(64, return_sequences=True)
26)(x)
27x = AttentionPooling()(x)
28outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
29
30model = tf.keras.Model(inputs, outputs)
31model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
32model.summary()

This is a good baseline for sequence classification.

Why return_sequences=True Matters

Attention needs access to the whole sequence of hidden states. If the Bi-LSTM returns only its final output, there is no sequence left to score.

That means this is correct:

tf.keras.layers.LSTM(64, return_sequences=True)

and this is not sufficient for sequence-level attention:

tf.keras.layers.LSTM(64, return_sequences=False)

The second form collapses the time dimension before attention ever gets a chance to use it.

What the Attention Layer Is Learning

The custom layer above learns a scalar score for each timestep. After softmax, those scores become weights that sum to 1 across the sequence.

That means important timesteps receive more influence in the final context vector. In text classification, those may be sentiment-bearing words. In time series, those may be spikes or regime changes.

This is why attention is often easier to interpret than using only the final hidden state.

Built-In Attention Versus Custom Pooling

Keras also has built-in attention layers such as Attention and AdditiveAttention, but those are most natural when you have separate query and value tensors, such as encoder-decoder models.

For a simple Bi-LSTM classifier, a lightweight attention-pooling layer is often easier to understand and debug.

If you move into sequence-to-sequence tasks, decoder attention becomes a better fit than this pooling-style summary.

Masking Is Not Optional

If sequences are padded, the attention mechanism must ignore padded positions. In the example above, the layer uses the incoming mask and applies a large negative value to padded logits before softmax.

Without masking, the model can waste probability mass on padding timesteps, which weakens both accuracy and interpretability.

That is why mask_zero=True on the embedding layer and mask-aware attention logic work well together.

A Small Training Example

The model compiles like any other Keras model.

1import numpy as np
2
3X = np.random.randint(1, 5000, size=(100, 50))
4y = np.random.randint(0, 2, size=(100,))
5
6model.fit(X, y, epochs=2, batch_size=16)

This is only synthetic data, but it shows that the architecture is runnable end to end.

When Attention Helps Most

Attention is especially useful when:

• not all positions in the sequence are equally informative
• the important evidence may appear anywhere in the sequence
• you want the model to produce a learned weighted summary rather than relying on the last state only

For very short sequences or simple tasks, a plain Bi-LSTM may already be enough. Attention helps most when the sequence contains distractors or dispersed signal.

Common Pitfalls

A common mistake is forgetting return_sequences=True, which removes the time dimension attention needs.

Another issue is applying attention without masking padded tokens. That often leads to weaker and less interpretable weights.

Developers also sometimes use built-in Keras attention layers without understanding whether they need query-key-value style attention or simple sequence pooling. Those are related but not identical.

Finally, do not expect attention to rescue a poor encoder automatically. If the Bi-LSTM representations are weak, the attention layer has little useful signal to weight.

Summary

• Add attention after a Bi-LSTM that returns the full sequence.
• Use return_sequences=True so attention can score each timestep.
• Mask padding positions before softmax.
• A simple attention-pooling layer is often enough for sequence classification.
• Attention helps the model focus on the most informative parts of the sequence.