Health Classification Head (3 Classes)

Learning Objectives

By the end of this section, you will:

1. Define the three health states and their RUL boundaries
2. Design the classification head architecture
3. Understand softmax output as probability distribution
4. Address class imbalance in health state labels
5. Connect classification to cross-entropy loss

Why This Matters: Health classification provides discrete checkpoints along the degradation trajectory. By predicting whether an engine is Healthy, Degrading, or Critical, the model learns to recognize qualitative state transitions—knowledge that transfers to improved RUL prediction.

Health State Definition

We discretize the continuous RUL into three health states based on remaining life.

State Definitions

Visual Representation

1RUL Timeline:
2  ∞ ←────────────────────────────────────────→ 0
3      │         │              │              │
4      │  Healthy │   Degrading  │   Critical  │ FAILURE
5      │  (Class 0)│   (Class 1)  │   (Class 2) │
6      │         │              │              │
7     125       125            50              0
8
9Health State Transition:
10  Healthy → Degrading → Critical → Failure
11     │           │           │
12    Long       Medium      Short
13   horizon    horizon     horizon

Why These Boundaries?

• 125 cycles: Aligns with piecewise linear RUL cap—beyond this, degradation is minimal
• 50 cycles: Approximately one standard deviation of typical failure prediction error—a reasonable "danger zone"
• Three classes: Sufficient granularity without over-complicating the auxiliary task

1def get_health_label(rul: float) -> int:
2    if rul > 125:
3        return 0  # Healthy
4    elif rul > 50:
5        return 1  # Degrading
6    else:
7        return 2  # Critical

Classification Head Architecture

The health classification head transforms the shared representation into class logits.

Architecture Overview

1Input: z ∈ ℝ²⁵⁶ (encoder output)
2            ↓
3┌─────────────────────────────────┐
4│  Linear(256, 64)                │
5│  256 × 64 + 64 = 16,448 params  │
6└─────────────────────────────────┘
7            ↓
8┌─────────────────────────────────┐
9│  ReLU                           │
10│  0 parameters                   │
11└─────────────────────────────────┘
12            ↓
13┌─────────────────────────────────┐
14│  Dropout(p=0.3)                 │
15│  0 parameters                   │
16└─────────────────────────────────┘
17            ↓
18┌─────────────────────────────────┐
19│  Linear(64, 3)                  │
20│  64 × 3 + 3 = 195 params        │
21└─────────────────────────────────┘
22            ↓
23Output: logits ∈ ℝ³ (class scores)

PyTorch Implementation

# Default configuration:
# input_dim=256  (from encoder)
# hidden_dim=64  (smaller than RUL head)
# num_classes=3  (Healthy, Degrading, Critical)
# dropout=0.3    (regularization)

# Data flow through Sequential:
# Input: z.shape = (32, 256)
#   ↓ Linear(256, 64)
# Shape: (32, 64)
#   ↓ ReLU()
# Shape: (32, 64)  # negative values → 0
#   ↓ Dropout(0.3)
# Shape: (32, 64)  # 30% zeroed during training
#   ↓ Linear(64, 3)
# Output: (32, 3)  # 3 class logits

# BEFORE: z.shape = (32, 256)
# Linear transformation: output = z @ W.T + b

nn.Linear(256, 64)
# Parameters: 256 × 64 + 64 = 16,448

# AFTER: shape = (32, 64)

# BEFORE: tensor([-0.5, 0.3, -0.1, 0.8])
nn.ReLU()
# AFTER:  tensor([ 0.0, 0.3,  0.0, 0.8])

# Negative values become 0
# Positive values unchanged

# Training mode (model.train()):
# BEFORE: tensor([0.5, 0.3, 0.7, 0.2, 0.8])
nn.Dropout(0.3)
# AFTER:  tensor([0.71, 0.0, 1.0, 0.29, 0.0])
#         ↑ scaled up      ↑ dropped

# Eval mode (model.eval()): No change

# BEFORE: shape = (32, 64)
nn.Linear(64, 3)
# Parameters: 64 × 3 + 3 = 195

# AFTER: shape = (32, 3)
# logits = tensor([[2.1, 0.8, -0.5],  # Sample 1
#                  [0.2, 1.5,  0.3],  # Sample 2
#                  ...])
# Column 0: Healthy score
# Column 1: Degrading score
# Column 2: Critical score

# Usage:
z = encoder(x)        # shape: (32, 256)
logits = health_head(z)  # shape: (32, 3)

# Example output:
# logits[0] = [2.1, 0.8, -0.5]
# Highest score (2.1) = Healthy
# Prediction: Class 0 (Healthy)

