TetherGate — how it works

The problem

A learned gate scales a residual branch: out = x + g · branch(x). When the loss does not pay for keeping the branch open — it hurts at initialisation, or its benefit comes later — the optimizer slowly pushes g to zero and the branch stops contributing. Two drifts cause it:

early in training, the control signal s drifts below a threshold theta that has not moved yet;
later, the threshold itself drifts upwards, because the loss-only equilibrium moves towards "closed".

The mechanism

Forward (TetherGate.forward, differentiable):

g_raw = sigmoid( clamp( beta · (s − theta), arg_min, +8 ) )
g     = F + (1 − F) · g_raw

beta is fixed; theta is the only learned parameter.
F (buffer floor) is an affine floor: with F = 1 the gate is fully open ("asleep") and theta receives no loss gradient; as F decays, g_raw takes over. Forward and backward use the same formula — the floor is not a backward-only trick.
In train() mode the module updates ema_s (buffer) with the mean of s (no_grad, no Python branching on tensors, so torch.compile can trace the forward). The first observation initialises the EMA (it starts as NaN).
With detach_signal=True (default) s is detached: the branch cannot learn to steer its own gate.

Gradient and clamp. d g / d theta = −(1 − F) · beta · σ'(a) for a = beta·(s − theta) inside [arg_min, +8], and exactly 0 outside it (the clamp passes no gradient). Inside the range the gradient stays alive even close to the floor; outside it the only force on theta is the spring.

After each optimizer step (tethergate_step, no_grad, outside the graph):

F      = floor_schedule(step)                       # 1 until ramp_start, then exp or cos decay to f_min
theta* = EMA(s) − logit(g_target) / beta            # relative anchor (per gate, or the mean over gates)
k      = kappa_sleep  if F == 1  else  lr · wd
theta ← theta − k · (theta − theta*)

The anchor is relative: it follows the signal's own drift, so the spring removes only the common one-sided drift; the gate's mean opening settles near g_target regardless of the learning rate (the LR sets the time constant), while the per-token selectivity driven by s remains.
anchor='global': one theta* = mean over all gates (gates may differentiate around it). anchor='local': each gate to its own theta*.
Gates whose EMA is still NaN are skipped (also for the floor update).

Choosing `wd`

Under a steady closing push the optimizer moves theta by about lr per step (Adam), while the spring pulls it back by lr · wd · (theta − theta*). Theta then settles about 1/wd from its anchor, whatever the LR, and the gate argument shifts by about beta / wd. With the default wd = 20 and beta = 6 that is 0.3: the gate stays near g_target. With wd = 0.5 a sustained push could shift it by 12, i.e. towards the floor; our long production run used wd = 0.5, which gives the loss much more freedom — with a small wd, monitor the mean gate opening. Stability requires lr · wd ≪ 1.

Floor schedule

mode='exp': F = f_min + (1 − f_min) · exp(−(step − ramp_start)/tau) — approaches f_min asymptotically.
mode='cos': cosine from 1 at ramp_start to f_min at ramp_end.
Driven by the training step only, never by the loss.

State, determinism, cost

State: theta (parameter), floor, ema_s (buffers, saved in state_dict).
Deterministic given the same inputs (no randomness in the module).
Cost: one sigmoid per gate element; tethergate_step reads each EMA as a Python float (one device sync per gate per step).

Tests

test_tethergate.py checks: the forward floor; a live gradient below the floor; zero gradient beyond the clamp and recovery through the spring; detach_signal; EMA updates only in train(); the spring reaching its target; an equilibrium independent of the LR; global vs local anchor.