Per-Step Gradient Norm Computation

The Sensor Inside The Loop

Walk into a modern car's ABS controller and you find, before any control logic, a wheel-speed sensor sampled many times per second. The controller can't modulate brake pressure unless it MEASURES wheel speed first. Throw the sensor away and the algorithm is just an opinion.

GABA has the same anatomy. The closed form $\lambda^*_i = g_j / (g_i + g_j)$ is the controller. Its sensor is the per-task gradient norm $g_i = \| \nabla_{\theta_s} \mathcal{L}_i \|_2$ on the shared backbone parameters $\theta_s$. Every training step starts with a measurement, then applies the rule. This section is about the measurement: how to compute $g_i$ exactly, cheaply, and without contaminating the rest of the optimisation.

What ‘Per-Step’ Means In GABA

The GABA algorithm (paper Algorithm 1) is a control loop running once per training step. The four stages, in order:

• Measure. Compute $g_i = \| \nabla_{\theta_s} \mathcal{L}_i \|_2$ for each task. THIS SECTION.
• Compute. Apply the closed form $\lambda_i = (S - g_i) / ((K{-}1) S)$ where $S = \sum_j g_j$ (§17.3 derivation).
• Stabilise. EMA-smooth (§18.2), apply the floor (§18.3), maybe pass through warmup (§18.4).
• Combine and step. Form $\mathcal{L} = \sum_i \lambda^*_i \mathcal{L}_i$ and let the main optimiser take one step.

Stages 1 and 4 both involve gradients on the same model, but they are NOT the same computation. Stage 1 needs the SCALAR $g_i$ per task; the actual gradient tensors get thrown away. Stage 4 needs the actual gradient of the combined loss applied to all parameters (including heads). The rest of this section makes Stage 1 explicit.

Selecting Shared Parameters

The first decision is: which parameters does $\theta_s$ actually contain? GABA balances task gradients on the shared backbone, not on every parameter in the model. Including the task heads is wrong: each head's gradient norm is large for its own task and zero for the others, so head parameters artificially inflate one side of the imbalance and deflate the other.

The paper's utility get_shared_params(model, head_names=("rul_head", "health_head")) in grace/core/gradient_utils.py:13 walks model.named_parameters() and excludes any parameter whose name contains a configured head substring. This is robust to PyTorch wrapping (EMA, DataParallel) because named_parameters() emits dotted-path names that retain the head substring.

Why substring matching, not type matching? Type-matching (e.g. ‘exclude all nn.Linear layers’) would also catch the inner backbone layers. Substring on the registered name is a stronger contract because the model author already chose those names to mark roles. The user can pass a different head_names tuple if their architecture differs.

Why The L2 Norm (Not L1 Or L∞)

Three norm choices satisfy ‘a positive scalar that increases with magnitude’:

The L2 norm is chosen for three reasons:

• Rotation invariance. If you rotate the parameter basis (e.g. SVD reparametrisation), $\| g \|_2$ does not change. L1 and L∞ do change — they are basis-dependent. Multi-task balance should not be basis-dependent.
• Standard practice in optimisation. PyTorch's torch.nn.utils.clip_grad_norm_, Adam's second-moment normaliser, and almost every gradient- based regulariser default to L2. Using the same norm downstream avoids subtle interactions.
• Smooth gradient. The L2 norm is differentiable everywhere except at zero (a measure- zero point). L1 has a kink at every coordinate axis; L∞ has many kinks. GABA itself is one-shot so this matters less than for GradNorm, which back-propagates through the norm.

Sum-Of-Squares Equals Concat-Then-Norm

Real backbones store gradients as a list of per-parameter tensors of different shapes (one weight matrix here, one bias vector there). The L2 norm of the full gradient vector $g \in \mathbb{R}^{D}$ with $D = \sum_p \dim(\theta_p)$ could be computed by concatenating everything into one big 1-D vector and calling np.linalg.norm:

$\| g \|_2 = \sqrt{ \sum_{p} \sum_{i} g_{p,i}^2 } = \sqrt{ \sum_{p} \| g_p \|_2^2 }$

The right-hand side is the paper's implementation choice: accumulate squared per-parameter norms, then take the square root once at the end. The two are numerically identical (we verify to 4 decimals in the Python demo). The sum-of-squares version has two practical advantages:

• No materialisation. Concatenation would allocate a fresh $D$ float buffer. For the paper's 3.5M-parameter backbone that is 14 MB of temporary memory per gradient norm. The sum-of-squares version reuses each parameter gradient in place.
• Pipelines with the autograd output. torch.autograd.grad returns a tuple of per-parameter gradient tensors already — the sum-of-squares loop walks them directly without first concatenating.

