Attention and transformers for forecasting

1. Sequence to sequence, and the distance problem

This lecture arrives at forecasting through a detour, and the detour is the point. Sequence-to-Sequence manipulation is a key problem in NLP (Natural Language Processing), with examples the deck lists as translation (X original text, Y translated text), question answering (X question, Y answer) and text completion.

Recurrent Neural Networks model sequences — chapter 5 built them. But there is a structural cost: to compare two samples x_i and x_j we need |i - j| steps, and therefore it is very easy to lose information in the process. In language, the distance between elements is crucial, and related concepts can be very far away from each other in the sentence.

The Attention Operator solves this issue. The mechanism, in the deck's words: during a single operation, each proposed output element looks at all elements of the input sequence in parallel. And the consequence stated as a complexity:

2. The attention operator: queries, keys, values

Each element in the sequence is projected into three different sets, and the deck gives each one a question to answer:

"Projected" means matrix multiplication: Q = XW_q, K = XW_k, V = XW_v, where the weight matrices are learnable parameters, with W ∈ ℝ^d×k.

import torch, torch.nn as nn, math

class Attention(nn.Module):
    def __init__(self, d, k):
        super(Attention, self).__init__()
        self.Wq, self.Wk, self.Wv = self.init_qkv(d, k)
        self.Wo = nn.Parameter(torch.randn(k, d))

    # d = input dim, k = projection dim
    def init_qkv(self, d, k):
        Wq, Wk, Wv = [nn.Parameter(torch.randn(d, k)) for _ in range(3)]
        return Wq, Wk, Wv

    # converts a torch array X into queries, keys, values
    def get_qkv(self, X):
        return X @ self.Wq, X @ self.Wk, X @ self.Wv

The deck notes the convention that makes this efficient: here we stack all the elements of the sequence in a single matrix X, and each output Q, K, V will have in row i the projection of x_i. One matrix multiplication does the whole sequence at once — which is precisely what the recurrent formulation could not do.

The web search analogy

The deck offers an analogy that is worth keeping, because it makes the three roles concrete. Single query: "Best pizza in Cesena". Google extracts the keys from each possibly relevant page (e.g. metadata). If there is a good match, it only shows us the page content. Query, key, value — asked, advertised, delivered.

3. Scaling, softmax and extracting the values

Having queries and keys, we can compare how much a query of an element and a key of another element match. The similarity function is the dot product, q_i · k_j, or in matrix form Q · K^T, so that in each cell (i, j) sits the similarity between q_i and k_j.

Two corrections follow, each with a stated reason.

• For stability purposes, it is better to scale the dot product by √d_k. Without it, large projection dimensions produce large dot products, softmax saturates, and gradients vanish — the same saturation argument as chapter 5.
• To understand these values better, normalize for rows with a softmax: softmax(Q·K^T / √d_k). Each row then sums to one, and can be read as how the corresponding query distributes its attention.

Finally, values: what does each element contain? Let us extract only the content of similar elements:

Attention(Q, K, V) = softmax( Q · Kᵀ / √dₖ ) · V
                     ⌊___ attention weights ___⌋

4. Self-attention, cross-attention, multihead

def self_attention(self, X):
    Q, K, V = self.get_qkv(X)
    return self.forward(Q, K, V)

def forward(self, Q, K, V):
    d_k = K.shape[1]
    sim = Q @ K.T
    att_weights = torch.softmax(sim / math.sqrt(d_k), dim=-1)
    out = att_weights @ V
    out = out @ self.Wo      # project back into d dimensions
    return out

def cross_attention(self, X, Y):
    Q, _, _ = self.get_qkv(X)
    _, K, V = self.get_qkv(Y)
    return self.forward(Q, K, V)

5. Positional encoding

Attention buys parallelism at a price, and the deck states the price plainly: attention evaluates each query with each key → no knowledge on position in the sequence! The operation is symmetric under permutation of the input. For a bag of words that might be tolerable; for a time series it is fatal.

The fix: we add "positional knowledge" to the inputs.

class PositionalEncoding(nn.Module):
    def __init__(self, d, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.pe = nn.Parameter(torch.randn(max_len, d))

    def forward(self, x):
        seq_len = x.size(1)
        pe = self.pe[:seq_len, :].unsqueeze(0)
        return x + pe

Note that this implementation makes the positional encoding a learned parameter (nn.Parameter(torch.randn(...))) rather than a fixed sinusoidal table: the network discovers for itself what "position 3" should mean. It is added to the embedded input, not concatenated — the position becomes part of the same vector the attention operates on.

6. The transformer architecture

The transformer is introduced with three bullets: the first Attention-only model; an Encoder-Decoder architecture; where the Encoder manipulates the inputs and the Decoder generates the output.

class EncoderBlock(nn.Module):
    def __init__(self, d, k):
        super(EncoderBlock, self).__init__()
        self.attn  = Attention(d, k)
        self.norm1 = nn.LayerNorm(d)
        self.norm2 = nn.LayerNorm(d)
        self.ffn   = nn.Sequential(nn.Linear(d, d * 4), nn.ReLU(),
                                   nn.Linear(d * 4, d))

    def forward(self, x):
        att_out = self.attn.self_attention(x)
        x = self.norm1(x + att_out)        # residual connection + layer norm
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)
        return x

