tethergate / TECHNICAL.md

Initial release: code, docs, tests (CC BY-NC 4.0)

d7070c6 22 days ago

4.04 kB

# 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:

1. early in training, the control signal `s` drifts below a threshold `theta` that has not moved yet;
2. 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.

	# 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:

	1. early in training, the control signal `s` drifts below a threshold `theta` that has not moved yet;
	2. 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.