One Forward Pass, K Backward Calls

For K=2 tasks, GABA needs TWO gradient norms per step: $g_{\text{rul}}$ and $g_{\text{health}}$. Naively that is two forward passes — expensive. The standard PyTorch idiom is one forward pass and two torch.autograd.grad calls with retain_graph=True:

• The first autograd.grad(rul_loss, ..., retain_graph=True) computes $\nabla L_{\text{rul}}$ AND keeps the autograd graph alive.
• The second autograd.grad(health_loss, ..., retain_graph=True) re-uses the same forward graph to compute $\nabla L_{\text{health}}$. Because the graph is still alive, no second forward pass is needed.
• After both calls, the trainer typically calls combined_loss.backward() to actually update the weights. This third pass also re-uses the same graph; on the third call, PyTorch finally frees it.

Interactive: Which Parameters Count?

Toggle individual parameters in or out of the gradient-norm aggregation. The default mirrors the paper's get_shared_params: backbone IN, heads OUT. Watch how including a head parameter inflates one side's norm and changes the resulting GABA λ.

Try this. Start with the paper default, then toggle rul_head.weight on. The RUL norm jumps from $\sim 110$ to $\sim 180$ because that one head parameter contributes a 145.2-unit gradient by itself. The health norm doesn't change (the rul_head is detached from the health loss). Now $\lambda_{\text{rul}}$ DROPS — you accidentally told GABA that the RUL task was even more dominant than it really is on the shared backbone, which is wrong.

Python: Aggregating Per-Parameter Norms From Scratch

Build both versions of the L2 norm in pure NumPy: the paper's sum-of-squares loop and the reference concat-then-norm. Run them on a 4-tensor synthetic backbone and verify they produce identical answers to 4-decimal precision.

param  | shape    |  ||g_rul_p|| | ||g_health_p||
------------------------------------------------------------
W1     | (3, 4)   |      22.7259 |       0.022468
b1     | (3,)     |       6.1986 |       0.011417
W2     | (2, 3)   |       9.4379 |       0.027483
b2     | (2,)     |       7.2857 |       0.006390

||g_rul||    sum-of-squares    = 26.4016
||g_rul||    concat-then-norm   = 26.4016
||g_health|| sum-of-squares    = 0.037833
||g_health|| concat-then-norm   = 0.037833

lambda_rul    = 0.001431
lambda_health = 0.998569

W1 grad: shape (3, 4) — each element ~ N(0, 1) × 5.0
b1 grad: shape (3,)  — each element ~ N(0, 1) × 5.0
W2 grad: shape (2, 3) — each element ~ N(0, 1) × 5.0
b2 grad: shape (2,)  — each element ~ N(0, 1) × 5.0

||W1|| ≈ 22.7259
||b1|| ≈  6.1986
||W2|| ≈  9.4379
||b2|| ≈  7.2857

||W1|| ≈ 0.022468
||b1|| ≈ 0.011417
||W2|| ≈ 0.027483
||b2|| ≈ 0.006390

def grad_norm_sum_of_squares_debug(grads, label='grads'):
    print(f'\n=== {label} ===')
    sq = 0.0
    for i, g in enumerate(grads):
        contribution = (g ** 2).sum()
        sq += contribution
        print(f'  param {i}: shape={str(g.shape):<8} '
              f'||g_p||={np.sqrt(contribution):>10.6f} '
              f'sq_running={sq:>12.6f}')
    norm = np.sqrt(sq)
    print(f'  FINAL ||g||_2 = {norm:.6f}')
    return norm

# Run on the same toy example
g_rul    = grad_norm_sum_of_squares_debug(shared_grads_rul,    'RUL')
g_health = grad_norm_sum_of_squares_debug(shared_grads_health, 'health')

=== RUL ===
  param 0: shape=(3, 4)  ||g_p||= 22.725891 sq_running=  516.466002
  param 1: shape=(3,)    ||g_p||=  6.198651 sq_running=  554.889266
  param 2: shape=(2, 3)  ||g_p||=  9.437935 sq_running=  643.963431
  param 3: shape=(2,)    ||g_p||=  7.285717 sq_running=  697.039406
  FINAL ||g||_2 = 26.401504

=== health ===
  param 0: shape=(3, 4)  ||g_p||=  0.022468 sq_running=    0.000505
  param 1: shape=(3,)    ||g_p||=  0.011417 sq_running=    0.000635
  param 2: shape=(2, 3)  ||g_p||=  0.027483 sq_running=    0.001390
  param 3: shape=(2,)    ||g_p||=  0.006390 sq_running=    0.001431
  FINAL ||g||_2 = 0.037833

