The GRACE Loss Equation

An Equation You Can Read In One Breath

A pharmacist filling a prescription does two things at once. They weigh each active ingredient against the others — more antibiotic, less antihistamine, an exact ratio — and they shape how much of the dose hits each tissue via the delivery vehicle. A capsule with the same total mass can be tuned to release fast or slow, in the stomach or the small intestine. The two adjustments live on different axes, and a working prescription engages both.

GRACE is a prescription. The outer ratio is GABA, deciding how the gradient budget is split between the RUL and health tasks. The inner shaping is the failure-biased weighted MSE, deciding which samples inside the RUL task carry the heaviest squared error. This section writes that prescription as a single equation, walks every symbol, traces the four-stage controller pipeline, computes a worked example end to end, and shows the same equation in production PyTorch.

Anatomy: Six Symbols, Two Axes

Six symbols carry every piece of GRACE's machinery. Three are per-task (outer), three are per-sample (inner):

Read the equation top to bottom: the outer factor $\lambda^*_i(t)$ is a step-varying scalar that depends on this batch's gradient ratio; the inner factor $(1/N) \sum_j w(y_j)\,(\hat y_j - y_j)^2$ is a per-batch scalar that depends only on the predictions and targets. The two factors multiply and then sum across the two tasks. That is the entire model.

Inner Expansion: The Weighted MSE

Expand the inner term explicitly. For one mini-batch of size $N$:

$\frac{1}{N}\sum_{j=1}^{N} w(y_j)\,(\hat{y}_j - y_j)^2 \;=\; \frac{1}{N}\sum_{j=1}^{N} \Big[1 + \mathrm{clip}\!\left(1 - \tfrac{y_j}{125},\ 0,\ 1\right)\Big]\,(\hat{y}_j - y_j)^2.$

The bracketed factor is the per-sample weight. It rewrites cleanly as a piecewise function:

The piecewise structure has a physical reading: a failure prediction error is twice as costly as the same error during early-life operation. The clip at $y \geq 125$ matches the paper's piecewise-linear RUL target — for engines that have not started visibly degrading yet, every cycle looks the same to the labels, so the loss treats every cycle the same too. There is no asymmetry to encode at the healthy end.

Outer Expansion: The Four-Stage GABA Pipeline

The outer factor is not a closed-form expression in the usual sense; it is the output of a four-stage controller. Stage by stage, with $K=2$ tasks:

Stage 1: Closed form

$\lambda^{\text{raw}}_i(t) \;=\; \dfrac{\sum_{j\neq i} g_j(t)}{(K{-}1)\,\sum_{j} g_j(t)}, \qquad g_i(t) \;=\; \bigl\| \nabla_{\theta_s}\mathcal{L}_i(t)\bigr\|_2.$

Inverse-ratio: the task with the smaller gradient gets the larger raw weight. For $K=2$ this collapses to $\lambda^{\text{raw}}_{\text{rul}} = g_{\text{health}}/(g_{\text{rul}}+g_{\text{health}})$.

Stage 2: Exponential moving average

$\bar\lambda_i(t) \;=\; \beta\,\bar\lambda_i(t{-}1) \;+\; (1-\beta)\,\lambda^{\text{raw}}_i(t), \qquad \bar\lambda_i(0) = 1/K.$

Smooths out per-batch noise in $\lambda^{\text{raw}}_i(t)$. Paper default $\beta = 0.99$ — an effective memory of $1/(1-\beta) \approx 100$ steps. The starting value $1/K$ is uniform.

Stage 3: Per-element floor

$\tilde\lambda_i(t) \;=\; \max\bigl(\bar\lambda_i(t),\ \varepsilon\bigr), \qquad \varepsilon = 0.05.$

Prevents either task from being driven to zero. Without this clamp, after enough EMA iterations the over-gradient task (RUL on C-MAPSS) would collapse to a numerically negligible weight and the joint optimisation would degenerate to single-task health training.

Stage 4: Renormalisation

$\lambda^*_i(t) \;=\; \dfrac{\tilde\lambda_i(t)}{\sum_{k=1}^{K} \tilde\lambda_k(t)}.$

Restores the sum-to-one property after the floor breaks it. The final $\lambda^*_i(t)$ is a proper probability distribution over the K tasks at every step.

What .backward() Actually Computes

The whole point of writing the equation is to differentiate it. The gradient with respect to a backbone parameter $\theta_s$ is

$\nabla_{\theta_s}\,\mathcal{L}_{\text{GRACE}}(t) \;=\; \lambda^*_{\text{rul}}(t)\,\nabla_{\theta_s}\,\mathcal{L}^{\text{WMSE}}_{\text{rul}}(t) \;+\; \lambda^*_{\text{health}}(t)\,\nabla_{\theta_s}\,\mathcal{L}_{\text{CE}}(t).$

