GABA as a Proportional Feedback Controller

Hook: Cruise Control On A Hilly Road

A car cruising at 100 km/h on a hilly highway does not actually run at 100 km/h. The road tilts up, the car slows, the speedometer drops to 96. The cruise-control unit reads the discrepancy — $e = 100 - 96 = 4\ \text{km/h}$ — and pushes the throttle harder by an amount proportional to that error. Over the next second, the car accelerates back toward 100. The loop runs ten times per second and the driver feels nothing.

That little box is a proportional feedback controller. It has four ingredients: a sensor (speedometer), a setpoint (100 km/h), a comparator (subtraction), and an actuator (throttle). The control action is proportional to the error. If you slept through control engineering, this is the only sentence you need.

GABA is doing exactly the same thing — not in km/h, but in gradient space. The plant is the joint multi-task optimiser. The sensor is $g_i = \| \nabla_\theta \mathcal{L}_i \|_2$ (per-task gradient norm on the shared backbone). The setpoint is equal contribution from all tasks, $r = 1/K$. The comparator subtracts the measured share from the setpoint. The control action — the task weight $\lambda_i^{*}$ — is emitted by an inverse-proportional law. Paper main.tex:387 names this explicitly: “GABA can be viewed as a proportional feedback controller operating in gradient space.” This section unpacks what that sentence means line by line.

Why this framing matters. Once you see GABA as a P-controller, three properties of the algorithm stop being coincidences and become consequences: (1) the weights always sum to 1, (2) the rule has no learnable parameters and just three constants, and (3) the bounded-weight guarantee proven in §19.3 is the anti-windup property of any well-designed feedback loop. Open-loop alternatives (notably GradNorm) do not get these for free — they diverged on 1/5 N-CMAPSS seeds (paper main.tex:553). The closed loop is what makes GABA robust.

Five Words Every Control Engineer Knows

Before we map cruise control onto GABA, fix the vocabulary. Any single-input single-output (SISO) feedback loop has these five pieces:

The diagram is always the same: r → (+) → controller → plant → sensor → back to (-) of the comparator. The minus sign is what makes it ‘negative feedback’: the loop subtracts the measurement from the reference. Without that minus sign you have positive feedback, the microphone-howl regime, which we explicitly do not want.

A proportional controller is the simplest non-trivial choice for $C(e)$: $u = K_p \cdot e$. The control command is directly proportional to the error, with gain $K_p$. Larger error ⇒ larger correction. Zero error ⇒ zero correction. Linear, memoryless, dead simple.

Mapping Those Five Words To GABA

Now we replace cruise-control nouns with GABA nouns. Same five positions in the loop, different physical interpretation.

The actuator output $\lambda_i^{*}$ — the task weight — is what gets multiplied against the per-task losses in the next forward step: $\mathcal{L} = \sum_i \lambda_i^{*} \mathcal{L}_i$. The gradient of that combined loss flows back into the plant (the optimiser updates $\theta$), the new $\theta$ changes the next batch's $g_i$, the sensor reads the new shares, and the loop closes.

The Math: GABA Is A P-Controller With Unity Gain

