Attention Mechanism Theory

Learning Objectives

By the end of this section, you will:

1. Understand the attention bottleneck problem that motivates attention mechanisms
2. Derive the general attention formula with queries, keys, and values
3. Compare attention mechanisms: additive, dot-product, and scaled dot-product
4. Apply temporal attention to sequence processing for RUL prediction
5. Preview self-attention as the foundation of transformer architectures
6. Understand attention weights as interpretable importance scores

Why This Matters: Our AMNL model uses temporal attention after the BiLSTM layer. Attention allows the model to focus on the most informative timesteps for RUL prediction—whether that's a sudden vibration spike or a gradual temperature trend. This section derives the mathematics that makes this focused reasoning possible.

Why Attention?

Despite solving the vanishing gradient problem, LSTMs still face a fundamental limitation we call the bottleneck problem.

The Bottleneck Problem

Consider a BiLSTM processing a 30-timestep sequence. The final hidden states are:

• $\overrightarrow{\mathbf{h}}_{30}$: Forward LSTM's final state (summarizes entire sequence)
• $\overleftarrow{\mathbf{h}}_1$: Backward LSTM's final state (also summarizes entire sequence)

If we only use these final states, all 30 timesteps must be compressed into a fixed-size vector of dimension $2H = 128$. This creates a bottleneck:

Information Loss

This compression forces the network to prioritize some information over others:

• Recent timesteps may dominate due to recency bias
• Subtle but important early signals can be lost
• The network cannot easily retrieve specific past information when needed

The Attention Solution

Attention addresses this by maintaining direct access to all timesteps:

Where $\alpha_t \in [0, 1]$ are attention weights that sum to 1. The context vector $\mathbf{c}$ is a weighted average of all hidden states, not just the final one.

Attention Fundamentals

At its core, attention is a mechanism for computing a weighted average where the weights depend on the content being aggregated.

The Query-Key-Value Framework

Modern attention mechanisms are formulated using three components:

The attention mechanism computes how well each key matches the query, then uses these matches to weight the values.

General Attention Formula

Given:

• Query: $\mathbf{q} \in \mathbb{R}^{d_k}$
• Keys: $\mathbf{K} = [\mathbf{k}_1, \ldots, \mathbf{k}_T]^{\top} \in \mathbb{R}^{T \times d_k}$
• Values: $\mathbf{V} = [\mathbf{v}_1, \ldots, \mathbf{v}_T]^{\top} \in \mathbb{R}^{T \times d_v}$

Attention computes:

Where the attention weights are:

The score function measures compatibility between query and key. The softmax ensures $\sum_t \alpha_t = 1$.

Types of Attention Mechanisms

Different score functions define different attention mechanisms. The three main types are additive, dot-product, and scaled dot-product.

1. Additive Attention (Bahdanau)

Introduced by Bahdanau et al. (2014) for neural machine translation:

Components:

• $\mathbf{W}_q \in \mathbb{R}^{d_a \times d_q}$: Projects query to attention dimension
• $\mathbf{W}_k \in \mathbb{R}^{d_a \times d_k}$: Projects keys to attention dimension
• $\mathbf{v} \in \mathbb{R}^{d_a}$: Learned vector that produces scalar score
• $\tanh$: Non-linearity that captures complex query-key interactions

Advantages: More expressive due to non-linearity; works well when query and key dimensions differ.

Disadvantages: Requires additional parameters; slower due to MLP computation.

2. Dot-Product Attention (Luong)

Simpler alternative using direct dot products:

Advantages: No additional parameters; highly efficient using matrix multiplication.

Disadvantages: Requires $d_q = d_k$; can have numerical issues for large dimensions.

3. Scaled Dot-Product Attention

Used in Transformers (Vaswani et al., 2017) to address the scaling issue:

Why Scaling?

Comparison Summary

Temporal Attention for Sequences

For RUL prediction, we apply attention over the temporal dimension to identify which timesteps are most relevant for the prediction.