The two x + ... are residual connections: the block learns a correction to its input rather than a replacement. Together with LayerNorm, they are what allows N blocks to be stacked without the signal degrading — the structural answer to the vanishing gradient that limited the RNNs of chapter 5.

class DecoderBlock(nn.Module):
    def __init__(self, d, k):
        super(DecoderBlock, self).__init__()
        self.self_attn  = Attention(d, k)
        self.cross_attn = Attention(d, k)
        self.norm1, self.norm2, self.norm3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d)
        self.ffn = nn.Sequential(nn.Linear(d, d * 4), nn.ReLU(), nn.Linear(d * 4, d))

    def forward(self, x, enc_out):
        x = self.norm1(x + self.self_attn.self_attention(x))
        x = self.norm2(x + self.cross_attn.cross_attention(x, enc_out))
        x = self.norm3(x + self.ffn(x))
        return x

class Transformer(nn.Module):
    def __init__(self, d, k, N):
        super(Transformer, self).__init__()
        self.encoder = Encoder(d, k, N)
        self.decoder = Decoder(d, k, N)

    def forward(self, src, tgt):
        enc_out = self.encoder(src)
        return self.decoder(tgt, enc_out)

    def generate(self, src, max_len=20):
        enc_out = self.encoder(src)
        dec_input = torch.zeros(src.size(0), 1, enc_out.size(-1)).to(src.device)
        generated = []
        for _ in range(max_len):
            dec_out = self.decoder(dec_input, enc_out)
            output  = dec_out[:, -1, :].unsqueeze(1)
            generated.append(output)
            dec_input = torch.cat([dec_input, output], dim=1)
        return torch.stack(generated, dim=1)

generate is the recursive forecasting of chapters 5 and 6 in a new costume: the encoder runs once, and then each generated point is appended to dec_input and fed back. Outputs are m scalar points, generated one at a time.

7. From NLP to univariate time series

Now the translation. Inputs: n scalar points. Outputs: m scalar points, generated one at a time. But a transformer operates on vectors of dimension d, not on scalars, so two conversions are needed:

into the model:   z  ← x · wᵢₙ + bᵢₙ
out of the model:  x̂ ← z · wₒᵤₜ + bₒᵤₜ

class Embedder(nn.Module):
    def __init__(self, d):
        super(Embedder, self).__init__()
        self.w_in  = nn.Parameter(torch.randn(1, d))
        self.b_in  = nn.Parameter(torch.zeros(d))
        # initialize so that (x * w_in) @ w_out = x
        self.w_out = nn.Parameter(self.w_in / (self.w_in.norm() + 1e-8))
        self.b_out = nn.Parameter(torch.zeros(1))

    def embed(self, x):
        return x * self.w_in + self.b_in

    def unembed(self, x):
        return x @ self.w_out.T + self.b_out

The comment on w_out is the elegant part: the un-embedding is initialised as the normalised embedding, so that at step zero the round trip scalar → vector → scalar is the identity. The network starts from "do nothing" and learns a departure from it, rather than starting from noise.

The output scaler

A second module conditions the output on the context, producing a scale and a bias:

class OutputScaler(nn.Module):
  def __init__(self, embed_dim):
    super().__init__()
    self.head  = nn.Sequential(nn.Linear(embed_dim, embed_dim * 4), nn.ReLU(),
                               nn.Linear(embed_dim * 4, embed_dim))
    self.scale = nn.Linear(embed_dim, embed_dim)
    self.bias  = nn.Linear(embed_dim, embed_dim)

  def forward(self, X, context):
    head_in = X[:, -1, :]
    out = self.head(head_in)
    s = torch.sigmoid(self.scale(context))     # sigmoid keeps the scale in (0,1)
    b = self.bias(context)
    return out * s + b                          # x ← out · s + b

class TimeSeriesTransformer(nn.Module):
    def __init__(self, d, k, N):
        super(TimeSeriesTransformer, self).__init__()
        self.transformer = Transformer(d, k, N)
        self.embedder, self.scaler = Embedder(d), OutputScaler(d)

    def forward(self, src, tgt):
        src_emb, tgt_emb = self.embedder.embed(src), self.embedder.embed(tgt)
        transformer_out = self.transformer(src_emb, tgt_emb)
        scaled_out = self.scaler(transformer_out, src_emb.mean(dim=1))
        return self.embedder.unembed(scaled_out)

Note what plays the role of context: src_emb.mean(dim=1), the average embedding of the input window. The scaler therefore adapts the output level to the level of the window it just saw — a learned defence against the normalization problem of the next section.

8. Normalization and the data pollution trap

The points are normalised to [0,1] with the min-max feature scaling of chapter 3:

xₙ ← (xᵢ - xᵐᵢₙ) / (xᵐₐₓ - xᵐᵢₙ)

