PyTorch Implementation

Import the main PyTorch library. This gives us access to tensor operations, neural network modules, and GPU acceleration.

import torch  # torch.tensor([1, 2, 3])

Import nn module which contains building blocks for neural networks like layers, loss functions, and containers.

nn.Linear(10, 5)  # Creates a linear layer

Import Optional for type hints. Optional[float] means the variable can be either float or None, improving code clarity.

self.ema: Optional[float] = None  # Can be float or None

Define the EMANormalizer class. This is a plain Python class (not nn.Module) because it only tracks statistics, not learnable parameters.

Initialize the normalizer with beta (smoothing factor). β=0.99 means 99% weight on history, 1% on new value. Higher β = smoother, slower adaptation.

normalizer = EMANormalizer(beta=0.99)  # ~100 step memory

Store the smoothing factor. Common values: 0.9 (~10 step window), 0.99 (~100 step window), 0.999 (~1000 step window).

self.beta = 0.99  # Effective window ≈ 1/(1-0.99) = 100 steps

Initialize EMA value to None. We use None to detect the first update and handle it specially (no history to blend with).

self.ema = None  # Will be set to first loss value

Track number of updates for bias correction. Without this, early EMA values would be biased toward zero.

self.steps = 0  # Incremented each update() call

Main method called each training step. Takes the current loss value (as a scalar float, not tensor) and returns the bias-corrected EMA for normalization.

norm = normalizer.update(loss.item())  # Returns ~1500.0

Increment step counter BEFORE computing. This is critical for correct bias correction formula (uses t, not t-1).

Step 1: steps=1, Step 2: steps=2, Step 3: steps=3

Check if this is the first update. If ema is None, we have no history to blend with.

For the first step, set EMA directly to the loss value. No blending needed since there is no prior history.

First loss=1500 → self.ema = 1500

The exponential moving average formula: μₜ = β·μₜ₋₁ + (1-β)·Lₜ. This blends 99% of history with 1% of new value.

ema=1500, loss=1600: new_ema = 0.99×1500 + 0.01×1600 = 1501

Apply bias correction: μ̂ = μ / (1 - βᵗ). Early EMA values are biased low because they start from 0. This corrects for that bias.

t=1: 1/(1-0.99¹)=100×, t=10: 1/(1-0.99¹⁰)≈10.5×, t=100: ≈1×

Return the bias-corrected EMA. This is used as the denominator when normalizing the loss: normalized_loss = loss / corrected.

loss=1500, corrected=1500 → normalized=1.0

Get the current bias-corrected EMA without updating it. Useful for logging or validation where you do not want to affect the running average.

If EMA has not been initialized yet (no updates called), return 1.0 as a safe default that will not affect loss scaling.

Same bias correction formula as in update(). Returns the current estimate without modifying state.

Reset the normalizer to initial state. Call this when starting a new training run to clear history from previous training.

normalizer.reset()  # Before new training run

Set EMA back to None so next update() will initialize fresh.

Reset step counter to 0 for correct bias correction in new training run.

Initialize with β=0.99. The effective memory window is ≈1/(1-β) = 100 steps. The EMA will be heavily influenced by the last ~100 loss values.

Typical RUL loss values might range from 1000-2000 (MSE with errors around 30-40 cycles). We simulate 5 training steps with varying losses.

Iterate through each loss value. enumerate(losses, 1) starts counting from 1 instead of 0 for human-readable step numbers.

step=1, loss=1500.0 | step=2, loss=1600.0 | ...

Capture the raw EMA value before update for comparison. On first step, ema is None so we use 0 for display.

Call update() with the loss value. This: (1) increments steps, (2) updates raw EMA, (3) returns bias-corrected value.

Step 1: corrected=1500.0, Step 2: corrected=1507.54

Divide loss by corrected EMA. Result oscillates around 1.0 regardless of absolute loss magnitude. This is the key benefit!

1600/1507.54 ≈ 1.06 (6% above average)

First step: EMA = loss = 1500. Bias correction: 1500/(1-0.99¹) = 1500/0.01×0.01 = 1500 (correction factor is 1.0 after simplification for first step).

EMA = 0.99×1500 + 0.01×1600 = 1501. Bias correction: 1501/(1-0.99²) = 1501/0.0199 = 1507.54. Loss 1600 is above average, so normalized > 1.