Set $K = 2$ for clarity (the paper's case; the K-task version is below). The measured share of task $i$ is

$s_i = \frac{g_i}{g_1 + g_2}, \qquad s_1 + s_2 = 1.$

The reference is $r = 1/2$. The error is

$e_i = r - s_i = \frac{1}{2} - \frac{g_i}{g_1 + g_2} = \frac{(g_1 + g_2) - 2 g_i}{2(g_1 + g_2)} = \frac{g_j - g_i}{2(g_1 + g_2)}$

where $j$ is ‘the other task’. Now the proportional law with unity gain emits

$\lambda_i^{*} = r + K_p \cdot e_i = \frac{1}{2} + 1 \cdot \frac{g_j - g_i}{2(g_1 + g_2)} = \frac{(g_1 + g_2) + (g_j - g_i)}{2(g_1 + g_2)} = \frac{2 g_j}{2(g_1 + g_2)} = \frac{g_j}{g_1 + g_2}.$

That is exactly the paper's closed form for $K = 2$: $\lambda_i^{*} = (\Sigma g - g_i)/\Sigma g = g_j / \Sigma g$. So GABA's closed form is mathematically identical to $r + K_p \cdot e_i$ with $K_p = 1$. Q.E.D.

For general $K$ tasks the algebra is the same with a slightly bigger normaliser:

$\lambda_i^{*} = \frac{\sum_j g_j - g_i}{(K - 1) \sum_j g_j} = \frac{1}{K} + \frac{1}{K - 1} \cdot e_i, \qquad e_i = \frac{1}{K} - \frac{g_i}{\sum_j g_j}.$

The proportional gain becomes $K_p = 1/(K - 1)$ — still constant, still proportional, still no learned parameters. For the paper's $K = 2$ we recover $K_p = 1$ exactly.

Three numerical sanity checks fall out of the math:

• Sum to 1. $\sum_i \lambda_i^{*} = \sum_i (\Sigma g - g_i)/((K-1)\Sigma g) = (K \Sigma g - \Sigma g)/((K-1)\Sigma g) = 1$. The controller output is automatically on the simplex. No renormalisation needed.
• Symmetry. Swap $g_1 \leftrightarrow g_2$ and the weights swap accordingly. The rule is permutation-equivariant in the tasks.
• Equal-grad fixed point. If $g_1 = g_2$, the error is zero and the controller emits $\lambda^{*} = (1/2, 1/2)$ — equal weights, no correction. Equal-gradient is a fixed point of the closed loop.

Interactive: The Closed Loop, Live

Move the two gradient-norm sliders below. Watch the controller emit $\lambda^{*}$ in real time. The block diagram shows the closed loop — reference, comparator, P-law, plant, sensor — with the live values flowing along the wires. The preset buttons jump to specific operating points: Equal, 10× imbalance, 500× imbalance (the paper's measured median), and 2,400× peak (the paper's observed peak ratio early in training).

Three things to verify by playing with the sliders. First, when you nudge $g_{\text{rul}}$ up, $\lambda_{\text{rul}}$ goes down — this is the inverse action of the controller. Second, the two weights always sum to 1 (read the bottom-right readout). Third, when both sliders sit at the same value, the controller emits exactly $(0.5, 0.5)$ — the equal-gradient fixed point.

Python: P-Controller From Scratch (NumPy)

We'll first build the controller as a textbook P-controller with explicit sense/error/control methods so every block of the loop has a name. This is pedagogical: the production code (next subsection) collapses these into a single closed-form expression, but the algebraic equivalence is exactly what we proved above.

reference  = 0.5000
shares     = [0.99800399 0.00199601]
errors     = [-0.49800399  0.49800399]
weights    = [0.00199601 0.99800399]
sums to 1  : 1.000000
manual P   = [0.00199601 0.99800399]    (matches lambda above)

1"""GABA's INNER LOOP as a textbook proportional feedback controller.
2
3Pure NumPy. No EMA, no floor (those are sections 19.2 and 19.3).
4This file shows that the closed-form weight rule of GABA
5
6    lambda_i = (sum(g) - g_i) / ((K - 1) * sum(g))
7
8is a P-controller of unity gain in the *measured share*
9
10    s_i = g_i / sum(g)
11
12that drives s_i toward the reference r = 1/K.
13"""
14
15import numpy as np
16
17
18class GabaPController:
19    """Proportional feedback controller in gradient space."""
20
21    def __init__(self, n_tasks=2, K_p=1.0):
22        self.K = n_tasks
23        self.K_p = K_p          # paper uses unity gain
24        self.reference = 1.0 / n_tasks
25
26    def sense(self, grad_norms):
27        """Sensor: convert raw gradient norms to per-task share."""
28        g = np.asarray(grad_norms, dtype=np.float64)
29        total = g.sum() + 1e-12
30        return g / total
31
32    def error(self, share):
33        """Error signal: how far each task is from the equal-share reference."""
34        return self.reference - share
35
36    def control(self, grad_norms):
37        """P-action: emit task weights from the closed-form inverse rule."""
38        g = np.asarray(grad_norms, dtype=np.float64)
39        total = g.sum() + 1e-12
40        # Closed form (paper eq. 4) ≡ r + K_p * error  for K=2.
41        return (total - g) / ((self.K - 1) * total)
42
43
44# ---- Single-step demo at the paper's measured 500x imbalance ----
45ctrl = GabaPController(n_tasks=2, K_p=1.0)
46
47g = [250.0, 0.5]                # paper-realistic g_rul, g_health
48share = ctrl.sense(g)           # measured share
49err   = ctrl.error(share)       # error from equal share
50lam   = ctrl.control(g)         # control action
51
52print(f"reference  = {ctrl.reference:.4f}")
53print(f"shares     = {share}")
54print(f"errors     = {err}")
55print(f"weights    = {lam}")
56print(f"sums to 1  : {lam.sum():.6f}")
57
58# ---- Sanity check: P-controller equivalent form (K=2 only) ----
59# For K=2:  lambda_i = r + K_p * (r - s_i)   with K_p = 1
60manual = ctrl.reference + ctrl.K_p * err
61print(f"manual P   = {manual}    (matches lambda above)")

The output confirms the math: at the paper's 500× imbalance, the controller emits $(\lambda_{\text{rul}}, \lambda_{\text{health}}) \approx (0.002, 0.998)$ — the inverse of the share — and the explicit P-form $\lambda = r + K_p \cdot e$ reproduces the same numbers exactly.

PyTorch: The Inner Loop Of GABALoss Verbatim

The paper's production code is in grace/core/gaba.py. To isolate the proportional inner loop — with no EMA, no floor, no warmup — we extract just the three lines of forward_k that compute the closed form. The result is bit-exact to the un-smoothed branch of GABALoss when $\beta = 1$ and $\lambda_{\min} = 0$.

Tensor (8, 3).

Tensor (8, 3). Per-example, per-class scores.

grad norms = [4.221, 0.064]
weights    = [0.0150, 0.9850]    (sums to 1.0000)
total loss = 1.0832
# (Exact numbers vary slightly with torch version due to autograd algebra,
#  but the qualitative pattern — λ inverse to grad_norms — is invariant.)

1"""The proportional inner loop of GABALoss in PyTorch.
2
3Strips out the EMA and floor (those are sec. 19.2, 19.3). What remains
4is exactly the P-controller: sensor + comparator + proportional law.
5
6This block is the inner three lines of grace/core/gaba.py:GABALoss.forward_k.
7"""
8
9import torch
10import torch.nn as nn
11
12
13def compute_task_grad_norm(loss, shared_params, retain_graph=True):
14    """L2 norm of grad(loss) on shared_params (paper grace/core/gradient_utils.py)."""
15    grads = torch.autograd.grad(loss, shared_params,
16                                 retain_graph=retain_graph,
17                                 create_graph=False, allow_unused=True)
18    total = torch.tensor(0.0, device=loss.device)
19    for g in grads:
20        if g is not None:
21            total = total + g.pow(2).sum()
22    return total.sqrt()
23
24
25class GabaProportionalInnerLoop(nn.Module):
26    """The inner P-controller of GABA, with no EMA, no floor.
27
28    Identical numerics to the un-smoothed branch of paper GABALoss when
29    beta = 1 (no new-measurement absorption) and min_weight = 0.
30    """
31
32    def __init__(self, n_tasks=2):
33        super().__init__()
34        self.n_tasks = n_tasks
35        self.register_buffer("reference",
36                              torch.full((n_tasks,), 1.0 / n_tasks))
37
38    def forward(self, losses, shared_params):
39        K = len(losses)
40        device = losses[0].device
41        # Sensor: per-task gradient norms on shared backbone.
42        grad_norms = torch.zeros(K, device=device)
43        for i, loss_i in enumerate(losses):
44            grad_norms[i] = compute_task_grad_norm(loss_i, shared_params)
45        # Closed-form proportional law (paper eq. 4 for K tasks).
46        total_norm = grad_norms.sum() + 1e-12
47        lam = (total_norm - grad_norms) / ((K - 1) * total_norm)
48        # Combined loss = sum(lambda_i * loss_i)
49        return sum(w * L for w, L in zip(lam, losses)), lam, grad_norms
50
51
52# ---- Tiny smoke test (no real model) ----
53torch.manual_seed(0)
54
55# Two losses sharing a (3,) backbone parameter:
56shared = torch.nn.Parameter(torch.randn(3))
57x = torch.randn(8, 3)
58y_rul    = torch.randn(8)
59y_health = torch.randint(0, 3, (8,))
60
61# Manufacture realistic per-task losses on shared
62rul_pred    = (x * shared).sum(dim=1)                           # regression head
63logits      = torch.outer(x.mean(dim=1), torch.ones(3)) * shared
64rul_loss    = ((rul_pred - y_rul) ** 2).mean()                  # MSE: large grad
65health_loss = nn.functional.cross_entropy(logits, y_health)     # CE: small grad
66
67inner = GabaProportionalInnerLoop(n_tasks=2)
68total, lam, gn = inner([rul_loss, health_loss], [shared])
69print(f"grad norms = {gn.tolist()}")
70print(f"weights    = {lam.tolist()}    (sums to {lam.sum().item():.4f})")
71print(f"total loss = {total.item():.4f}")

The smoke test demonstrates the controller in action on real autograd: the regression task's gradient norm exceeds the classification task's by ~70× on this random tensor, and the controller emits weights ~(0.015, 0.985) — the exact inverse-share. On the paper's real 3.5M-parameter backbone the imbalance is closer to 500× (paper main.tex:64) and the controller pushes $\lambda_{\text{health}}$ to ~0.995 (paper main.tex:647) before the EMA filter slows it down (sec. 19.2).

The Same P-Controller In Other Domains

The control-theoretic framing is not just a metaphor. The same proportional inner loop — sensor, comparator, proportional law, plant, feedback wire — appears in physical, biological, and digital systems whose engineering history goes back centuries. Recognising the pattern across domains is what tells you that GABA's stability story is on solid ground.

Two patterns repeat. First, the P-only controller is rarely the whole story — in every application a real engineer adds memory (the I in PID), prediction (the D in PID), or a bounding mechanism. GABA does the same: the EMA filter (§19.2) is a bounded-memory smoother, and the floor (§19.3) is the bounding mechanism. Second, in every well-designed loop the controller emits a value inverse to whatever the sensor reads — if the engine slows, throttle up; if the gradient is too large, weight down. That is precisely the inverse-proportional rule of paper eq. 4.

Cross-domain analogy: federated learning. A 2024 line of work on robust federated averaging weights edge-device gradients inversely to their L2 norm to prevent a single chatty client from dominating the global update. The mechanism is identical to GABA's inverse-share rule, the goal is identical (equal effective contribution to the shared model), and the derivation works through the same five-piece feedback loop. Recognise the closed loop in one domain and you can transfer the analysis to another.

Why ‘Feedback’ Matters: GradNorm Has Open-Loop Failure Modes

The strongest argument for the closed-loop framing is what happens without it. GradNorm (Chen et al. 2018) tries to equalise per-task gradient magnitudes by adding an auxiliary loss that is itself differentiated and minimised alongside the task losses. From a control-theoretic standpoint this is an open-loop design: the auxiliary loss does not directly observe the current measured share and emit a proportional weight; instead it feeds the error into a parameter that is itself updated by gradient descent — a cascade of two optimisation processes whose stability depends on the relative time constants of the inner and outer loops.

The paper observed the consequence empirically (paper main.tex:553): “GradNorm diverged on 1 of 5 N-CMAPSS seeds (seed 789 produced NaN gradients); its auxiliary gradient-balancing loss can amplify oscillations when the imbalance exceeds 500×, whereas GABA's direct inverse-ratio formulation with EMA smoothing avoids this failure mode.”

The take-away is structural, not numerical. Closed-loop design with a fixed proportional gain and an explicit bounding mechanism is more robust than open-loop design with a learnable parameter, particularly when the disturbance amplitude is 500× the reference. This is a textbook control-engineering result — well-tuned P controllers beat parameter-adapted ones in high-disturbance, low-prior settings — and GABA's 5/5 vs. GradNorm's 4/5 reproduction rate on N-CMAPSS is the empirical confirmation in the prognostics setting.

Pitfalls When Treating GABA As A Controller

Takeaway

• GABA's closed-form rule IS a textbook P-controller. The algebraic equivalence is exact: $\lambda_i^{*} = r + K_p \cdot e_i$ with $r = 1/K$, $K_p = 1/(K - 1)$, and $e_i$ the deviation of measured share from the equal-share reference.
• The five blocks are concrete. Plant = joint optimiser. Sensor = $\| \nabla_\theta \mathcal{L}_i \|_2$. Reference = $1/K$. Comparator = subtraction absorbed into the closed form. Controller = paper eq. 4. Feedback wire = the next batch's autograd graph.
• Unity gain is the algebraic optimum. The closed-form rule fixes $K_p = 1/(K-1)$ — the exact value at which a one-step proportional correction would zero the share error in the absence of EMA smoothing. No tuning, no learnable parameters.
• Three guarantees fall out of the framing. Weights sum to 1 (simplex constraint). The rule is symmetric under task permutation. Equal gradients is a fixed point.
• Closed loop > open loop in disturbance regimes. GABA's 5/5 reproduction rate vs. GradNorm's 4/5 on N-CMAPSS (paper main.tex:553) is the empirical confirmation of a textbook control-engineering result: well-designed feedback rejects high-magnitude disturbances better than parameter-adapted alternatives.
• The next two sections close the analysis. §19.2 identifies the EMA buffer as a first-order IIR low-pass filter with time constant $\tau = 1/(1-\beta) = 100$ steps, completing the noise-rejection story. §19.3 identifies the floor as the anti-windup mechanism and proves the bounded-weight guarantee in $[\lambda_{\min}, 1 - \lambda_{\min}]$ that GradNorm cannot match.