And then the deck states the rule and immediately flags the difficulty, in a slide that ends with a deliberate ellipsis: to avoid data pollution, normalization should only employ points from the train series. However… The next slide answers it with three words on a chart: values unseen during training!

def normalize_series(train_series, test_series=None):
    s_max = np.max(train_series)
    s_min = np.min(train_series)
    train_series = (train_series - s_min) / (s_max - s_min)
    if test_series is not None:
        test_series = (test_series - s_min) / (s_max - s_min)   # TRAIN stats
        return train_series, test_series
    return train_series

The code is correct — it applies the training min and max to the test series — and that correctness is exactly what creates the problem. The airline series rises throughout its history, so every test point lies above the training maximum. Normalised with training statistics, the test values exceed 1. The model has never seen an input greater than 1 in its life, and is now asked to forecast in that region.

9. Training, and what the loss curves confess

The dataset is the airline series once more, this time loaded through seaborn:

def get_airline_passenger_data() -> np.ndarray:
    flights = sns.load_dataset('flights')
    series = flights.pivot(index='year', columns='month', values='passengers')
    return series.values.flatten()

def create_dataset(series, input_len, output_len):
    X, Y = [], []
    for i in range(len(series) - input_len - output_len + 1):
        X.append(series[i:i+input_len])
        Y.append(series[i+input_len:i+input_len+output_len])
    return torch.utils.data.TensorDataset(
        torch.tensor(X, dtype=torch.float32), torch.tensor(Y, dtype=torch.float32))

def get_datasets(in_len=12, out_len=6):
    series = get_airline_passenger_data()
    train_series, test_series = train_test_split(series)          # 0.8
    train_series, test_series = normalize_series(train_series, test_series)
    return create_dataset(train_series, in_len, out_len), \
           create_dataset(test_series,  in_len, out_len)

This is the sliding window of chapter 5 with one change: the target is a block of output_len points, not a single one. Input 12, output 6 — the model is trained to produce half a year at a time.

def train_model(model, train_ds, test_ds, epochs=400, lr=5e-4, batch_size=8):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = torch.nn.MSELoss()
    ...
    for epoch in range(epochs):
        avg_loss = train_loop(model, train_ds, optimizer, criterion)
        train_loss.append(avg_loss)
        avg_loss = test_loop(model, test_ds, criterion)
        test_loss.append(avg_loss)
        if (epoch + 1) % 100 == 0:
            ...
            lr /= 2                      # learning rate halved every 100 epochs

if __name__ == "__main__":
    train_ds, test_ds = get_datasets()
    model = TimeSeriesTransformer(d=6, k=12, N=1)
    train_model(model, train_ds, test_ds)

The model is deliberately tiny: d = 6, k = 12, N = 1 — a six-dimensional embedding, a twelve-dimensional projection, one encoder and one decoder block. And the printed log is the most instructive output of the entire lecture:

Forecasting the whole test span

Because the model emits out_len points at a time, a long forecast is assembled in chunks, with the context window rolled forward each time:

def forecast_sequence(model, x_seq, out_len=None) -> np.ndarray:
    model.eval()
    context = x_seq.detach().clone()
    preds, remaining = [], out_len
    with torch.no_grad():
        while remaining > 0:
            y_chunk = model.generate(context.unsqueeze(0).unsqueeze(-1),
                                     max_len=remaining).squeeze(0)
            preds.append(y_chunk)
            context = torch.cat([context, y_chunk], dim=0)[-x_seq.size(0):]
            remaining -= remaining
    return torch.cat(preds, dim=0).detach().cpu().numpy()

The line context = torch.cat([context, y_chunk], dim=0)[-x_seq.size(0):] is recursive forecasting once more: append the predictions, keep only the last input_len values. Same accumulation of error, same widening uncertainty, in a fourth architecture.

10. Chronos, out of the box

The lecture closes with an Extra: Chronos by Amazon (arXiv 2403.07815), a pretrained forecasting model usable without training anything yourself.

> pip install chronos-forecasting

from chronos import ChronosPipeline

def get_model(device_map=None):
    if device_map is None:
        device_map = "cuda" if torch.cuda.is_available() else "cpu"
    return ChronosPipeline.from_pretrained("amazon/chronos-t5-tiny",
                                           device_map=device_map)

def predict(pipeline, input_series, max_len=6, num_samples=20):
    context = torch.tensor(input_series, dtype=torch.float32).unsqueeze(0)
    forecast = pipeline.predict(context, prediction_length=max_len,
                                num_samples=num_samples)
    if forecast.ndim == 3:
        forecast = np.median(forecast, axis=1)      # median over samples
    return forecast[0]

Two design choices are worth reading off that snippet. Chronos is probabilistic: it draws num_samples trajectories rather than one, and the point forecast is their median. That is the chapter 2 attribute — forecasts are not just numbers, they carry a dispersion — realised directly in the API, and the sample cloud is the prediction interval of chapter 4. And it needs no fitting at all: the pipeline is loaded pretrained and handed a context window.

Check your understanding