EMA updates toward 1400. Corrected value slowly adapts. Loss 1400 is below average → normalized = 0.93 (7% below average).

After 5 steps, EMA has stabilized around 1500. Bias correction factor is now close to 1. Future losses will be normalized relative to this running average.

Inherit from nn.Module to make this a proper PyTorch loss function. This allows it to be used seamlessly in training loops and supports GPU acceleration.

Initialize with r_max (maximum RUL). For C-MAPSS, RUL is typically capped at 125 cycles since engines are healthy above this threshold.

loss_fn = WeightedMSELoss(r_max=125.0)

Call parent nn.Module constructor. Required for proper PyTorch module initialization and parameter tracking.

Store maximum RUL as instance variable. Used to compute sample weights based on how close to failure each sample is.

Define the forward pass. This is called when you use the loss function: loss = loss_fn(pred, target).

Reshape predictions to 1D using .view(-1). The -1 means 'infer this dimension'. Handles both (batch,) and (batch, 1) shapes from different model outputs.

# Case 1: Model outputs shape (batch, 1)
pred = tensor([[100.0],
               [45.0],
               [80.0],
               [120.0]])  # shape = (4, 1)

pred.view(-1)
# Result: tensor([100.0, 45.0, 80.0, 120.0])  # shape = (4,)

# Case 2: Already flat - no change
pred = tensor([100.0, 45.0, 80.0, 120.0])  # shape = (4,)
pred.view(-1)  # Same: tensor([100.0, 45.0, 80.0, 120.0])

Same flattening for targets. Both tensors now have identical shape (batch_size,) regardless of input format, enabling element-wise operations.

target = tensor([[90.0],
                 [50.0],
                 [85.0],
                 [110.0]])  # shape = (4, 1)

target.view(-1)
# Result: tensor([90.0, 50.0, 85.0, 110.0])  # shape = (4,)

# Now both aligned for element-wise operations:
# pred   = [100.0, 45.0, 80.0, 120.0]
# target = [ 90.0, 50.0, 85.0, 110.0]
# diff   = [ 10.0, -5.0, -5.0,  10.0]

Clamp targets at r_max for weight computation. Samples with RUL > 125 all get weight = 1.0 (minimum weight).

# BEFORE: target tensor with values exceeding r_max
target = tensor([50.0, 125.0, 200.0, 80.0])

# torch.clamp(target, max=125.0)
capped_target = tensor([50.0, 125.0, 125.0, 80.0])
#                                  ↑
#                          200 → 125 (capped)

# Values ≤ 125: unchanged
# Values > 125: clamped to 125

Linear decay formula: w = 2 - y/R_max. At RUL=0: w=2.0 (max), at RUL=125: w=1.0 (min). Critical samples weighted 2× more.

# capped_target = [50.0, 125.0, 125.0, 80.0]
# r_max = 125.0

# Step-by-step: weights = 2.0 - capped_target / r_max
#   2.0 - 50/125  = 2.0 - 0.4 = 1.6
#   2.0 - 125/125 = 2.0 - 1.0 = 1.0
#   2.0 - 125/125 = 2.0 - 1.0 = 1.0
#   2.0 - 80/125  = 2.0 - 0.64 = 1.36

weights = tensor([1.60, 1.00, 1.00, 1.36])
# Higher weight = more important sample

Standard MSE component: (pred - target)². Element-wise operation produces a tensor of squared errors for each sample.

# pred   = tensor([100.0, 45.0, 80.0, 120.0])
# target = tensor([ 90.0, 50.0, 85.0, 110.0])

# Step 1: Compute differences (pred - target)
diff = tensor([10.0, -5.0, -5.0, 10.0])

# Step 2: Square the differences
squared_errors = tensor([100.0, 25.0, 25.0, 100.0])

# Each element: (pred[i] - target[i])²

Compute weighted mean: Σ(w × e²) / Σw. Normalizing by sum of weights ensures the loss scale is consistent regardless of weight values.

# weights        = tensor([1.60, 1.00, 1.00, 1.36])
# squared_errors = tensor([100.0, 25.0, 25.0, 100.0])

# Step 1: Element-wise multiplication
weighted = weights * squared_errors
# weighted = [160.0, 25.0, 25.0, 136.0]

