Supervised learning: deep learning 2 / RNNs and Attention¶

In the previous practical we built feed-forward and convolutional networks for MNIST digit classification. Here we look at two more architectures: recurrent neural networks (RNNs) and attention. We use the same dataset so we can compare all four approaches directly.

RNNs read sequences one step at a time. An image isn't obviously a sequence, but we can make it one: treat each 28×28 digit as a sequence of 28 rows, each row a 28-pixel vector. A SimpleRNN or LSTM reads the image row by row, top to bottom, and builds up a hidden representation before classifying it, differently from a CNN's local, translation-invariant filters.

We'll then add attention on top, which learns how much weight to give each row's hidden state. That gives us something a CNN's filters don't hand us directly: a per-row importance score we can overlay on the digit and inspect visually.

Setup: installing TensorFlow¶

TensorFlow does not yet publish packages for very new Python versions (e.g. 3.13/3.14) — pip install tensorflow will fail with "No matching distribution found" on those. Use a dedicated conda environment pinned to Python 3.11 (any version in the 3.9–3.11 range works):

conda create -n daml python=3.11 -y
conda activate daml
pip install numpy matplotlib tensorflow
conda install jupyter -y
jupyter notebook
In [1]:
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

Let's set the seed value and use the same number as below to reproduce the same results.

In [2]:
tf.random.set_seed(45)
np.random.seed(45)

Data preparation¶

As in the previous practical, we load MNIST and normalize pixel values to (0, 1). We won't add a channel dimension this time.

In [3]:
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train / 255
x_test = x_test / 255
In [4]:
def plot_img(img, cmap="gray_r", **kwargs):
    plt.imshow(img, cmap=cmap, **kwargs)
    plt.axis("off")
    plt.show()
In [5]:
def plot_history(history, title):
    acc = history.history["accuracy"]
    val_acc = history.history["val_accuracy"]
    loss = history.history["loss"]
    val_loss = history.history["val_loss"]
    x = range(1, len(acc) + 1)

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(x, acc, "b", label="Training accuracy")
    plt.plot(x, val_acc, "r", label="Validation accuracy")
    plt.title(f"{title} — accuracy")
    plt.xlabel("epoch")
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(x, loss, "b", label="Training loss")
    plt.plot(x, val_loss, "r", label="Validation loss")
    plt.title(f"{title} — loss")
    plt.xlabel("epoch")
    plt.legend()
    plt.tight_layout()
    plt.show()

1. Each MNIST image already has shape (28, 28), no reshaping needed. Explain why this shape already works as RNN input, and how it differs from what the CNN in the previous practical required.¶

Recall a Keras RNN layer expects input of shape (timesteps, features_per_timestep).

Baseline: feed-forward network¶

A quick reference point, as in the previous practical.

2. Build and train a feed-forward network: flatten the input, a Dense(64, activation="relu") hidden layer, then a Dense(10, activation="softmax") output layer. Compile with loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"], and train for 5 epochs with validation_split=0.2. Plot the training history with plot_history().¶

Recurrent neural networks¶

3. Build a model with a SimpleRNN(64) layer taking the (28, 28) images directly as input, followed by a Dense(10, activation="softmax") output layer. Compile the same way, train for 5 epochs, and plot the training history with plot_history().¶

A plain SimpleRNN struggles to carry information over long sequences. Its hidden state gets overwritten at every step. LSTM layers add gates that let the network learn what to keep and what to forget, which usually works better, especially over 28 timesteps.

4. Rebuild the same model, but replace SimpleRNN(64) with LSTM(64). Train it the same way, plot its training history, and compare its test accuracy to the SimpleRNN and feed-forward baselines.¶

5. Row-scan vs column-scan: does the direction the RNN reads the image matter? Build the same LSTM(64) model as Q4, but feed it the transposed images (columns read top-to-bottom instead of rows) using np.transpose(x, axes=(0, 2, 1)). Train and evaluate it (plot its training history too), then compare its test accuracy to the row-scanning LSTM from Q4.¶

Attention¶

An LSTM with return_sequences=True outputs a hidden state for every row, not just the last one. Attention learns a weight for each row's hidden state, i.e., how much that row should contribute to the final prediction, and combines them into a single weighted vector.

6. Below is a small custom attention layer implementing additive (Bahdanau-style) attention. Read through it and make sure you understand what each part computes.¶

In [6]:
class AttentionLayer(layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        # input_shape: (batch, timesteps, hidden_dim)
        self.W = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="glorot_uniform",
            trainable=True,
            name="attention_W"
        )
        self.v = self.add_weight(
            shape=(self.units, 1),
            initializer="glorot_uniform",
            trainable=True,
            name="attention_v"
        )

    def call(self, hidden_states):
        # hidden_states: (batch, timesteps, hidden_dim)
        score = tf.nn.tanh(tf.matmul(hidden_states, self.W))    # (batch, timesteps, units)
        score = tf.matmul(score, self.v)                         # (batch, timesteps, 1)
        weights = tf.nn.softmax(score, axis=1)                   # (batch, timesteps, 1)
        context = tf.reduce_sum(weights * hidden_states, axis=1)   # (batch, hidden_dim)
        return context, weights

7. Build a model using LSTM(64, return_sequences=True) followed by the AttentionLayer above, then a Dense(10, activation="softmax") output layer on the resulting context vector. Since this model has two outputs (context and weights) but we only train on the classification output, use the functional API rather than Sequential. Compile and train it the same way as the other models, and plot its training history.¶

8. Extract the attention weights for a few test digits using a second model that outputs attention_weights instead of the final prediction. For 3 examples, plot the original digit next to a horizontal bar chart of its 28 row-attention weights, so the bars line up with the image rows. Which rows does the model attend to most?¶

9. Attention on a misclassified example: does attention look any different when the model gets it wrong? Find a test digit the attention model misclassifies (compare attn_model.predict() to y_test), and plot its row-attention weights next to a correctly classified example of the same true digit.¶

10. OPTIONAL (if time permits): Compare all four models (feed-forward, SimpleRNN, LSTM, and LSTM+Attention) on test accuracy in a single table.¶