Warmup (First 100 Steps)

The Cold-Start Problem

On a cold morning your car's engine doesn't go straight to peak efficiency. It runs richer for a minute or two, the catalytic converter waits, the emissions-control loop holds back. The on-board computer knows that running the full closed-loop fuel-injection algorithm against COLD sensors produces nonsense, so it gates the controller off until the engine reaches a sensible operating point.

GABA has the same cold-start problem. The closed form $\lambda^*_i = g_j / (g_i + g_j)$ assumes $g_i$ reflects the steady-state task structure of the loss landscape. At step 0 the model is at random initialisation: the gradients reflect random projections of random labels, not the 500× imbalance the paper characterises in §12.3. Feeding those transient gradients into the closed form — and into the EMA on top — produces garbage $\lambda^*$ for the first few hundred steps.

Why Step-0 Gradients Are Untrustworthy

Three reasons the first ~50 training steps look nothing like the steady-state regime characterised in §12.3:

• Random init dominates. At step 0 the backbone is at Kaiming-init values; logits are random. Cross-entropy on random logits is $\approx \ln K$ regardless of the true labels. MSE on random RUL predictions is dominated by random-prediction variance, not by useful regression structure.
• The 500× imbalance is a steady-state property. Paper main.tex:319 measures it on $n = 4{,}120$ epoch-level samples FROM TRAINED MODELS — not at init. At step 0 the ratio can be anywhere from 1× to 100×, depending on which random seed was drawn.
• Adam's own bias-correction kicks in over the first ~100 steps. Adam's first-moment estimator is biased toward zero for $t \lesssim 1/(1-\beta_1) = 10$ steps; the second-moment estimator for $t \lesssim 1/(1-\beta_2) = 1{,}000$ steps but with smaller variance. Combining adaptive loss weighting with adaptive optimiser estimates during THEIR transient phase amplifies oscillation.

Trying to use GABA from step 1 produces wild oscillations in $\lambda^*$ as the EMA chases transient noise. The visualization below shows this directly: with $W = 0$ (no warmup) the trajectory starts dropping from 0.5 immediately and overshoots before settling. With paper's $W = 100$ the trajectory stays at 0.5 until the model has had a chance to find a stable operating point.

The Warmup Gate (Paper Algorithm 1)

Paper main.tex:362-374 specifies the gate:

$\lambda^*_i = \begin{cases} 1/K & \text{if } t \leq W \\ \text{(GABA pipeline)} & \text{if } t > W \end{cases}$

Where $W = 100$ is the warmup duration and $t$ is a 1-indexed step counter. Two important details:

• The comparison is inclusive. $t \leq W$ means step 100 is still in warmup; step 101 is the first active step. Paper code: self.step_count.item() <= self.warmup_steps.
• The EMA buffer is NOT updated during warmup. The closed form is not computed; there is nothing to feed into the EMA. The buffer stays at its initial value $\hat{\lambda}^{(0)}_i = 1/K$ (paper Algorithm 1 line 2). When the gate flips at step 101, the EMA starts from a sensible 1/K value — which matches the warmup output, so there's no discontinuity in $\lambda^*$ itself.

What Uniform 1/K Means In Practice

For $K = 2$ (RUL + health), uniform weighting means $\lambda^*_{\text{rul}} = \lambda^*_{\text{health}} = 0.5$. The combined loss during warmup is exactly:

$\mathcal{L} = 0.5 \cdot \mathcal{L}_{\text{rul}} + 0.5 \cdot \mathcal{L}_{\text{health}}$

In other words, GABA falls back to the ‘Fixed Baseline’ method (paper §3.5) for the first 100 steps. This is intentional. Fixed Baseline is the simplest, safest, most-studied multi-task scheme: every published MTL paper has experience with it, every deep-learning library has well-understood behaviour for it, and Adam's bias correction has been engineered for losses that look like this.

Why warmup doesn't hurt the 500× imbalance regime. One could worry that 100 steps of equal weighting lets the RUL gradient dominate and train backbone features that are useless for health. Empirically the paper's ablations show no measurable harm: 100 steps is a tiny fraction of the full 500-epoch training horizon, and the active GABA controller has 400+ epochs to re-shape the backbone afterwards.

Why W = 100 Steps

The choice $W = 100$ is paper canonical (main.tex:362). Three justifications:

The robustness ablation in paper §5.8 shows results are statistically indistinguishable for $W \in [50, 200]$; the canonical 100 sits comfortably in the middle.

What Happens At Step W + 1

At step $W + 1$ the gate flips and three things happen in one update:

