Dual-Task vs Single-Task: +304% Degradation

1class SingleTaskRULModel(nn.Module):
2    """
3    Single-task RUL prediction model (for ablation).
4
5    Same architecture as AMNL but without health classification head.
6    """
7
8    def __init__(
9        self,
10        input_size: int = 17,
11        sequence_length: int = 30,
12        hidden_size: int = 256,
13        dropout: float = 0.2,
14        use_attention: bool = True
15    ):
16        super().__init__()
17
18        # CNN Feature Extractor (same as AMNL)
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 (same as AMNL)
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        # Multi-Head Attention (same as AMNL)
39        self.use_attention = use_attention
40        if use_attention:
41            self.attention = nn.MultiheadAttention(
42                embed_dim=hidden_size * 2,
43                num_heads=8,
44                dropout=dropout,
45                batch_first=True
46            )
47
48        # RUL Prediction Head ONLY (no health head)
49        self.rul_head = nn.Sequential(
50            nn.Linear(hidden_size * 2, 128),
51            nn.ReLU(),
52            nn.Dropout(dropout),
53            nn.Linear(128, 1)
54        )
55
56        # NOTE: No health_head - this is the ablation!
57
58    def forward(self, x: torch.Tensor) -> torch.Tensor:
59        # CNN: [batch, seq, features] -> [batch, features, seq]
60        x = x.transpose(1, 2)
61        x = self.cnn(x)
62        x = x.transpose(1, 2)  # Back to [batch, seq, features]
63
64        # BiLSTM
65        lstm_out, _ = self.lstm(x)
66
67        # Attention
68        if self.use_attention:
69            attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out)
70            lstm_out = lstm_out + attn_out  # Residual
71
72        # Pool and predict RUL
73        pooled = lstm_out[:, -1, :]  # Take last timestep
74        rul_pred = self.rul_head(pooled)
75
76        # Return ONLY RUL prediction (no health logits)
77        return rul_pred

1def train_single_task(
2    model: SingleTaskRULModel,
3    train_loader: DataLoader,
4    optimizer: torch.optim.Optimizer,
5    device: torch.device,
6    use_weighted_mse: bool = True
7) -> float:
8    """
9    Train single-task model for one epoch.
10
11    Note: No health labels, no health loss, no dual-task balancing.
12    """
13    model.train()
14    total_loss = 0.0
15
16    for batch_x, batch_rul in train_loader:
17        batch_x = batch_x.to(device)
18        batch_rul = batch_rul.to(device)
19
20        optimizer.zero_grad()
21
22        # Forward pass - single output
23        rul_pred = model(batch_x)
24
25        # Loss - RUL only (no health component)
26        if use_weighted_mse:
27            weights = 1.0 + (125.0 - batch_rul.clamp(0, 125)) / 125.0
28            loss = (weights * (rul_pred.squeeze() - batch_rul) ** 2).mean()
29        else:
30            loss = F.mse_loss(rul_pred.squeeze(), batch_rul)
31
32        # Backward pass
33        loss.backward()
34        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
35        optimizer.step()
36
37        total_loss += loss.item()
38
39    return total_loss / len(train_loader)

1def run_dual_task_ablation(
2    datasets: List[str] = ['FD001', 'FD002', 'FD003', 'FD004'],
3    seeds: List[int] = [42, 123, 456],
4    epochs: int = 300
5) -> Dict:
6    """
7    Compare dual-task AMNL vs single-task RUL.
8
9    Returns results for statistical comparison.
10    """
11    results = {
12        'dual_task': [],
13        'single_task': []
14    }
15
16    for dataset in datasets:
17        for seed in seeds:
18            print(f"Dataset: {dataset}, Seed: {seed}")
19
20            # Train dual-task AMNL
21            amnl_model = DualTaskAMNL(...)
22            amnl_result = train_amnl(amnl_model, dataset, seed, epochs)
23            results['dual_task'].append({
24                'dataset': dataset,
25                'seed': seed,
26                'rmse': amnl_result['rmse'],
27                'nasa_score': amnl_result['nasa_score']
28            })
29
30            # Train single-task (ablation)
31            single_model = SingleTaskRULModel(...)
32            single_result = train_single_task_model(
33                single_model, dataset, seed, epochs
34            )
35            results['single_task'].append({
36                'dataset': dataset,
37                'seed': seed,
38                'rmse': single_result['rmse'],
39                'nasa_score': single_result['nasa_score']
40            })
41
42    # Calculate degradation
43    for ds in datasets:
44        dual_rmse = np.mean([
45            r['rmse'] for r in results['dual_task']
46            if r['dataset'] == ds
47        ])
48        single_rmse = np.mean([
49            r['rmse'] for r in results['single_task']
50            if r['dataset'] == ds
51        ])
52        degradation = (single_rmse - dual_rmse) / dual_rmse * 100
53        print(f"{ds}: Dual={dual_rmse:.2f}, Single={single_rmse:.2f}, "
54              f"Degradation={degradation:+.1f}%")
55
56    return results