Two structural facts to internalise:

• The lambdas $\lambda^*_i(t)$ are constants from autograd's point of view. The implementation calls .detach() before the multiply (gaba.py:131), so PyTorch never tries to differentiate through the gradient-norm computation. Forgetting this turns GABA into a meta-gradient method (closer to GradNorm) with ~2× memory and a different convergence law.
• The inner per-sample weights $w(y_j)$ depend only on $y_j$, not on $\hat y_j$, so they are also constants for autograd. The gradient simplifies to $(2/N)\sum_j w(y_j)\,(\hat y_j - y_j)\,\nabla_{\theta_s}\hat y_j$ — the standard MSE backward, multiplied per-sample by $w(y_j)$.

Interactive: Deconstructing The Equation

Click any symbol pill below the rendered equation to see its definition, code reference, and current numeric value. The three sliders drive the OUTER controller live: vary the gradient ratio, the EMA $\beta$, and the number of EMA steps to watch $\lambda^*_i$ evolve from $(0.5, 0.5)$ at $t=0$ toward the floor-clamped converged value at large $t$.

Try this. Set $\log_{10}(g_{\text{rul}}/g_{\text{health}}) = 2.84$ (the C-MAPSS measured value), $\beta = 0.99$, and slide steps from 0 → 1000. You will see $\lambda^*_{\text{rul}}$ drift down from 0.5 toward the floor at 0.0477; the floor activates around step 400 and the renormalisation kicks in. That trajectory is what the paper's gradient_logger.py records during real training.

Python: One Training Step, Symbol By Symbol

Walk every line. The script reproduces the four-stage GABA pipeline by hand, computes the failure-biased MSE explicitly, and prints every intermediate value — raw, smoothed, floored, renormalised — so the equation can be checked numerically before we trust it to PyTorch.

y_true       = [10, 30, 60, 90, 110, 5, 15, 80]
y_pred       = [15, 26, 68, 87, 112, -7, 24, 74]
residual^2   = [25, 16, 64, 9, 4, 144, 81, 36]
w(y)         = [1.92, 1.76, 1.52, 1.28, 1.12, 1.96, 1.88, 1.36]
w * r^2      = [48.0, 28.16, 97.28, 11.52, 4.48, 282.24, 152.28, 48.96]
L_rul^WMSE   = 84.1150

g_rul        = 26.4016    g_health = 0.037833
raw lambdas  = (0.001431, 0.998569)
EMA lambdas  = (0.495014, 0.504986)    (β=0.99, prev=0.5)
floored      = (0.495014, 0.504986)    (ε=0.05)
final λ*     = (0.495014, 0.504986)    (renormalised)

L_GRACE      = 0.4950*84.1150 + 0.5050*0.6069
             = 41.9446

1"""GRACE loss equation — every symbol, every value, in NumPy."""
2
3import numpy as np
4
5
6# ---------- Mini-batch ----------
7y_true = np.array([10, 30, 60, 90, 110,  5, 15, 80], dtype=float)
8y_pred = np.array([15, 26, 68, 87, 112, -7, 24, 74], dtype=float)
9N      = len(y_true)
10L_health = 0.6069                    # cross-entropy from same forward pass
11
12
13# ---------- Inner axis: weighted MSE ----------
14def w_failure(y, max_rul=125.0):
15    """Per-sample weight w(y) = 1 + clip(1 - y/max_rul, 0, 1)."""
16    return 1.0 + np.clip(1.0 - y / max_rul, 0.0, 1.0)
17
18
19w           = w_failure(y_true)
20residual_sq = (y_pred - y_true) ** 2
21L_rul_w     = (w * residual_sq).mean()
22
23
24# ---------- Outer axis: GABA closed form, EMA, floor, renorm ----------
25def gaba_step(g_rul, g_health, prev, beta=0.99, eps=0.05):
26    """One step of the four-stage GABA pipeline. Returns (lam_rul, lam_h)."""
27    # Stage 1: closed form (raw)
28    S = g_rul + g_health
29    raw = np.array([g_health / S, g_rul / S])
30    # Stage 2: EMA smoothing
31    smoothed = beta * np.array(prev) + (1.0 - beta) * raw
32    # Stage 3: floor
33    floored = np.maximum(smoothed, eps)
34    # Stage 4: renormalise so weights sum to 1
35    normed = floored / floored.sum()
36    return normed, raw, smoothed, floored
37
38
39# Real measurements from the same forward pass (chapter 18 §1)
40g_rul, g_health = 26.4016, 0.037833
41prev_lam        = (0.5, 0.5)
42lam, raw, sm, fl = gaba_step(g_rul, g_health, prev_lam)
43
44
45# ---------- Compose: outer * inner ----------
46L_GRACE = lam[0] * L_rul_w + lam[1] * L_health
47
48
49# ---------- Inspect every symbol ----------
50print(f"y_true       = {y_true.astype(int).tolist()}")
51print(f"y_pred       = {y_pred.astype(int).tolist()}")
52print(f"residual^2   = {residual_sq.astype(int).tolist()}")
53print(f"w(y)         = {np.round(w, 3).tolist()}")
54print(f"w * r^2      = {np.round(w * residual_sq, 2).tolist()}")
55print(f"L_rul^WMSE   = {L_rul_w:.4f}")
56print()
57print(f"g_rul        = {g_rul:.4f}    g_health = {g_health:.6f}")
58print(f"raw lambdas  = ({raw[0]:.6f}, {raw[1]:.6f})")
59print(f"EMA lambdas  = ({sm[0]:.6f}, {sm[1]:.6f})    (β=0.99, prev=0.5)")
60print(f"floored      = ({fl[0]:.6f}, {fl[1]:.6f})    (ε=0.05)")
61print(f"final λ*     = ({lam[0]:.6f}, {lam[1]:.6f})    (renormalised)")
62print()
63print(f"L_GRACE      = {lam[0]:.4f}*{L_rul_w:.4f} + {lam[1]:.4f}*{L_health:.4f}")
64print(f"             = {L_GRACE:.4f}")

