Introduction

If you use the keras-self-attention package with an LSTM, the easiest way to visualize attention is to ask the attention layer to return its attention matrix and then plot that matrix for a sample sequence. The package supports this directly through the return_attention=True option on SeqSelfAttention.

That matters because an attention model without attention outputs is hard to inspect. You may get a prediction, but you cannot tell which timesteps influenced it unless you expose the attention weights explicitly.

Build the Model with return_attention=True

The SeqSelfAttention layer normally returns the transformed sequence. When return_attention=True is enabled, it returns two outputs: the transformed sequence and the attention matrix.

1import numpy as np
2import tensorflow as tf
3from tensorflow import keras
4from tensorflow.keras import layers
5from keras_self_attention import SeqSelfAttention
6
7max_len = 8
8vocab_size = 100
9
10inputs = keras.Input(shape=(max_len,), dtype="int32")
11x = layers.Embedding(input_dim=vocab_size, output_dim=16, mask_zero=True)(inputs)
12x = layers.Bidirectional(layers.LSTM(32, return_sequences=True))(x)
13att_output, att_weights = SeqSelfAttention(
14    attention_activation="sigmoid",
15    return_attention=True,
16    name="self_attention"
17)(x)
18pooled = layers.GlobalMaxPooling1D()(att_output)
19prediction = layers.Dense(1, activation="sigmoid", name="prediction")(pooled)
20
21model = keras.Model(inputs=inputs, outputs=prediction)
22inspect_model = keras.Model(inputs=inputs, outputs=[prediction, att_weights])

The important object here is inspect_model. It gives you both the model prediction and the attention matrix for the same input.

Understand the Shape of the Attention Output

With this package, self-attention weights are a matrix over sequence positions. For one example, the shape is usually time_steps x time_steps, wrapped in a batch dimension.

That means each row shows how strongly one timestep attends to all timesteps in the sequence. In a bidirectional LSTM, this is especially useful because the model can attend forward and backward across the entire sequence.

Run a Sample Through the Inspection Model

Use one padded input sequence and collect both outputs.

1sample = np.array([
2    [4, 8, 15, 16, 23, 42, 0, 0]
3], dtype="int32")
4
5pred, attention = inspect_model.predict(sample, verbose=0)
6
7print("prediction:", pred[0, 0])
8print("attention shape:", attention.shape)
9print(attention[0])

For a single sequence of length 8, attention[0] is the matrix you want to visualize.

Plot the Attention Matrix

A heatmap is the clearest first visualization.

1import matplotlib.pyplot as plt
2
3sequence_tokens = ["4", "8", "15", "16", "23", "42", "PAD", "PAD"]
4weights = attention[0]
5
6plt.figure(figsize=(6, 5))
7plt.imshow(weights, cmap="viridis")
8plt.colorbar(label="attention weight")
9plt.xticks(range(len(sequence_tokens)), sequence_tokens, rotation=45)
10plt.yticks(range(len(sequence_tokens)), sequence_tokens)
11plt.xlabel("Attended timestep")
12plt.ylabel("Query timestep")
13plt.title("Self-attention heatmap")
14plt.tight_layout()
15plt.show()

This plot helps you see whether the model is concentrating on a few positions, spreading attention broadly, or mostly attending to padding by mistake.

Keep Token Mapping Around

The heatmap only becomes interpretable if you can map sequence positions back to original tokens. If your model uses integer encoding, keep the tokenizer or vocabulary lookup available when you inspect predictions.

For example, if position 3 receives consistently high attention, you need to know whether that position corresponds to a meaningful word, a delimiter, or padding. Without that mapping, the plot is technically correct but not very useful.

Save and Load Models Correctly

If you save a model that uses SeqSelfAttention, load it with the package's custom objects so Keras can reconstruct the layer.

1from tensorflow import keras
2from keras_self_attention import SeqSelfAttention
3
4loaded = keras.models.load_model(
5    "attention_model.keras",
6    custom_objects=SeqSelfAttention.get_custom_objects()
7)

That is not part of visualization itself, but it matters if you want to inspect attention weights later from a saved model rather than from the training process.

Common Pitfalls

• Forgetting return_attention=True, which leaves you with predictions but no attention matrix to plot.
• Plotting timestep indices without mapping them back to actual tokens.
• Interpreting high attention on padded positions as meaningful model behavior.
• Building only the training model and forgetting to create an inspection model that exposes attention outputs.
• Assuming attention alone proves model correctness instead of checking predictions and errors too.

Summary

• 'SeqSelfAttention can return attention weights directly when return_attention=True is enabled.'
• Build a secondary inspection model that outputs both predictions and attention matrices.
• Plot attention[0] as a heatmap to visualize timestep-to-timestep focus.
• Keep token mappings so the visualization can be interpreted by humans.
• Validate attention patterns alongside model accuracy, not instead of it.