Self-Attention

The Library Analogy

Walk into a library with a question you cannot quite phrase — you know it has something to do with nineteenth-century French opera but you do not remember more. The librarian listens, walks the shelves, and pulls a stack of books that overlap with your fuzzy query. From those books she produces a single answer that is a weighted blend of the most relevant pages.

That is precisely what self-attention does to a sequence. Each timestep produces a query; the system compares it against every timestep's key; the matches determine how much each timestep's value should contribute to the answer. For RUL prediction this lets cycle 27 attend to cycle 5 if cycle 5 carried the early-warning signal — something neither convolution (only sees K cycles) nor a vanilla LSTM (forgets gradually) can reliably do.

Queries, Keys, Values

Given an input sequence $\mathbf{X} \in \mathbb{R}^{T \times d}$ (T cycles, d-dim features), three projection matrices produce three new sequences:

$\mathbf{Q} = \mathbf{X} W^Q, \quad \mathbf{K} = \mathbf{X} W^K, \quad \mathbf{V} = \mathbf{X} W^V.$

$W^Q, W^K, W^V \in \mathbb{R}^{d \times d_k}$ are learnable. After the projections, $\mathbf{Q}, \mathbf{K} \in \mathbb{R}^{T \times d_k}$ and $\mathbf{V} \in \mathbb{R}^{T \times d_v}$.

Scaled Dot-Product Attention

Compute pairwise similarity between every query and every key, then normalise:

$\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) \;=\; \text{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d_k}}\right) \mathbf{V}.$

Three things in one expression. (1) $\mathbf{Q}\mathbf{K}^{\top} \in \mathbb{R}^{T \times T}$ gives every query-key dot product. (2) Divide by $\sqrt{d_k}$ to keep variance roughly 1 regardless of $d_k$; without this, large$d_k$ pushes softmax into a near-one-hot regime where gradients vanish. (3) Softmax row-wise to produce a valid probability distribution, then multiply by $\mathbf{V}$ to compute the weighted sum.

Interactive: Pick a Query, See Where It Looks

Below is a $6 \times 6$ attention map computed on random Q, K, V matrices (matching the NumPy run further down so you can verify the numbers). Click any row to pick a query; the bar chart on the right shows where its probability mass goes. Toggle between raw scores and softmax-normalised attention.

Try sliding the scale factor down to 0.5: the softmax row collapses to one-hot — the model fixates on a single key, which kills gradient flow to the others. Push it up to 8: the row becomes near-uniform, and self-attention degenerates into uniform averaging. The choice $\text{scale} = \sqrt{d_k}$ is the sweet spot.

Interactive: The Macro Flow

Zooming out, the four-stage flow:

Multi-Head Attention

One attention head learns one kind of relationship. Multi-head attention runs $H$ attention computations in parallel with separate $W^Q_h, W^K_h, W^V_h$ for each head $h$, then concatenates the outputs and projects:

$\text{MHA}(X) \;=\; \text{Concat}(\text{head}_1, \ldots, \text{head}_H) \, W^O,$

with $\text{head}_h = \text{Attention}(X W^Q_h,\, X W^K_h,\, X W^V_h)$. Each head sees a lower-dim slice ($d_k = d_{\text{model}} / H$) and can specialise: one head might learn “late cycles attend to early cycles” while another learns “adjacent-cycle smoothing”. The paper and our backbone use 8 heads.

Python: Attention From Scratch

Twenty lines of NumPy implement scaled-dot-product attention end to end. The output values match what the interactive heatmap displays when you select query 0 with scale = sqrt(4) = 2.