# Step 2: Sum weighted errors
sum_weighted = 160.0 + 25.0 + 25.0 + 136.0 = 346.0

# Step 3: Sum weights
sum_weights = 1.60 + 1.00 + 1.00 + 1.36 = 4.96

# Step 4: Compute weighted average
weighted_loss = 346.0 / 4.96 = 69.76

Return the scalar loss tensor. Gradients will flow through this for backpropagation.

Instantiate with r_max=125. This sets the threshold where samples are considered healthy and receive minimum weight.

Create test targets spanning the full RUL range. Note: 150 is above r_max to test capping behavior.

Cap all values at r_max. Any RUL > 125 becomes 125 for weight calculation. This prevents negative weights.

150.0 → 125.0 (capped)

Apply w = 2 - y/125. This linear formula gives: RUL=0 → w=2, RUL=125 → w=1. Simple and effective.

Weight = 2.0 (maximum). Samples at imminent failure are weighted 2× more than healthy samples. Errors here cost twice as much.

Weight = 1.6. Moderate emphasis on mid-degradation samples. Still 60% more important than healthy samples.

Weight = 1.0 (minimum). Healthy samples are still learned, but errors here are acceptable. The model can afford to be less accurate.

After capping: treated same as RUL=125, weight=1.0. This prevents any sample from having less than baseline importance.

The main AMNL loss class. Inherits from nn.Module for PyTorch compatibility. This is the core contribution of our research.

Configure AMNL behavior: r_max (RUL cap), beta (EMA smoothing), lambda values (task weights), and number of health classes.

EMA smoothing factor. β=0.99 provides stable normalization with ~100-step memory. Lower values adapt faster but may be noisy.

beta=0.99: stable, beta=0.9: responsive, beta=0.999: very smooth

Weight for RUL loss in final combination. Default 0.5 gives equal importance to both tasks. Increase for RUL-focused training.

lambda_rul=0.7, lambda_health=0.3 → 70% RUL focus

Weight for health classification loss. λ_rul + λ_health typically sum to 1.0, but this is not required.

Create WeightedMSELoss instance for RUL prediction. Sample weighting emphasizes critical (low RUL) samples.

Standard cross-entropy for 3-class health classification (Healthy, Degrading, Critical). No sample weighting needed here.

Separate normalizer for RUL loss. Tracks running average of RUL loss magnitude for scale-invariant normalization.

Separate normalizer for health loss. Each task has its own EMA to handle different loss scales independently.

Main computation. Takes predictions and targets for both tasks, returns total loss and metrics dictionary for logging.

Call WeightedMSELoss. Returns a tensor with gradients attached. Typical values: 500-3000 depending on prediction quality.

pred=[100, 80], target=[90, 70] → rul_loss ≈ 150.0

Cross-entropy between logits and class labels. Typical values: 0.1-2.0. Much smaller scale than RUL loss!

logits=[[2,0,-1]], target=[0] → health_loss ≈ 0.4

Use .item() to extract scalar Python float for logging. This does not affect gradients - just copies the value.

Update EMA with raw RUL loss value. Returns bias-corrected average for normalization denominator.

rul_loss=1500 → rul_norm≈1480 (running average)

Same for health loss. Separate EMA tracks health-specific statistics independently.

health_loss=0.8 → health_norm≈0.75 (running average)

Prevent division by zero. If EMA is tiny (early training, very good model), clamp to minimum value 1e-8.

KEY STEP: Divide loss by its running average. Result oscillates around 1.0 regardless of absolute magnitude. Gradients flow through!

rul_loss=1500, rul_norm=1480 → normalized≈1.01

Same normalization for health. Now both normalized losses are on same scale (~1.0), enabling fair combination.

health_loss=0.8, health_norm=0.75 → normalized≈1.07

Final AMNL loss: λ_rul × normalized_rul + λ_health × normalized_health. With equal weights: total ≈ 0.5×1.01 + 0.5×1.07 = 1.04

total_loss = 0.5×1.01 + 0.5×1.07 = 1.04

Return comprehensive metrics for logging/visualization. Track raw losses, normalized losses, EMAs, and total loss.

Return tuple of (loss_tensor, metrics_dict). Loss tensor has gradients for backprop. Metrics dict is for logging only.

Clear EMA state when starting fresh training. Important to call between separate training runs.