Setup

After the BiLSTM, we have hidden states for each timestep:

Where $T = 30$ (window size) and $2H = 128$ (BiLSTM output dimension).

Temporal Attention Formulation

We use a learned attention mechanism to compute importance weights:

Interpreting Each Component

Attention Interpretation

The attention weights $\alpha_t$ have a powerful interpretation: they tell us which timesteps the model considers most important for the prediction.

Self-Attention Preview

While our model uses the temporal attention mechanism described above, it's worth understanding self-attention—the foundation of modern Transformer architectures.

From Attention to Self-Attention

In the attention we described, the query typically comes from a different source (e.g., decoder hidden state querying encoder states). In self-attention, queries, keys, and values all come from the same sequence:

Where $\mathbf{H}$ is the sequence of hidden states. Each position can attend to every other position, including itself.

Self-Attention Formula

Key Difference: All-to-All Attention

Self-attention computes $T \times T$ attention scores, allowing every timestep to directly attend to every other timestep. This enables:

• Long-range dependencies: Timestep 1 can directly influence timestep 30 (no gradient path through intermediate states)
• Parallel computation: All attention scores computed simultaneously
• Rich representations: Each position gets a context-aware embedding

Attention in Our Architecture

Let's trace how attention integrates into the AMNL model pipeline.

Data Flow

1. CNN Output: $\mathbf{X} \in \mathbb{R}^{30 \times 64}$(30 timesteps, 64 features)
2. BiLSTM Output: $\mathbf{H} \in \mathbb{R}^{30 \times 128}$(forward + backward hidden states)
3. Attention: $\mathbf{c} \in \mathbb{R}^{128}$(weighted combination)
4. Prediction Heads: RUL regression and health state classification

Attention Dimension Design

Common choices for attention hidden dimension $d_a$:

Complete Attention Layer

The attention computation can be written as:

1# Attention layer (PyTorch-style pseudocode)
2# Input: H of shape (batch, T, 2*hidden_dim)
3# Output: context of shape (batch, 2*hidden_dim)
4
5# Step 1: Project to attention space
6U = tanh(self.W_a(H))  # (batch, T, attention_dim)
7
8# Step 2: Compute alignment scores
9e = self.v_a(U).squeeze(-1)  # (batch, T)
10
11# Step 3: Softmax to get weights
12alpha = softmax(e, dim=1)  # (batch, T)
13
14# Step 4: Weighted sum
15context = torch.sum(alpha.unsqueeze(-1) * H, dim=1)  # (batch, 2*hidden_dim)

Why Attention Helps RUL Prediction

For RUL prediction specifically, attention provides three key benefits:

1. Selective focus: Weights identify which parts of the degradation history matter most
2. Noise robustness: The model can downweight timesteps with anomalous sensor readings
3. Interpretability: Attention weights explain which time periods influenced the prediction

Summary

In this section, we explored attention mechanisms as a solution to the sequence modeling bottleneck:

1. Bottleneck problem: Fixed-size hidden states must compress entire sequences, losing information
2. Attention solution: Compute weighted averages over all hidden states: $\mathbf{c} = \sum_t \alpha_t \mathbf{h}_t$
3. Query-Key-Value: Framework where queries search keys to retrieve values
4. Score functions: Additive (learned), dot-product (efficient), scaled dot-product (stable)
5. Temporal attention: Applies attention over time dimension for sequence summarization
6. Self-attention: Each position attends to all positions (foundation of Transformers)

Looking Ahead: We have now covered the individual components of our model: CNN for local feature extraction, BiLSTM for temporal dependencies, and attention for focused summarization. But our AMNL model has one more innovation—it performs multi-task learning, predicting both continuous RUL and discrete health states simultaneously. In the next section, we'll derive the multi-task learning framework and understand how combining tasks improves overall performance.

With attention understood, we are ready to explore how multi-task learning enables our model to leverage complementary prediction targets.