Attention Mechanism Impact

1class AMNLWithoutAttention(nn.Module):
2    """
3    AMNL model with attention mechanism removed (for ablation).
4
5    Uses only the BiLSTM output without self-attention.
6    """
7
8    def __init__(
9        self,
10        input_size: int = 17,
11        sequence_length: int = 30,
12        hidden_size: int = 256,
13        num_health_states: int = 3,
14        dropout: float = 0.2
15    ):
16        super().__init__()
17
18        # CNN Feature Extractor
19        self.cnn = nn.Sequential(
20            nn.Conv1d(input_size, 64, kernel_size=3, padding=1),
21            nn.BatchNorm1d(64),
22            nn.ReLU(),
23            nn.Conv1d(64, 128, kernel_size=3, padding=1),
24            nn.BatchNorm1d(128),
25            nn.ReLU(),
26        )
27
28        # BiLSTM Temporal Encoder
29        self.lstm = nn.LSTM(
30            input_size=128,
31            hidden_size=hidden_size,
32            num_layers=2,
33            batch_first=True,
34            bidirectional=True,
35            dropout=dropout
36        )
37
38        # NOTE: No attention mechanism!
39        # self.attention = nn.MultiheadAttention(...)  # REMOVED
40
41        # Simple pooling instead of attention
42        # Option 1: Last timestep (used in ablation)
43        # Option 2: Mean pooling
44        # Option 3: Max pooling
45        self.pooling_type = 'last'
46
47        # Prediction heads
48        lstm_output_size = hidden_size * 2  # Bidirectional
49
50        self.rul_head = nn.Sequential(
51            nn.Linear(lstm_output_size, 128),
52            nn.ReLU(),
53            nn.Dropout(dropout),
54            nn.Linear(128, 1)
55        )
56
57        self.health_head = nn.Sequential(
58            nn.Linear(lstm_output_size, 64),
59            nn.ReLU(),
60            nn.Dropout(dropout),
61            nn.Linear(64, num_health_states)
62        )
63
64    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
65        # CNN: [batch, seq, features] -> [batch, features, seq]
66        x = x.transpose(1, 2)
67        x = self.cnn(x)
68        x = x.transpose(1, 2)  # Back to [batch, seq, features]
69
70        # BiLSTM
71        lstm_out, _ = self.lstm(x)
72
73        # Simple pooling (no attention)
74        if self.pooling_type == 'last':
75            pooled = lstm_out[:, -1, :]
76        elif self.pooling_type == 'mean':
77            pooled = lstm_out.mean(dim=1)
78        elif self.pooling_type == 'max':
79            pooled = lstm_out.max(dim=1)[0]
80
81        # Predictions
82        rul_pred = self.rul_head(pooled)
83        health_pred = self.health_head(pooled)
84
85        return rul_pred, health_pred

1class AMNLWithAttention(nn.Module):
2    """
3    Full AMNL model with multi-head self-attention.
4    """
5
6    def __init__(self, ...):
7        # ... same as above ...
8
9        # Multi-Head Self-Attention
10        self.attention = nn.MultiheadAttention(
11            embed_dim=hidden_size * 2,
12            num_heads=8,
13            dropout=dropout,
14            batch_first=True
15        )
16
17        # Layer norm for residual connection
18        self.attention_norm = nn.LayerNorm(hidden_size * 2)
19
20    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
21        # ... CNN and BiLSTM same as above ...
22
23        # Multi-Head Self-Attention
24        attn_out, attn_weights = self.attention(
25            query=lstm_out,
26            key=lstm_out,
27            value=lstm_out
28        )
29
30        # Residual connection + layer norm
31        lstm_out = self.attention_norm(lstm_out + attn_out)
32
33        # Weighted sum using attention to last timestep
34        # Or: pooled = lstm_out[:, -1, :]
35        pooled = lstm_out[:, -1, :]
36
37        # ... predictions same as above ...
38
39        return rul_pred, health_pred

1def visualize_attention_patterns(
2    model: AMNLWithAttention,
3    test_loader: DataLoader,
4    device: torch.device,
5    num_samples: int = 10
6) -> Dict:
7    """
8    Extract and visualize attention patterns.
9
10    Returns attention weights for analysis.
11    """
12    model.eval()
13    attention_data = []
14
15    with torch.no_grad():
16        for batch_x, batch_y in test_loader:
17            batch_x = batch_x.to(device)
18
19            # Get attention weights
20            # Need to modify forward pass to return weights
21            rul_pred, health_pred, attn_weights = model.forward_with_attention(
22                batch_x
23            )
24
25            for i in range(min(len(batch_x), num_samples)):
26                attention_data.append({
27                    'rul_true': batch_y[i].item(),
28                    'rul_pred': rul_pred[i].item(),
29                    'attention_weights': attn_weights[i].cpu().numpy(),
30                    'health_state': get_health_state(batch_y[i].item())
31                })
32
33            if len(attention_data) >= num_samples:
34                break
35
36    # Analyze patterns
37    healthy_attn = np.stack([
38        d['attention_weights']
39        for d in attention_data
40        if d['health_state'] == 'healthy'
41    ])
42    critical_attn = np.stack([
43        d['attention_weights']
44        for d in attention_data
45        if d['health_state'] == 'critical'
46    ])
47
48    print("Healthy samples: attention entropy =", entropy(healthy_attn.mean(0)))
49    print("Critical samples: attention entropy =", entropy(critical_attn.mean(0)))
50
51    return {
52        'healthy_mean': healthy_attn.mean(0),
53        'critical_mean': critical_attn.mean(0),
54        'all_data': attention_data
55    }

1def run_attention_ablation(
2    datasets: List[str] = ['FD002', 'FD004'],
3    seeds: List[int] = [42, 123, 456],
4    epochs: int = 300
5) -> pd.DataFrame:
6    """
7    Compare models with and without attention.
8    """
9    results = []
10
11    for dataset in datasets:
12        for seed in seeds:
13            # With attention
14            model_attn = AMNLWithAttention(use_attention=True)
15            result_attn = train_model(model_attn, dataset, seed, epochs)
16
17            results.append({
18                'dataset': dataset,
19                'seed': seed,
20                'attention': True,
21                'rmse': result_attn['rmse'],
22                'nasa_score': result_attn['nasa_score'],
23                'r2': result_attn['r2']
24            })
25
26            # Without attention
27            model_no_attn = AMNLWithoutAttention()
28            result_no_attn = train_model(model_no_attn, dataset, seed, epochs)
29
30            results.append({
31                'dataset': dataset,
32                'seed': seed,
33                'attention': False,
34                'rmse': result_no_attn['rmse'],
35                'nasa_score': result_no_attn['nasa_score'],
36                'r2': result_no_attn['r2']
37            })
38
39    df = pd.DataFrame(results)
40
41    # Statistical test
42    for dataset in datasets:
43        attn_rmse = df[(df['dataset'] == dataset) & (df['attention'])]['rmse']
44        no_attn_rmse = df[(df['dataset'] == dataset) & (~df['attention'])]['rmse']
45
46        t_stat, p_val = stats.ttest_ind(attn_rmse, no_attn_rmse)
47        effect_size = (no_attn_rmse.mean() - attn_rmse.mean()) / np.sqrt(
48            (attn_rmse.var() + no_attn_rmse.var()) / 2
49        )
50
51        print(f"{dataset}: t={t_stat:.2f}, p={p_val:.4f}, d={effect_size:.2f}")
52
53    return df