import numpy as np

np.random.seed(1)
grads = [np.random.randn(3, 4) * 5.0,
         np.random.randn(3)    * 5.0,
         np.random.randn(2, 3) * 5.0,
         np.random.randn(2)    * 5.0]

# Method A — paper's sum of squares
sq = sum((g ** 2).sum() for g in grads)
norm_a = np.sqrt(sq)

# Method B — concat then L2 norm
flat   = np.concatenate([g.reshape(-1) for g in grads])
norm_b = np.linalg.norm(flat, ord=2)

print(f'Method A (sum of squares): {norm_a:.10f}')
print(f'Method B (concat-then-norm): {norm_b:.10f}')
print(f'identical to 1e-12: {np.isclose(norm_a, norm_b, atol=1e-12)}')

Method A (sum of squares): 26.4015038258
Method B (concat-then-norm): 26.4015038258
identical to 1e-12: True

1"""Per-task gradient norm - sum-of-squares equals concat-then-norm."""
2
3import numpy as np
4
5np.random.seed(1)
6
7
8# ---------- A miniature shared backbone ----------
9shared_grads_rul = [
10    np.random.randn(3, 4) * 5.0,    # W1 (12 params)
11    np.random.randn(3)    * 5.0,    # b1 (3 params)
12    np.random.randn(2, 3) * 5.0,    # W2 (6 params)
13    np.random.randn(2)    * 5.0,    # b2 (2 params)
14]
15shared_grads_health = [
16    np.random.randn(3, 4) * 0.01,
17    np.random.randn(3)    * 0.01,
18    np.random.randn(2, 3) * 0.01,
19    np.random.randn(2)    * 0.01,
20]
21
22
23def grad_norm_sum_of_squares(grads):
24    """L2 norm via per-parameter squared-sum aggregation (paper code)."""
25    sq = 0.0
26    for g in grads:
27        sq += (g ** 2).sum()
28    return np.sqrt(sq)
29
30
31def grad_norm_via_concat(grads):
32    """Equivalent: flatten + concat + np.linalg.norm."""
33    flat = np.concatenate([g.reshape(-1) for g in grads])
34    return np.linalg.norm(flat, ord=2)
35
36
37# ---------- Per-parameter contributions ----------
38names = ["W1", "b1", "W2", "b2"]
39print(f"{'param':<6} | {'shape':<8} | {'||g_rul_p||':>12} | {'||g_health_p||':>14}")
40print("-" * 60)
41for name, gr, gh in zip(names, shared_grads_rul, shared_grads_health):
42    print(f"{name:<6} | {str(gr.shape):<8} | {np.linalg.norm(gr.reshape(-1)):>12.4f} | {np.linalg.norm(gh.reshape(-1)):>14.6f}")
43
44
45# ---------- Aggregate ----------
46g_rul_sos    = grad_norm_sum_of_squares(shared_grads_rul)
47g_rul_cat    = grad_norm_via_concat(shared_grads_rul)
48g_health_sos = grad_norm_sum_of_squares(shared_grads_health)
49g_health_cat = grad_norm_via_concat(shared_grads_health)
50
51print(f"\n||g_rul||    sum-of-squares    = {g_rul_sos:.4f}")
52print(f"||g_rul||    concat-then-norm   = {g_rul_cat:.4f}")
53print(f"||g_health|| sum-of-squares    = {g_health_sos:.6f}")
54print(f"||g_health|| concat-then-norm   = {g_health_cat:.6f}")
55
56
57# ---------- GABA closed form ----------
58S = g_rul_sos + g_health_sos
59print(f"\nlambda_rul    = {g_health_sos / S:.6f}")
60print(f"lambda_health = {g_rul_sos    / S:.6f}")

PyTorch: The Paper's compute_task_grad_norm()

Now the production version — line-for-line from paper_ieee_tii/grace/core/gradient_utils.py. A toy $14 \to 32 \to 16$ shared backbone with two heads; one forward pass; two gradient norms via torch.autograd.grad with create_graph=False and retain_graph=True; finally the K=2 closed form.

||g_rul||    = 114.4296
||g_health|| = 0.2164
ratio        = 528.8x

lambda_rul    = 0.001888
lambda_health = 0.998112

backbone.fc1.weight   shape (32, 14) — 448 scalars
backbone.fc1.bias     shape (32,)    —  32 scalars
backbone.fc2.weight   shape (16, 32) — 512 scalars
backbone.fc2.bias     shape (16,)    —  16 scalars
Total D = 1 008 shared scalars

