Introduction

Adding attention to a Bi-LSTM usually means asking the model to learn which time steps in the sequence matter most before making its final prediction. The core requirement is simple: the Bi-LSTM must return a sequence of hidden states, and the attention block must turn that sequence into a weighted context vector.

What Changes When You Add Attention

A plain Bi-LSTM for classification often looks like this:

• embedding or feature input
• 'Bidirectional(LSTM(...))'
• dense output layer

Without attention, many models use only the final hidden representation. With attention, the model keeps the full sequence of hidden states and learns weights over time steps.

That means the Bi-LSTM must use return_sequences=True.

A Small Attention Layer in Keras

One practical way to add attention is to define a custom attention layer that scores each time step, normalizes the scores with softmax, and returns the weighted sum.

1import tensorflow as tf
2from tensorflow.keras import layers
3
4class AttentionLayer(layers.Layer):
5    def __init__(self, **kwargs):
6        super().__init__(**kwargs)
7
8    def build(self, input_shape):
9        self.W = self.add_weight(
10            shape=(input_shape[-1], 1),
11            initializer="glorot_uniform",
12            trainable=True,
13            name="attention_weight",
14        )
15        self.b = self.add_weight(
16            shape=(input_shape[1], 1),
17            initializer="zeros",
18            trainable=True,
19            name="attention_bias",
20        )
21        super().build(input_shape)
22
23    def call(self, inputs):
24        score = tf.tanh(tf.matmul(inputs, self.W) + self.b)
25        weights = tf.nn.softmax(score, axis=1)
26        context = tf.reduce_sum(inputs * weights, axis=1)
27        return context

This layer expects input shaped like batch x timesteps x features and produces a context vector shaped like batch x features.

Building a Bi-LSTM with Attention

Here is a runnable text-classification style model using that attention layer:

1import tensorflow as tf
2from tensorflow.keras import Sequential, layers
3
4max_tokens = 5000
5embedding_dim = 64
6sequence_length = 100
7
8model = Sequential([
9    layers.Embedding(input_dim=max_tokens, output_dim=embedding_dim, input_length=sequence_length),
10    layers.Bidirectional(layers.LSTM(64, return_sequences=True)),
11    AttentionLayer(),
12    layers.Dense(32, activation="relu"),
13    layers.Dense(1, activation="sigmoid"),
14])
15
16model.compile(
17    optimizer="adam",
18    loss="binary_crossentropy",
19    metrics=["accuracy"],
20)
21
22model.summary()

The crucial part is return_sequences=True. If you forget that, the LSTM returns only one vector and there is no time-axis information for attention to weight.

A Minimal Training Example

To make the structure concrete, here is a small synthetic training run:

1import numpy as np
2
3X = np.random.randint(0, max_tokens, size=(256, sequence_length))
4y = np.random.randint(0, 2, size=(256, 1))
5
6model.fit(X, y, epochs=2, batch_size=32, validation_split=0.2)

This is not a meaningful dataset, but it is enough to verify that the model compiles and trains end to end.

How the Attention Weights Work

The attention layer computes a score for each Bi-LSTM time step. Those scores are normalized so they sum to 1 across the sequence. The output context vector is then the weighted combination of all hidden states.

The effect is:

• highly relevant time steps get larger weights
• less relevant time steps still contribute, but less
• the downstream dense layers receive a summary focused on important positions

This often works better than relying only on the final hidden state, especially when important evidence may appear anywhere in the sequence.

Functional API Version

If you want more control, the Functional API is usually better than Sequential, especially when you later want to inspect attention outputs or branch the model.

1import tensorflow as tf
2from tensorflow.keras import layers, Model
3
4inputs = layers.Input(shape=(sequence_length,))
5x = layers.Embedding(max_tokens, embedding_dim)(inputs)
6x = layers.Bidirectional(layers.LSTM(64, return_sequences=True))(x)
7x = AttentionLayer()(x)
8x = layers.Dense(32, activation="relu")(x)
9outputs = layers.Dense(3, activation="softmax")(x)
10
11model = Model(inputs, outputs)
12model.compile(
13    optimizer="adam",
14    loss="sparse_categorical_crossentropy",
15    metrics=["accuracy"],
16)

This version is often easier to extend to multi-class classification or to expose intermediate tensors for debugging.

Built-In Attention Layers vs Custom Attention

Keras also provides built-in attention-related layers, but a custom layer is often easier to understand when you are starting out. It makes the sequence scoring and weighted sum explicit.

Once the concept is clear, you can decide whether a built-in attention mechanism or a more advanced architecture such as self-attention is a better fit for the task.

Common Pitfalls

The most common mistake is forgetting return_sequences=True on the Bi-LSTM, which removes the time dimension that attention needs. Another is applying attention directly to token ids or embeddings before the recurrent layer when the intent was to weigh contextual hidden states. Developers also sometimes mix binary and multi-class output configurations, such as pairing a sigmoid output with categorical loss. A final issue is expecting attention to solve a weak data pipeline automatically; it can improve representation learning, but it does not replace correct preprocessing, padding, and label setup.

Summary

• To add attention to a Bi-LSTM, keep the full sequence output from the recurrent layer.
• Attention learns weights over time steps and produces a context vector.
• 'return_sequences=True is required for the Bi-LSTM.'
• A custom attention layer is a clear way to understand the mechanism in Keras.
• Match the final output layer and loss function to the real prediction task.