1import numpy as np
2
3# ----- Scaled dot-product attention from scratch -----
4def softmax_rowwise(x: np.ndarray) -> np.ndarray:
5    shifted = x - x.max(axis=-1, keepdims=True)    # numerical stability
6    expd = np.exp(shifted)
7    return expd / expd.sum(axis=-1, keepdims=True)
8
9
10def attention(Q, K, V):
11    """Q, K: (T, d_k). V: (T, d_v). Returns (T, d_v) and (T, T) attention map."""
12    d_k = Q.shape[-1]
13    scores = Q @ K.T / np.sqrt(d_k)                # (T, T)
14    attn   = softmax_rowwise(scores)               # (T, T), rows sum to 1
15    out    = attn @ V                              # (T, d_v)
16    return out, attn
17
18
19# ----- Run on a tiny synthetic batch -----
20np.random.seed(0)
21T, d_k = 6, 4
22Q = np.random.randn(T, d_k).astype(np.float32)
23K = np.random.randn(T, d_k).astype(np.float32)
24V = np.random.randn(T, d_k).astype(np.float32)
25
26out, attn = attention(Q, K, V)
27
28print("scores[0]      :", (Q @ K.T / np.sqrt(d_k))[0].round(3).tolist())
29print("attn[0]        :", attn[0].round(4).tolist())
30print("attn[0].sum()  :", round(attn[0].sum(), 6))
31print("out[0]         :", out[0].round(3).tolist())
32
33# scores[0]      : [1.524, 2.145, -1.174, 0.797, 0.142, -0.279]
34# attn[0]        : [0.2611, 0.4863, 0.0176, 0.1263, 0.0656, 0.043]
35# attn[0].sum()  : 1.0
36# out[0]         : [-0.781, -0.694, -0.437, 0.121]

PyTorch: scaled_dot_product_attention

PyTorch 2.0+ ships a single fused kernel for attention, optionally backed by Flash-Attention on CUDA. For multi-head usage the nn.MultiheadAttention module wraps the projections and head-splitting boilerplate.

1import torch
2import torch.nn.functional as F
3
4torch.manual_seed(0)
5B, T, d_k = 1, 6, 4
6
7# Same shapes as the NumPy version, with a leading batch axis
8Q = torch.randn(B, T, d_k)
9K = torch.randn(B, T, d_k)
10V = torch.randn(B, T, d_k)
11
12# Single line - PyTorch 2.0+ ships an optimised kernel (Flash-Attention on
13# CUDA, plain math fallback on CPU)
14out = F.scaled_dot_product_attention(Q, K, V)
15
16print("Q.shape  :", tuple(Q.shape))    # (1, 6, 4)
17print("out.shape:", tuple(out.shape))  # (1, 6, 4)
18
19
20# ----- Multi-head attention via nn.MultiheadAttention -----
21import torch.nn as nn
22
23mha = nn.MultiheadAttention(
24    embed_dim=64, num_heads=8, batch_first=True   # 64 / 8 = 8 dims per head
25)
26
27# (B, T, embed_dim)
28x = torch.randn(2, 30, 64)
29y, attn_weights = mha(x, x, x, need_weights=True)
30
31print("input  x        :", tuple(x.shape))           # (2, 30, 64)
32print("output y        :", tuple(y.shape))           # (2, 30, 64)
33print("attn_weights    :", tuple(attn_weights.shape)) # (2, 30, 30)

Attention Beyond RUL

The same five-line attention computation underpins all of them. The mathematics is exact; what changes is the modality and the size of the (T, d) tensor.

The Three Pitfalls

The point. Self-attention lets every cycle see every other cycle in O(T^2) time, learn which to weight, and produce a context-rich output. Combined with Conv1D (local patterns) and BiLSTM (sequential dynamics), it gives our backbone its long-range modelling power.

Takeaway

• Self-attention is differentiable lookup. Each query selects a weighted blend of values via similarity to keys.
• The whole formula is one line. $\text{softmax}(QK^{\top}/\sqrt{d_k})\, V$ — three matrix multiplies plus a softmax.
• Scale by sqrt(d_k). Keeps softmax in its well-behaved regime regardless of the per-head dimension.
• Multi-head runs H attentions in parallel. Each head sees a lower-dim slice and learns its own type of relationship.
• PyTorch ships a fused kernel. F.scaled_dot_product_attention dispatches to Flash-Attention automatically on CUDA — never write the softmax(Q K^T)/sqrt(d) yourself in production code.