Static utility method to convert continuous RUL values to discrete health class labels. Used to create health targets from RUL targets.

Start with all zeros (Healthy class). torch.zeros_like creates a tensor of same shape as input, filled with zeros.

# Input RUL tensor
rul = tensor([150.0, 100.0, 40.0, 200.0])

# torch.zeros_like(rul, dtype=torch.long)
health = tensor([0, 0, 0, 0])  # dtype=long for class indices

# All samples start as class 0 (Healthy)

Boolean indexing: health[condition] = value. All samples where RUL ≤ 125 get reassigned to class 1 (Degrading).

# rul    = tensor([150.0, 100.0, 40.0, 200.0])
# health = tensor([0, 0, 0, 0])  # Before

# Boolean mask: rul <= 125.0
mask = tensor([False, True, True, False])
#               150>125  100≤125  40≤125  200>125

# health[mask] = 1
health = tensor([0, 1, 1, 0])  # After
#                   ↑  ↑
#            These changed to 1 (Degrading)

Second pass: samples where RUL ≤ 50 override to class 2 (Critical). This overwrites the Degrading class for very low RUL.

# rul    = tensor([150.0, 100.0, 40.0, 200.0])
# health = tensor([0, 1, 1, 0])  # After first assignment

# Boolean mask: rul <= 50.0
mask = tensor([False, False, True, False])
#               150>50  100>50  40≤50  200>50

# health[mask] = 2
health = tensor([0, 1, 2, 0])  # Final result
#                      ↑
#            40 → Critical (overwrites Degrading)

# Final mapping:
# RUL=150 → 0 (Healthy, above threshold)
# RUL=100 → 1 (Degrading)
# RUL=40  → 2 (Critical)
# RUL=200 → 0 (Healthy)

Create AMNL loss with recommended hyperparameters. These values work well for C-MAPSS without tuning.

Model predictions for RUL. Shape (batch_size,). Values represent predicted remaining cycles until failure.

pred=[100, 45, 80, 120] cycles

Ground truth RUL values. Errors: [10, -5, -5, 10] cycles. Mix of over and under predictions.

Model outputs for health classification (3 classes). Raw scores before softmax. Random for this example.

Use static method to convert RUL targets to health labels. RUL=90,85,110 → Degrading (1), RUL=50 → Critical (2).

rul=[90,50,85,110] → health=[1,2,1,1]

Compute AMNL loss. Returns tensor (for backprop) and metrics dict (for logging). This is where all the magic happens.

Weighted MSE before normalization. Scale depends on prediction quality. Here: 112.5 (small errors).

Cross-entropy before normalization. Much smaller scale (~1.3) than RUL loss (~112). This is why we normalize!

After dividing by EMA: ~1.0. Scale-invariant! The 112.5 raw loss is now directly comparable to health loss.

Also ~1.0 after normalization. Now both tasks contribute equally to gradients, regardless of raw magnitudes.

Final combined loss: 0.5×1.0 + 0.5×1.0 = 1.0. This is what gets backpropagated. Balanced contribution from both tasks.

Compute gradients. Normalization is differentiable (division by constant), so gradients flow correctly to both task heads.

Takes model, data, optimizer, loss function, and device. Returns dictionary of averaged metrics for the epoch.

Enable training mode. This activates dropout, batch normalization uses batch statistics. Critical for correct training behavior.

Dictionary to accumulate per-batch metrics. Will average at the end for epoch-level statistics.

Iterate over batches. enumerate gives batch index for logging. Each iteration processes one batch of samples.

Batch 0: 32 samples, Batch 1: 32 samples, ...

Transfer input tensor to GPU (if available). Shape is (batch, sequence_length, features) for time series data.

x.shape = (32, 30, 17) on CUDA

Transfer RUL targets. Must be on same device as predictions for loss computation.

Convert continuous RUL to discrete health classes. Uses static method, then moves result to correct device.

rul_target=[100,40,80] → health_target=[1,2,1]

Clear gradients from previous batch. PyTorch accumulates gradients by default, so this is required before each backward pass.

Run input through model. Returns two outputs: RUL predictions and health classification logits.

rul_pred.shape=(32,), health_logits.shape=(32,3)

Core AMNL computation. Normalizes both losses, combines with weights, returns loss tensor and metrics.