PyTorch: The Same Equation In Production Code

Now the equation as the paper actually runs it. GABALoss is an nn.Module — a stateful controller that keeps the EMA buffer between steps and exposes inspection helpers. The whole composition is a single call on line 49: L_GRACE = gaba(L_rul_w, L_health, shared_params=shared).

L_rul^WMSE = ~50000
L_health   = ~1.10
g_rul      = ~30000
g_health   = ~0.5
raw_λ_rul  = ~0.000017
final λ*   = (0.4950, 0.5050)
L_GRACE    = ~24750

1"""GRACE loss equation — production composition with the paper helpers."""
2
3import torch
4import torch.nn as nn
5import torch.nn.functional as F
6
7from grace.core.gaba           import GABALoss
8from grace.core.weighted_mse   import moderate_weighted_mse_loss
9
10torch.manual_seed(0)
11
12
13# ---------- Tiny dual-task model ----------
14class TinyDualHead(nn.Module):
15    def __init__(self):
16        super().__init__()
17        self.backbone    = nn.Linear(4, 6)
18        self.rul_head    = nn.Linear(6, 1)
19        self.health_head = nn.Linear(6, 3)
20
21    def forward(self, x):
22        feat = torch.relu(self.backbone(x))
23        return self.rul_head(feat).squeeze(-1), self.health_head(feat)
24
25    def get_shared_params(self):
26        return [p for n, p in self.named_parameters() if "head" not in n]
27
28
29# ---------- One mini-batch forward pass ----------
30model = TinyDualHead()
31x         = torch.randn(8, 4)
32y_true    = torch.tensor([10., 30., 60., 90., 110., 5., 15., 80.])
33hp_target = torch.tensor([0, 1, 2, 0, 2, 0, 0, 2])
34
35y_pred, hp_logits = model(x)
36
37
38# ---------- Inner axis: failure-biased MSE on RUL ----------
39L_rul_w  = moderate_weighted_mse_loss(y_pred, y_true, max_rul=125.0)
40L_health = F.cross_entropy(hp_logits, hp_target)
41
42
43# ---------- Outer axis: GABA controller ----------
44gaba = GABALoss(beta=0.99, warmup_steps=0, min_weight=0.05, n_tasks=2)
45shared = model.get_shared_params()
46
47# One step of GABA: the controller computes per-task g_i internally,
48# applies EMA + floor + renorm, and returns the COMBINED loss.
49L_GRACE = gaba(L_rul_w, L_health, shared_params=shared)
50
51
52# ---------- Inspect the controller state ----------
53weights      = gaba.get_weights()                        # final λ*
54grad_stats   = gaba.get_gradient_stats()                 # g_i + raw_λ_i
55
56print(f"L_rul^WMSE = {L_rul_w.item():.4f}")
57print(f"L_health   = {L_health.item():.4f}")
58print(f"g_rul      = {grad_stats['grad_norm_rul']:.4f}")
59print(f"g_health   = {grad_stats['grad_norm_health']:.4f}")
60print(f"raw_λ_rul  = {grad_stats['raw_weight_rul']:.6f}")
61print(f"final λ*   = ({weights['rul_weight']:.6f}, {weights['health_weight']:.6f})")
62print(f"L_GRACE    = {L_GRACE.item():.4f}")

A Worked Numerical Example, End To End

The same 8-sample mini-batch as section 21·1, with every intermediate value at every stage. Inputs first, then the inner axis, then the four-stage outer pipeline, then the composed loss. Numbers come from the printed output of the NumPy script above.