rul_head.weight (1, 16) + bias (1,) = 17 scalars
health_head.weight (3, 16) + bias (3,) = 51 scalars
— filtered out by get_shared_params on lines 37–38.

x          ~ N(0, 1) shape (64, 14)
rul_target ~ U(0, 125)  shape (64, 1)
hp_target  ~ randint(0, 3) shape (64,)

Tuple of 4 tensors:
  grad of fc1.weight  shape (32, 14)
  grad of fc1.bias    shape (32,)
  grad of fc2.weight  shape (16, 32)
  grad of fc2.bias    shape (16,)

def compute_task_grad_norm_debug(loss, shared_params, label='task',
                                  retain_graph=True):
    print(f'\n=== {label} ===')
    grads = torch.autograd.grad(loss, shared_params,
                                retain_graph=retain_graph,
                                create_graph=False,
                                allow_unused=True)
    total = torch.tensor(0.0, device=loss.device)
    for i, (p, g) in enumerate(zip(shared_params, grads)):
        if g is None:
            print(f'  [{i}] shape={tuple(p.shape)} grad=None (skipped)')
            continue
        contribution = g.pow(2).sum()
        total = total + contribution
        print(f'  [{i}] shape={str(tuple(p.shape)):<10} '
              f'||g_p||={contribution.sqrt().item():>10.4f} '
              f'total_running={total.item():>12.4f}')
    norm = total.sqrt()
    print(f'  FINAL ||g||_2 = {norm.item():.4f}')
    return norm

# Replace the two calls on lines 65-66 with:
g_rul    = compute_task_grad_norm_debug(rul_loss,    shared, 'RUL',    retain_graph=True)
g_health = compute_task_grad_norm_debug(health_loss, shared, 'health', retain_graph=True)

=== RUL ===
  [0] shape=(32, 14)   ||g_p||=   20.8104 total_running=    433.0707
  [1] shape=(32,)      ||g_p||=   21.0337 total_running=    875.4880
  [2] shape=(16, 32)   ||g_p||=   92.0929 total_running=   9356.5840
  [3] shape=(16,)      ||g_p||=   61.1354 total_running=  13094.1270
  FINAL ||g||_2 = 114.4296

=== health ===
  [0] shape=(32, 14)   ||g_p||=    0.0916 total_running=      0.0084
  [1] shape=(32,)      ||g_p||=    0.0361 total_running=      0.0097
  [2] shape=(16, 32)   ||g_p||=    0.1669 total_running=      0.0375
  [3] shape=(16,)      ||g_p||=    0.0964 total_running=      0.0468
  FINAL ||g||_2 = 0.2164

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