Compute gradients via backpropagation. Gradients flow through both task heads back to shared encoder.

Clip gradient norms to max_norm=1.0. Prevents exploding gradients which can destabilize training, especially with RNNs.

If ||grad|| > 1.0, scale down: grad = grad * 1.0/||grad||

Apply gradients to update model parameters. Uses optimizer learning rate and momentum settings.

Store each batch metrics in lists for later averaging. Handles dynamic keys from metrics dictionary.

Compute mean of each metric across all batches. Returns epoch-level statistics for logging.

50 batches × 7 metrics → 7 averaged values

Takes predicted and target RUL tensors, returns scalar score. Lower is better. This is the official C-MAPSS evaluation metric.

d = pred - target. Positive d means predicted RUL > actual RUL (late prediction, dangerous!). Negative d means early prediction.

# pred   = tensor([100.0, 50.0, 80.0, 70.0])
# target = tensor([ 90.0, 60.0, 80.0, 85.0])

# Element-wise subtraction
d = pred - target
d = tensor([10.0, -10.0, 0.0, -15.0])
#            ↑      ↑     ↑     ↑
#          late  early  perfect early

# Interpretation:
# d=10:  predicted 100, actual 90 → 10 cycles LATE (dangerous!)
# d=-10: predicted 50, actual 60  → 10 cycles EARLY (safe but costly)
# d=0:   predicted 80, actual 80  → PERFECT prediction
# d=-15: predicted 70, actual 85  → 15 cycles EARLY

torch.where applies different formulas based on error sign. This penalizes late predictions more than early ones.

# torch.where(condition, value_if_true, value_if_false)

# d = tensor([10.0, -10.0, 0.0, -15.0])

# For each element:
# if d < 0:  use exp(-d/13) - 1  (early formula)
# if d >= 0: use exp(d/10) - 1   (late formula)

# Result: applies different penalty to each element

Formula: exp(-d/13) - 1. Denominator 13 means gentler penalty. Predicting failure too early is costly but not catastrophic.

# For d = -10 (early by 10 cycles):
# exp(-(-10)/13) - 1 = exp(10/13) - 1
# = exp(0.769) - 1
# = 2.158 - 1
# = 1.158

# For d = -15 (early by 15 cycles):
# exp(15/13) - 1 = exp(1.154) - 1 = 2.17

Formula: exp(d/10) - 1. Denominator 10 means harsher penalty. Late predictions risk missing actual failure!

# For d = 10 (late by 10 cycles):
# exp(10/10) - 1 = exp(1) - 1
# = 2.718 - 1
# = 1.718

# For d = 0 (perfect):
# exp(0/10) - 1 = exp(0) - 1 = 0

# Compare same magnitude:
# d=-10 (early): score = 1.158
# d=+10 (late):  score = 1.718
# Late is 48% worse than early!

Sum individual sample scores. Total NASA score for the test set. Lower is better. State-of-art on FD001: ~250.

# d = tensor([10.0, -10.0, 0.0, -15.0])

# Individual scores (after torch.where):
# d=10 (late):    exp(10/10)-1   = 1.718
# d=-10 (early):  exp(10/13)-1   = 1.158
# d=0 (perfect):  exp(0/10)-1    = 0.000
# d=-15 (early):  exp(15/13)-1   = 2.170

score = tensor([1.718, 1.158, 0.000, 2.170])

# Sum all scores
total = score.sum().item()
# total = 1.718 + 1.158 + 0.0 + 2.170 = 5.046

Same magnitude error (10 cycles) but different direction: late prediction scored 1.72, early scored 1.16. Late is ~50% worse!

# Complete scoring example:
# pred   = [100, 50, 80]
# target = [ 90, 60, 80]
# d      = [ 10,-10,  0]

# Scores:
# d=10 (late):   exp(10/10)-1 = e¹-1  = 1.718
# d=-10 (early): exp(10/13)-1 = e⁰·⁷⁷-1 = 1.158
# d=0 (perfect): exp(0)-1     = 1-1   = 0.000

# Total NASA score = 1.718 + 1.158 + 0 = 2.876

# Key insight: Same |error|=10, but:
# Late (dangerous):  1.718
# Early (safe):      1.158
# Ratio: 1.718/1.158 = 1.48 (48% worse)