Inner axis on this batch

Standard MSE would be $379/8 = 47.375$; weighted MSE is $672.92/8 = 84.115$. The ratio $84.115/47.375 \approx 1.78$ depends on the y-distribution of THIS batch — for a batch of all-healthy engines (every $y_j \geq 125$) the two losses would be identical. The inner axis bites only when the batch contains failure-region samples.

Outer axis on this batch

The two EMA rows show the controller's two regimes. Early in training (n=1), $\lambda^*$ is essentially uniform — GABA is still ‘learning’ the gradient balance. After convergence (n=∞), the floor activates, fixing $\lambda^*_{\text{rul}}$ at 0.0477 and $\lambda^*_{\text{health}}$ at 0.9523. Production GRACE training spends most of its 500 epochs in the converged regime; the warmup_steps parameter (default 100) explicitly disables adaptation until enough EMA history exists.

Composed GRACE loss

$\mathcal{L}_{\text{GRACE}}(t) \;=\; 0.0477 \cdot 84.115 \;+\; 0.9523 \cdot 0.6069 \;=\; 4.011 + 0.578 \;=\; 4.589 \quad (t = \infty).$

Compare this with the AMNL cell from section 21·1 (cell B): same inner WMSE, but fixed lambdas at $(0.5, 0.5)$ give $0.5 \cdot 84.115 + 0.5 \cdot 0.6069 = 42.36$. The OUTER axis shifts the optimisation target by an order of magnitude. That shift — mediated by the four-stage controller — is what produces the 224 NASA score on FD002 instead of 356.

The Same Equation In Other Fields

The shape $\sum_i \lambda^*_i(t) \cdot \frac{1}{N}\sum_j w_j \ell_{ij}$ appears any time a learner has multiple competing objectives and non-uniform sample importance. A few canonical examples:

In every row, replacing the outer factor with GABA and the inner weight with a domain-specific failure or rarity ramp produces a GRACE-shaped algorithm. The equation is universal; the choice of $w(\cdot)$ is the domain knowledge.

Pitfalls When Reading The Equation

Pitfall 1: confusing the four lambdas

Four different lambda symbols appear in this section: $\lambda^{\text{raw}}_i$ (closed form), $\bar\lambda_i$ (EMA), $\tilde\lambda_i$ (floored), $\lambda^*_i$ (renormalised). Only $\lambda^*_i$ multiplies the loss. The other three are intermediate. Logging the wrong one (especially $\bar\lambda_i$ instead of $\lambda^*_i$) is the most common source of confusion when debugging GABA-style controllers — the controller's actual decision is post-floor, post-renorm.

Pitfall 2: thinking$w(y_j)$needs a gradient

$w(y_j)$ depends on $y_j$, the ground truth, not on $\hat y_j$. No backward pass flows through the weight. Treating the weights as needing a gradient (e.g. by making them learnable) is a different algorithm — closer to attention-weighted MSE — and produces qualitatively different training dynamics.

Pitfall 3: forgetting the (1/N) factor

The mean over the mini-batch makes the inner term scale-invariant in batch size. If the implementation uses reduction='sum' instead of reduction='mean', the per-task gradient norms grow with batch size and the closed-form lambdas drift accordingly. The paper consistently uses mean reduction; check before changing the loss aggregation.

Pitfall 4: assuming the equation is differentiable in all symbols

It is differentiable in $\hat y_j$ and the health logits. It is not differentiable in $y_j$ (one-sided clip), $\beta$, or $\varepsilon$ by design — these are controller hyperparameters, set per-experiment, never learned in the same loop as the model. Hyperparameter search exists in a separate outer loop (chapter 22 §2 covers the GRACE search protocol).

Takeaway

• GRACE is one equation: $\mathcal{L}_{\text{GRACE}}(t) = \lambda^*_{\text{rul}}(t) \cdot \mathcal{L}^{\text{WMSE}}_{\text{rul}}(t) + \lambda^*_{\text{health}}(t) \cdot \mathcal{L}_{\text{CE}}$.
• The inner factor is a per-sample failure ramp $w(y_j) = 1 + \mathrm{clip}(1 - y_j/125, 0, 1)$ times the squared residual, averaged.
• The outer factor is the output of a four-stage controller: closed form → EMA → floor → renormalise. Only the final $\lambda^*_i(t)$ multiplies the loss.
• The lambdas are detached before the multiply — autograd treats them as constants. Forgetting this turns GABA into a different algorithm (GradNorm-style meta-gradients) with double the memory.
• On the 8-sample worked example: converged $\lambda^* = (0.0477, 0.9523)$, inner WMSE = 84.12, composed $\mathcal{L}_{\text{GRACE}} \approx 4.59$ — an order of magnitude below the AMNL cell at the same batch.