class DualHead(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone    = nn.Sequential(nn.Linear(14, 32), nn.ReLU(),
                                          nn.Linear(32, 16))
        self.rul_head    = nn.Linear(16, 1)
        self.health_head = nn.Linear(16, 3)
    def forward(self, x):
        f = self.backbone(x)
        return self.rul_head(f), self.health_head(f)

def shared_params(model):
    return [p for n, p in model.named_parameters()
            if 'rul_head' not in n and 'health_head' not in n
               and p.requires_grad]

def grad_norm(loss, params):
    grads = torch.autograd.grad(loss, params,
                                retain_graph=True,
                                create_graph=False,
                                allow_unused=True)
    total = torch.tensor(0.0, device=loss.device)
    for g in grads:
        if g is not None:
            total = total + g.pow(2).sum()
    return total.sqrt()

model = DualHead()
shared = shared_params(model)

x          = torch.randn(64, 14)
rul_target = torch.rand(64, 1) * 125.0
hp_target  = torch.randint(0, 3, (64,))

rul_pred, hp_logits = model(x)
rul_loss    = ((rul_pred - rul_target) ** 2).mean()
health_loss = F.cross_entropy(hp_logits, hp_target)

g_rul    = grad_norm(rul_loss,    shared)
g_health = grad_norm(health_loss, shared)
S        = g_rul + g_health

print(f'||g_rul||    = {g_rul.item():.4f}')
print(f'||g_health|| = {g_health.item():.4f}')
print(f'ratio        = {(g_rul / g_health).item():.1f}x')
print(f'lambda_rul   = {(g_health / S).item():.6f}')
print(f'lambda_health= {(g_rul    / S).item():.6f}')

||g_rul||    = 114.4296
||g_health|| = 0.2164
ratio        = 528.8x
lambda_rul   = 0.001888
lambda_health= 0.998112

1"""Paper code: per-task gradient norm on shared backbone (compute_task_grad_norm)."""
2
3import torch
4import torch.nn as nn
5
6torch.manual_seed(0)
7
8
9class TinyBackbone(nn.Module):
10    def __init__(self):
11        super().__init__()
12        self.fc1 = nn.Linear(14, 32)
13        self.fc2 = nn.Linear(32, 16)
14
15    def forward(self, x):
16        return self.fc2(torch.relu(self.fc1(x)))
17
18
19class DualHead(nn.Module):
20    def __init__(self):
21        super().__init__()
22        self.backbone    = TinyBackbone()
23        self.rul_head    = nn.Linear(16, 1)
24        self.health_head = nn.Linear(16, 3)
25
26    def forward(self, x):
27        feat = self.backbone(x)
28        return self.rul_head(feat), self.health_head(feat)
29
30
31def get_shared_params(model, head_names=("rul_head", "health_head")):
32    """Return parameters that belong to the shared backbone."""
33    out = []
34    for name, p in model.named_parameters():
35        if not p.requires_grad:
36            continue
37        if any(h in name for h in head_names):
38            continue
39        out.append(p)
40    return out
41
42
43def compute_task_grad_norm(loss, shared_params, retain_graph=True):
44    """L2 norm of grad(loss) on shared_params. create_graph=False for speed."""
45    grads = torch.autograd.grad(loss, shared_params, retain_graph=retain_graph, create_graph=False, allow_unused=True)
46    total = torch.tensor(0.0, device=loss.device)
47    for g in grads:
48        if g is not None:
49            total = total + g.pow(2).sum()
50    return total.sqrt()
51
52
53# ---------- One forward pass, two gradient norms ----------
54model  = DualHead()
55shared = get_shared_params(model)
56
57x          = torch.randn(64, 14)
58rul_target = torch.rand(64, 1) * 125.0
59hp_target  = torch.randint(0, 3, (64,))
60
61rul_pred, hp_logits = model(x)
62rul_loss    = ((rul_pred - rul_target) ** 2).mean()
63health_loss = nn.functional.cross_entropy(hp_logits, hp_target)
64
65g_rul    = compute_task_grad_norm(rul_loss,    shared, retain_graph=True)
66g_health = compute_task_grad_norm(health_loss, shared, retain_graph=True)
67
68print(f"||g_rul||    = {g_rul.item():.4f}")
69print(f"||g_health|| = {g_health.item():.4f}")
70print(f"ratio        = {(g_rul / g_health).item():.1f}x")
71
72S = g_rul + g_health
73print(f"\nlambda_rul    = {(g_health / S).item():.6f}")
74print(f"lambda_health = {(g_rul    / S).item():.6f}")

Measured On A Real Backbone

On the paper's actual CNN-BiLSTM-Attention backbone (3.5 M parameters), the same compute_task_grad_norm utility produces the empirical distribution that motivates GABA. Quoting the paper directly (paper main.tex:319):

“During joint training with standard MSE and cross-entropy losses, regression (RUL) gradients exceed classification (health) gradients on shared backbone parameters by ${\sim} 500\times$ (median across $n = 4{,}120$ epoch-level gradient samples from 20 training runs).”

The takeaway: compute_task_grad_norm running across an entire training run produces a SIGNAL, not a single number. GABA's job (in §18.2) is to smooth that signal with EMA so the resulting $\lambda^*$ does not oscillate.

The Same Pattern In Other Fields

In every row, a control mechanism reads a per-step L2 norm and feeds it back into the next step's decision. GABA's contribution is the closed form plugged into stage 2; the measurement (this section) is a cross-disciplinary pattern.

Pitfalls In Per-Step Norm Computation

Takeaway

• GABA's sensor is $g_i = \| \nabla_{\theta_s} L_i \|_2$ on the shared backbone. Computed every step; drives the closed-form weight rule.
• Shared parameters are selected by name. get_shared_params(model) excludes the heads via substring match on named_parameters(). Including heads contaminates the imbalance reading.
• The L2 norm is the right choice. Rotation-invariant, differentiable, standard across PyTorch's clipping / Adam infrastructure.
• Aggregate via sum-of-squares. $\| g \|_2 = \sqrt{ \sum_p \| g_p \|_2^2 }$ gives the same answer as concat-then-norm with dramatically less memory.
• One forward pass serves K backward calls. retain_graph=True on every per-task autograd.grad plus the final combined_loss.backward().
• create_graph=False is what makes GABA cheap. No second-order autograd, no double memory. This is the operational gap vs GradNorm.
• The paper's 480-parameter toy backbone reproduces the 500× imbalance. The ratio is structural — a property of the loss families and target range, not architecture-specific.