1class HealthHead(nn.Module):
2    """
3    Health Classification Head: 256 → 64 → 3
4
5    Two-layer MLP that predicts health state class.
6    Output is raw logits (softmax applied in loss function).
7    """
8
9    def __init__(
10        self,
11        input_dim: int = 256,
12        hidden_dim: int = 64,
13        num_classes: int = 3,
14        dropout: float = 0.3
15    ):
16        super().__init__()
17
18        self.head = nn.Sequential(
19            nn.Linear(input_dim, hidden_dim),    # 256 → 64
20            nn.ReLU(),
21            nn.Dropout(dropout),
22            nn.Linear(hidden_dim, num_classes)   # 64 → 3
23        )
24
25    def forward(self, z: torch.Tensor) -> torch.Tensor:
26        """
27        Forward pass.
28
29        Args:
30            z: Encoder output (batch, 256)
31
32        Returns:
33            Class logits (batch, 3)
34        """
35        return self.head(z)

Design Rationale

Softmax and Probability Output

The head outputs raw logits; softmax converts these to probabilities.

Logits to Probabilities

Where:

• $z_i$: Logit (raw score) for class i
• $p_i$: Probability of class i
• $\sum_i p_i = 1$: Probabilities sum to 1

Example Computation

Why Not Apply Softmax in the Head?

PyTorch's CrossEntropyLoss expects raw logits, not probabilities:

• Numerical stability: CrossEntropyLoss uses log-sum-exp trick internally
• Efficiency: Avoids computing softmax twice (once in head, once in loss)
• Convention: Standard practice in PyTorch classification

# z.shape = (32, 256)  from encoder
logits = health_head(z)
# logits.shape = (32, 3)

# Example:
# logits[0] = tensor([2.1, 0.8, -0.5])
#             Healthy  Degrading  Critical

# logits = tensor([[2.1, 0.8, -0.5],
#                   [0.2, 1.5,  0.3]])
# labels = tensor([0, 1])  # True classes

loss = F.cross_entropy(logits, labels)
# loss ≈ 0.45  (scalar)

# Internally computes:
# 1. softmax(logits) → probabilities
# 2. -log(prob of correct class)
# 3. average over batch

# logits[0] = tensor([2.1, 0.8, -0.5])

probs = F.softmax(logits, dim=-1)
# probs[0] = tensor([0.74, 0.20, 0.06])
#                    ↑sum = 1.0

# Interpretation:
# 74% confident → Healthy
# 20% confident → Degrading
#  6% confident → Critical

# probs = tensor([[0.74, 0.20, 0.06],
#                  [0.15, 0.70, 0.15]])

predicted_class = probs.argmax(dim=-1)
# predicted_class = tensor([0, 1])

# Sample 0: max at index 0 → Healthy
# Sample 1: max at index 1 → Degrading

1# Training: use raw logits
2logits = health_head(z)           # (batch, 3)
3loss = F.cross_entropy(logits, labels)
4
5# Inference: apply softmax for probabilities
6probs = F.softmax(logits, dim=-1) # (batch, 3)
7predicted_class = probs.argmax(dim=-1)  # (batch,)

Handling Class Imbalance

The three health classes are not equally represented in the data.

Class Distribution

In a typical C-MAPSS training set, the distribution is approximately:

Imbalance Mitigation Strategies

We use focal loss (Chapter 11) to address class imbalance:

Where:

• $\alpha_t$: Class weight for class t
• $(1 - p_t)^\gamma$: Focusing factor (down-weights easy examples)
• $\gamma = 2$: Typical focusing parameter

Class Weights

Alternatively, class weights can balance the loss:

1# Inverse frequency weighting
2class_counts = [4500, 3000, 2500]  # Example counts
3total = sum(class_counts)
4class_weights = torch.tensor([total / c for c in class_counts])
5class_weights = class_weights / class_weights.sum() * 3  # Normalize
6
7# Use in loss
8loss = F.cross_entropy(logits, labels, weight=class_weights)

Summary

In this section, we designed the health classification head:

1. Three health states: Healthy (RUL > 125), Degrading (50-125), Critical (≤50)
2. Architecture: Two-layer MLP (256 → 64 → 3)
3. Output: Raw logits (softmax in loss function)
4. Class imbalance: Addressed via focal loss or class weights
5. Parameters: ~17K

Looking Ahead: We have designed both prediction heads. The next section assembles the complete model—connecting encoder, heads, and showing the full forward pass from raw sensors to dual predictions.

With both heads designed, we now assemble the complete AMNL model.