• Closed form computes. For the paper-realistic 500× imbalance the raw $\lambda$ is $[0.002, 0.998]$.
• EMA absorbs 1% of the new value. $\hat{\lambda} = 0.99 \cdot [0.5, 0.5] + 0.01 \cdot [0.002, 0.998] = [0.49502, 0.50498]$. Tiny shift — the EMA has a long memory.
• Floor is inactive. Both EMA values are above $\lambda_{\min} = 0.05$; clamp is a no-op; renormalisation divides by 1.

So the FIRST active step changes $\lambda^*$ from $0.5$ to $0.49502$ — a 0.5% relative shift, invisible to the optimiser. Convergence to the steady-state 0.04762 takes another $\sim 3\tau = 300$ steps. Total: warmup + post-warmup settling = 400 steps to reach the floor-bound regime, out of typically $50{,}000$+ training steps.

Interactive: Slide The Warmup Boundary

Drag the W slider. The amber band marks the warmup region; inside it $\lambda^*_{\text{rul}}$ is pinned at 0.5. Outside the band the GABA pipeline runs. The dashed grey line is the same simulation with $W = 0$ for comparison. Watch how W = 0 immediately starts dropping while W = 100 stays flat through the window.

Try this. Set W = 0 and watch the blue and grey traces collapse onto each other — no warmup, immediate GABA. Set W = 300 and watch the trace stay at 0.5 for the entire visible window — warmup eats the whole simulation. Paper's W = 100 is the sweet spot: enough warmup to dodge the cold-start transient, short enough to leave the rest of training in active mode.

Python: Warmup Gate From Scratch

Implement the full GABA per-step update with the warmup gate in pure NumPy. Run for 500 steps and print the $\lambda_{\text{rul}}$ trajectory at milestone steps so the warmup plateau and the post-warmup settling are visible.

 step |    lam_rul | regime
----------------------------------------
    0 |   0.500000 | warmup
   50 |   0.500000 | warmup
   99 |   0.500000 | warmup
  100 |   0.495051 | active
  101 |   0.490150 | active
  150 |   0.301712 | active
  200 |   0.185016 | active
  300 |   0.070340 | active
  499 |   0.048225 | active

1"""Warmup gate from scratch: paper Algorithm 1, lines 4-6."""
2
3import numpy as np
4
5np.random.seed(0)
6
7# Paper canonical hyperparameters
8W       = 100        # warmup steps
9beta    = 0.99       # EMA coefficient
10lam_min = 0.05       # floor
11K       = 2          # tasks
12
13
14def gaba_step(step_count, ema_w, g_rul, g_health):
15    """One full GABA step with warmup gate."""
16    if step_count <= W:
17        return np.full(K, 1.0 / K), ema_w   # warmup: uniform weights
18    # Closed form (eq. 4 specialised to K=2)
19    S = g_rul + g_health
20    raw = np.array([g_health / S, g_rul / S])
21    # EMA (eq. 5)
22    ema_w = beta * ema_w + (1.0 - beta) * raw
23    # Floor + renormalise (eq. 6)
24    clamped = np.maximum(ema_w, lam_min)
25    return clamped / clamped.sum(), ema_w
26
27
28# ---------- Synthetic gradients that drift toward the 500x regime ----------
29def synth(t):
30    blend = min(1.0, t / 50.0)
31    g_rul    = max((2.0 + blend * 3.0) + 0.5  * np.random.randn(), 0.01)
32    g_health = max((0.5 - blend * 0.49) + 0.05 * np.random.randn(), 1e-6)
33    return g_rul, g_health
34
35
36# ---------- Run 500 steps with the paper's warmup ----------
37ema_w = np.array([0.5, 0.5])
38trace = []
39for t in range(500):
40    step_count = t + 1
41    g_rul, g_health = synth(t)
42    weights, ema_w = gaba_step(step_count, ema_w, g_rul, g_health)
43    trace.append(weights[0])
44
45
46# ---------- Print milestones ----------
47print(f"{'step':>5} | {'lam_rul':>10} | regime")
48print("-" * 40)
49for t in [0, 50, 99, 100, 101, 150, 200, 300, 499]:
50    regime = "warmup" if (t + 1) <= W else "active"
51    print(f"{t:>5} | {trace[t]:>10.6f} | {regime}")

PyTorch: The Paper's Branching Code

The actual paper code lives at grace/core/gaba.py:107-108 and is two lines: if shared_params is None or self.step_count.item() <= self.warmup_steps: followed by weights = torch.ones(K, device=device) / K. We extract the gate (and its surrounding state-management machinery) into a standalone class for clarity.

step |    lam_rul | lam_health | regime
--------------------------------------------------
   1 |   0.500000 |   0.500000 | warmup
  50 |   0.500000 |   0.500000 | warmup
  99 |   0.500000 |   0.500000 | warmup
 100 |   0.500000 |   0.500000 | warmup
 101 |   0.495020 |   0.504980 | active
 105 |   0.475570 |   0.524430 | active
 110 |   0.452029 |   0.547971 | active
 130 |   0.367892 |   0.632108 | active
 150 |   0.301880 |   0.698120 | active

1"""Paper code: warmup gate as in grace/core/gaba.py:107-108."""
2
3import torch
4import torch.nn as nn
5
6
7class WarmupGabaLoss(nn.Module):
8    """Minimal GABA module showing the warmup gate from paper code.
9
10    Mirrors lines 107-108 of grace/core/gaba.py:
11
12        if shared_params is None or self.step_count.item() <= self.warmup_steps:
13            weights = torch.ones(K, device=device) / K
14        else:
15            ...full GABA pipeline...
16    """
17
18    def __init__(self, beta: float = 0.99, warmup_steps: int = 100,
19                 min_weight: float = 0.05, n_tasks: int = 2) -> None:
20        super().__init__()
21        self.beta         = beta
22        self.warmup_steps = warmup_steps
23        self.min_weight   = min_weight
24        self.n_tasks      = n_tasks
25        self.register_buffer("ema_weights", torch.ones(n_tasks) / n_tasks)
26        self.register_buffer("step_count",  torch.tensor(0, dtype=torch.long))
27
28    def forward(self, raw_weights: torch.Tensor, shared_params=None) -> torch.Tensor:
29        K = self.n_tasks
30        device = raw_weights.device
31        self.step_count += 1
32
33        # WARMUP GATE — paper Algorithm 1 lines 4-6
34        if shared_params is None or self.step_count.item() <= self.warmup_steps:
35            return torch.ones(K, device=device) / K
36
37        # ACTIVE branch — closed form is already in raw_weights; do EMA + floor.
38        ema_w = self.beta * self.ema_weights + (1.0 - self.beta) * raw_weights
39        self.ema_weights = ema_w.detach()
40        clamped = ema_w.clamp(min=self.min_weight)
41        return clamped / clamped.sum()
42
43
44# ---------- Smoke test: walk steps 1..150 with constant raw input ----------
45torch.manual_seed(0)
46gaba = WarmupGabaLoss(beta=0.99, warmup_steps=100, min_weight=0.05, n_tasks=2)
47raw_lambda = torch.tensor([0.002, 0.998])
48shared = [torch.zeros(1)]      # placeholder so the gate doesn't fall through
49
50print(f"{'step':>4} | {'lam_rul':>10} | {'lam_health':>10} | regime")
51print("-" * 50)
52for step in range(1, 151):
53    out = gaba(raw_lambda, shared_params=shared)
54    if step in (1, 50, 99, 100, 101, 105, 110, 130, 150):
55        regime = "warmup" if step <= gaba.warmup_steps else "active"
56        print(f"{step:>4} | {out[0].item():>10.6f} | {out[1].item():>10.6f} | {regime}")

Warmup Patterns In Other Pipelines

The pattern — gate the adaptive controller off until the system has reached a sensible operating point — recurs across nearly every adaptive learning pipeline. GABA inherits the pattern; the implementation is just a step counter and an inclusive comparison.

Pitfalls In Warmup Gating

Takeaway

• Warmup gates the adaptive controller off for the first W = 100 steps. Paper Algorithm 1 lines 4-6: if $t \leq W$ return $1/K$ uniform weights; otherwise run the GABA pipeline.
• W = 100 matches the EMA time constant. One τ = 1/(1−β) = 100 steps. This alignment isn't coincidence: warmup gives the EMA exactly one time constant of clean uniform initialisation before adaptive logic kicks in.
• The EMA buffer is not updated during warmup. It stays at the initial $1/K$, which matches the warmup output, so no discontinuity in $\lambda^*$ at step W+1.
• The transition is smooth. First active step shifts $\lambda^*$ by $\sim 0.5\%$ for the realistic 500× imbalance — well below Adam's noise floor.
• step_count must be a buffer. So checkpoint resume doesn't restart the warmup counter mid-training.
• The pattern is universal. Linear LR warmup, BatchNorm running stats, replay-buffer fill, gradual unfreezing — every adaptive controller benefits from a startup gate.